from __future__ import annotations

import hashlib
import sqlite3
import time
from collections import defaultdict
from pathlib import Path
from statistics import mean, pstdev

from .logs import THREAT_ACTIONS, is_utm_event
from .models import LogEvent


def _number(value: str | None) -> int:
    try:
        return int(float(value or 0))
    except ValueError:
        return 0


class BaselineStore:
    """Persistent five-minute behavior baseline, implemented with stdlib SQLite."""

    def __init__(self, path: str, *, bucket_seconds: int = 300) -> None:
        self.path = Path(path)
        self.path.parent.mkdir(parents=True, exist_ok=True)
        self.bucket_seconds = bucket_seconds
        with self._connect() as connection:
            connection.executescript(
                """
                create table if not exists seen_events (fingerprint text primary key);
                create table if not exists source_buckets (
                  source_ip text not null, bucket_start integer not null,
                  events integer not null, bytes integer not null, hitcount integer not null,
                  denies integer not null, utm integer not null,
                  primary key (source_ip, bucket_start)
                );
                create table if not exists source_values (
                  source_ip text not null, kind text not null, value text not null,
                  seen_count integer not null, primary key (source_ip, kind, value)
                );
                """
            )

    def _connect(self) -> sqlite3.Connection:
        return sqlite3.connect(self.path)

    def ingest(self, events: list[LogEvent], *, observed_at: int | None = None) -> int:
        observed_at = observed_at or int(time.time())
        pending: dict[tuple[str, int], list[int]] = defaultdict(lambda: [0, 0, 0, 0, 0])
        inserted = 0
        with self._connect() as connection:
            for event in events:
                if not event.src_ip:
                    continue
                fingerprint = hashlib.sha256(event.raw.encode("utf-8", errors="replace")).hexdigest()
                if connection.execute("insert or ignore into seen_events values (?)", (fingerprint,)).rowcount != 1:
                    continue
                bucket = observed_at - (observed_at % self.bucket_seconds)
                values = pending[(event.src_ip, bucket)]
                values[0] += 1
                values[1] += _number(event.fields.get("sentbyte")) + _number(event.fields.get("rcvdbyte"))
                values[2] += _number(event.fields.get("hitcount"))
                values[3] += int(event.action in THREAT_ACTIONS)
                values[4] += int(is_utm_event(event))
                for kind, value in (("destination", event.dst_ip), ("destination_port", event.fields.get("dstport"))):
                    if value:
                        connection.execute(
                            """insert into source_values values (?, ?, ?, 1)
                            on conflict(source_ip, kind, value) do update set seen_count=seen_count+1""",
                            (event.src_ip, kind, value),
                        )
                inserted += 1
            for (source_ip, bucket), values in pending.items():
                connection.execute(
                    """insert into source_buckets values (?, ?, ?, ?, ?, ?, ?)
                    on conflict(source_ip, bucket_start) do update set
                      events=events+excluded.events, bytes=bytes+excluded.bytes,
                      hitcount=hitcount+excluded.hitcount, denies=denies+excluded.denies, utm=utm+excluded.utm""",
                    (source_ip, bucket, *values),
                )
        return inserted

    def profiles(self, source_ips: set[str]) -> dict[str, dict[str, object]]:
        profiles: dict[str, dict[str, object]] = {}
        with self._connect() as connection:
            for source_ip in source_ips:
                rows = connection.execute(
                    "select events, bytes, hitcount from source_buckets where source_ip=? order by bucket_start", (source_ip,)
                ).fetchall()
                # Do not compare the current (often incomplete) bucket to itself.
                rows = rows[:-1]
                if len(rows) < 12:
                    continue
                rates = [row[0] * 60 / self.bucket_seconds for row in rows]
                hit_rates = [row[2] * 60 / self.bucket_seconds for row in rows]
                known = connection.execute(
                    "select kind, value from source_values where source_ip=?", (source_ip,)
                ).fetchall()
                profiles[source_ip] = {
                    "samples": len(rows),
                    "event_rate_mean": mean(rates), "event_rate_stddev": pstdev(rates) or 1.0,
                    "hitcount_rate_mean": mean(hit_rates), "hitcount_rate_stddev": pstdev(hit_rates) or 1.0,
                    "known_destinations": [value for kind, value in known if kind == "destination"],
                    "known_destination_ports": [value for kind, value in known if kind == "destination_port"],
                }
        return profiles
