from __future__ import annotations
import json
import sqlite3
import time
from pathlib import Path
from typing import Any

RETENTION_SECONDS = 7 * 24 * 60 * 60


class TlsTestDB:
    def __init__(self, filepath: str | Path | None = None):
        if filepath is None:
            filepath = Path.cwd() / "databases" / "tls_test.db"
        self.filepath = str(filepath)
        Path(self.filepath).parent.mkdir(parents=True, exist_ok=True)
        self._init_schema()

    def _conn(self) -> sqlite3.Connection:
        conn = sqlite3.connect(self.filepath, timeout=30)
        conn.execute("PRAGMA journal_mode=WAL")
        conn.execute("PRAGMA busy_timeout=10000")
        return conn

    def _init_schema(self) -> None:
        conn = self._conn()
        try:
            conn.executescript(
                """
                CREATE TABLE IF NOT EXISTS tests (
                    id              TEXT PRIMARY KEY,
                    target          TEXT NOT NULL,
                    client_ip       TEXT,
                    status          TEXT NOT NULL,
                    created_at      INTEGER NOT NULL,
                    started_at      INTEGER,
                    finished_at     INTEGER,
                    expires_at      INTEGER NOT NULL,
                    rank            TEXT,
                    score           REAL,
                    result_json     TEXT,
                    error_message   TEXT
                );
                CREATE INDEX IF NOT EXISTS idx_tests_expires ON tests(expires_at);
                CREATE INDEX IF NOT EXISTS idx_tests_ip_created ON tests(client_ip, created_at);

                CREATE TABLE IF NOT EXISTS test_progress (
                    test_id         TEXT NOT NULL,
                    seq             INTEGER NOT NULL,
                    ts              INTEGER NOT NULL,
                    phase           TEXT NOT NULL,
                    detail          TEXT,
                    progress        REAL,
                    severity        TEXT,
                    PRIMARY KEY(test_id, seq)
                );
                """
            )
            conn.commit()
        finally:
            conn.close()

    def create_job(self, test_id: str, target: str, client_ip: str | None) -> None:
        now = int(time.time())
        conn = self._conn()
        try:
            conn.execute(
                "INSERT INTO tests(id,target,client_ip,status,created_at,expires_at)"
                " VALUES (?,?,?,?,?,?)",
                (test_id, target, client_ip, "queued", now, now + RETENTION_SECONDS),
            )
            conn.commit()
        finally:
            conn.close()

    def update_status(
        self,
        test_id: str,
        status: str,
        *,
        started_at: int | None = None,
        finished_at: int | None = None,
        rank: str | None = None,
        score: float | None = None,
        result_json: str | None = None,
        error_message: str | None = None,
    ) -> None:
        sets = ["status = ?"]
        args: list[Any] = [status]
        if started_at is not None:
            sets.append("started_at = ?"); args.append(started_at)
        if finished_at is not None:
            sets.append("finished_at = ?"); args.append(finished_at)
        if rank is not None:
            sets.append("rank = ?"); args.append(rank)
        if score is not None:
            sets.append("score = ?"); args.append(score)
        if result_json is not None:
            sets.append("result_json = ?"); args.append(result_json)
        if error_message is not None:
            sets.append("error_message = ?"); args.append(error_message)
        args.append(test_id)
        conn = self._conn()
        try:
            conn.execute(f"UPDATE tests SET {', '.join(sets)} WHERE id = ?", args)
            conn.commit()
        finally:
            conn.close()

    def append_progress(
        self,
        test_id: str,
        seq: int,
        phase: str,
        detail: str,
        progress: float,
        severity: str,
    ) -> None:
        conn = self._conn()
        try:
            conn.execute(
                "INSERT OR REPLACE INTO test_progress(test_id,seq,ts,phase,detail,progress,severity)"
                " VALUES (?,?,?,?,?,?,?)",
                (test_id, seq, int(time.time()), phase, detail, progress, severity),
            )
            conn.commit()
        finally:
            conn.close()

    def get_job(self, test_id: str) -> dict[str, Any] | None:
        conn = self._conn()
        try:
            cur = conn.cursor()
            cur.execute("SELECT * FROM tests WHERE id = ?", (test_id,))
            row = cur.fetchone()
            if row is None:
                return None
            cols = [c[0] for c in cur.description]
            d = dict(zip(cols, row))
            if d.get("result_json"):
                try:
                    d["result"] = json.loads(d["result_json"])
                except json.JSONDecodeError:
                    d["result"] = None
            return d
        finally:
            conn.close()

    def get_progress(self, test_id: str) -> list[dict[str, Any]]:
        conn = self._conn()
        try:
            cur = conn.cursor()
            cur.execute(
                "SELECT seq,ts,phase,detail,progress,severity FROM test_progress"
                " WHERE test_id = ? ORDER BY seq ASC",
                (test_id,),
            )
            rows = cur.fetchall()
            return [
                {
                    "seq": r[0],
                    "ts": r[1],
                    "phase": r[2],
                    "detail": r[3],
                    "progress": r[4],
                    "severity": r[5],
                }
                for r in rows
            ]
        finally:
            conn.close()

    def count_ip_in_window(self, client_ip: str, window_seconds: int) -> int:
        cutoff = int(time.time()) - window_seconds
        conn = self._conn()
        try:
            cur = conn.cursor()
            cur.execute(
                "SELECT COUNT(*) FROM tests WHERE client_ip = ? AND created_at >= ?",
                (client_ip, cutoff),
            )
            return cur.fetchone()[0]
        finally:
            conn.close()

    def count_ip_active(self, client_ip: str) -> int:
        conn = self._conn()
        try:
            cur = conn.cursor()
            cur.execute(
                "SELECT COUNT(*) FROM tests WHERE client_ip = ? AND status IN ('queued','running')",
                (client_ip,),
            )
            return cur.fetchone()[0]
        finally:
            conn.close()

    def delete_expired(self) -> int:
        now = int(time.time())
        conn = self._conn()
        try:
            cur = conn.cursor()
            cur.execute("SELECT id FROM tests WHERE expires_at < ?", (now,))
            ids = [r[0] for r in cur.fetchall()]
            for tid in ids:
                conn.execute("DELETE FROM test_progress WHERE test_id = ?", (tid,))
                conn.execute("DELETE FROM tests WHERE id = ?", (tid,))
            conn.commit()
            return len(ids)
        finally:
            conn.close()
