from __future__ import annotations

import json
import sqlite3
import uuid
from collections.abc import Iterator
from contextlib import contextmanager
from datetime import UTC, datetime, timedelta
from pathlib import Path
from typing import Any

from app.config import MAX_ARCHIVE_RETENTION_DAYS


def utcnow() -> datetime:
    return datetime.now(UTC)


def timestamp(value: datetime | None = None) -> str:
    return (value or utcnow()).isoformat()


class Database:
    """Small SQLite persistence layer with lease-based job claiming."""

    def __init__(self, path: Path) -> None:
        self.path = path

    @contextmanager
    def connection(self, *, timeout: float = 15.0) -> Iterator[sqlite3.Connection]:
        connection = sqlite3.connect(self.path, timeout=timeout, isolation_level=None)
        connection.row_factory = sqlite3.Row
        connection.execute("PRAGMA foreign_keys = ON")
        connection.execute(f"PRAGMA busy_timeout = {int(timeout * 1_000)}")
        try:
            yield connection
        finally:
            connection.close()

    def initialize(self) -> None:
        with self.connection() as connection:
            connection.execute("PRAGMA journal_mode = WAL")
            connection.execute("PRAGMA synchronous = NORMAL")
            connection.executescript(
                """
                CREATE TABLE IF NOT EXISTS jobs (
                    id TEXT PRIMARY KEY,
                    source_url TEXT NOT NULL,
                    normalized_url TEXT NOT NULL,
                    owner_token TEXT NOT NULL DEFAULT '',
                    config_json TEXT NOT NULL,
                    state TEXT NOT NULL,
                    phase TEXT NOT NULL,
                    message TEXT NOT NULL,
                    error_code TEXT,
                    error_message TEXT,
                    warnings_json TEXT NOT NULL DEFAULT '[]',
                    pages_found INTEGER NOT NULL DEFAULT 0,
                    pages_fetched INTEGER NOT NULL DEFAULT 0,
                    assets_found INTEGER NOT NULL DEFAULT 0,
                    assets_fetched INTEGER NOT NULL DEFAULT 0,
                    bytes_downloaded INTEGER NOT NULL DEFAULT 0,
                    files_written INTEGER NOT NULL DEFAULT 0,
                    artifact_path TEXT,
                    report_path TEXT,
                    archive_size INTEGER,
                    entry_point TEXT,
                    cancel_requested INTEGER NOT NULL DEFAULT 0,
                    worker_id TEXT,
                    lease_expires_at TEXT,
                    attempt_count INTEGER NOT NULL DEFAULT 0,
                    created_at TEXT NOT NULL,
                    updated_at TEXT NOT NULL,
                    started_at TEXT,
                    completed_at TEXT,
                    expires_at TEXT
                );

                CREATE INDEX IF NOT EXISTS jobs_state_created_idx
                    ON jobs(state, created_at);

                CREATE INDEX IF NOT EXISTS jobs_lease_idx
                    ON jobs(lease_expires_at);

                CREATE TABLE IF NOT EXISTS job_events (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    job_id TEXT NOT NULL REFERENCES jobs(id) ON DELETE CASCADE,
                    kind TEXT NOT NULL,
                    payload_json TEXT NOT NULL,
                    created_at TEXT NOT NULL
                );

                CREATE INDEX IF NOT EXISTS job_events_job_id_idx
                    ON job_events(job_id, id);
                """
            )
            columns = {
                str(row["name"])
                for row in connection.execute("PRAGMA table_info(jobs)").fetchall()
            }
            if "owner_token" not in columns:
                connection.execute(
                    "ALTER TABLE jobs ADD COLUMN owner_token TEXT NOT NULL DEFAULT ''"
                )
            connection.execute(
                "CREATE INDEX IF NOT EXISTS jobs_owner_created_idx "
                "ON jobs(owner_token, created_at DESC)"
            )
            cancelled_rows = connection.execute(
                "SELECT id FROM jobs WHERE state = 'running' AND cancel_requested = 1"
            ).fetchall()
            if cancelled_rows:
                now = timestamp()
                connection.execute(
                    """
                    UPDATE jobs
                    SET state = 'cancelled', phase = 'Cancelled',
                        message = 'Cancellation requested',
                        worker_id = NULL, lease_expires_at = NULL, completed_at = ?, updated_at = ?
                    WHERE state = 'running' AND cancel_requested = 1
                    """,
                    (now, now),
                )
                for row in cancelled_rows:
                    self._append_event(
                        connection,
                        str(row["id"]),
                        "state",
                        {
                            "state": "cancelled",
                            "phase": "Cancelled",
                            "message": "Cancellation requested",
                        },
                    )

    def create_job(
        self,
        source_url: str,
        normalized_url: str,
        config: dict[str, Any],
        owner_token: str = "",
    ) -> dict[str, Any]:
        job_id = f"cap_{uuid.uuid4().hex}"
        now = timestamp()
        payload = json.dumps(config, separators=(",", ":"))
        with self.connection() as connection:
            connection.execute(
                """
                INSERT INTO jobs (
                    id, source_url, normalized_url, owner_token, config_json, state, phase, message,
                    created_at, updated_at, expires_at
                ) VALUES (
                    ?, ?, ?, ?, ?, 'queued', 'Queued', 'Waiting for a local worker', ?, ?, NULL
                )
                """,
                (
                    job_id,
                    source_url,
                    normalized_url,
                    owner_token,
                    payload,
                    now,
                    now,
                ),
            )
            self._append_event(
                connection,
                job_id,
                "state",
                {"state": "queued", "phase": "Queued", "message": "Waiting for a local worker"},
            )
        return self.get_job(job_id) or {}

    def get_job(self, job_id: str) -> dict[str, Any] | None:
        with self.connection() as connection:
            row = connection.execute("SELECT * FROM jobs WHERE id = ?", (job_id,)).fetchone()
        return self._row_to_job(row) if row else None

    def get_owned_job(self, job_id: str, owner_token: str) -> dict[str, Any] | None:
        with self.connection() as connection:
            row = connection.execute(
                "SELECT * FROM jobs WHERE id = ? AND owner_token = ?", (job_id, owner_token)
            ).fetchone()
        return self._row_to_job(row) if row else None

    def list_jobs(self, limit: int = 30) -> list[dict[str, Any]]:
        with self.connection() as connection:
            rows = connection.execute(
                "SELECT * FROM jobs ORDER BY created_at DESC LIMIT ?", (limit,)
            ).fetchall()
        return [self._row_to_job(row) for row in rows]

    def list_owned_jobs(self, owner_token: str, limit: int = 30) -> list[dict[str, Any]]:
        with self.connection() as connection:
            rows = connection.execute(
                """
                SELECT * FROM jobs
                WHERE owner_token = ?
                ORDER BY created_at DESC
                LIMIT ?
                """,
                (owner_token, limit),
            ).fetchall()
        return [self._row_to_job(row) for row in rows]

    def public_stats(self) -> dict[str, int]:
        with self.connection() as connection:
            row = connection.execute(
                """
                SELECT
                    COUNT(CASE WHEN state = 'ready' THEN 1 END) AS websites_cloned,
                    COALESCE(SUM(bytes_downloaded), 0) AS bytes_scraped,
                    COUNT(CASE WHEN state IN ('queued', 'running') THEN 1 END) AS active_crawls,
                    COALESCE(SUM(files_written), 0) AS files_archived
                FROM jobs
                """
            ).fetchone()
        return {
            "websites_cloned": int(row["websites_cloned"] or 0),
            "bytes_scraped": int(row["bytes_scraped"] or 0),
            "active_crawls": int(row["active_crawls"] or 0),
            "files_archived": int(row["files_archived"] or 0),
        }

    def get_events(self, job_id: str, after_id: int = 0) -> list[dict[str, Any]]:
        with self.connection() as connection:
            rows = connection.execute(
                """
                SELECT id, kind, payload_json, created_at
                FROM job_events
                WHERE job_id = ? AND id > ?
                ORDER BY id ASC
                LIMIT 100
                """,
                (job_id, after_id),
            ).fetchall()
        return [
            {
                "id": row["id"],
                "kind": row["kind"],
                "payload": json.loads(row["payload_json"]),
                "created_at": row["created_at"],
            }
            for row in rows
        ]

    def claim_next(self, worker_id: str, lease_seconds: int) -> dict[str, Any] | None:
        now = utcnow()
        now_text = timestamp(now)
        lease = timestamp(now + timedelta(seconds=lease_seconds))
        with self.connection() as connection:
            try:
                connection.execute("BEGIN IMMEDIATE")
                cancelled_rows = connection.execute(
                    """
                    SELECT id FROM jobs
                    WHERE cancel_requested = 1
                      AND (
                        state = 'queued'
                        OR (
                            state = 'running'
                            AND lease_expires_at IS NOT NULL
                            AND lease_expires_at < ?
                        )
                      )
                    """,
                    (now_text,),
                ).fetchall()
                connection.execute(
                    """
                    UPDATE jobs
                    SET state = 'cancelled', phase = 'Cancelled',
                        message = 'Cancelled before starting',
                        completed_at = ?, updated_at = ?
                    WHERE cancel_requested = 1
                      AND (
                        state = 'queued'
                        OR (
                            state = 'running'
                            AND lease_expires_at IS NOT NULL
                            AND lease_expires_at < ?
                        )
                      )
                    """,
                    (now_text, now_text, now_text),
                )
                for cancelled_row in cancelled_rows:
                    self._append_event(
                        connection,
                        cancelled_row["id"],
                        "state",
                        {
                            "state": "cancelled",
                            "phase": "Cancelled",
                            "message": "Cancelled before starting",
                        },
                    )
                row = connection.execute(
                    """
                    SELECT * FROM jobs
                    WHERE cancel_requested = 0
                      AND (
                        state = 'queued'
                        OR (
                            state = 'running'
                            AND lease_expires_at IS NOT NULL
                            AND lease_expires_at < ?
                        )
                      )
                    ORDER BY created_at ASC
                    LIMIT 1
                    """,
                    (now_text,),
                ).fetchone()
                if not row:
                    connection.execute("COMMIT")
                    return None

                job_id = row["id"]
                message = (
                    "Recovering a leased capture"
                    if row["state"] == "running"
                    else "Preparing capture"
                )
                connection.execute(
                    """
                    UPDATE jobs
                    SET state = 'running', phase = 'Preparing', message = ?, worker_id = ?,
                        lease_expires_at = ?, attempt_count = attempt_count + 1,
                        started_at = COALESCE(started_at, ?), updated_at = ?
                    WHERE id = ?
                    """,
                    (message, worker_id, lease, now_text, now_text, job_id),
                )
                self._append_event(
                    connection,
                    job_id,
                    "state",
                    {"state": "running", "phase": "Preparing", "message": message},
                )
                claimed = connection.execute(
                    "SELECT * FROM jobs WHERE id = ?", (job_id,)
                ).fetchone()
                connection.execute("COMMIT")
                return self._row_to_job(claimed)
            except Exception:
                connection.execute("ROLLBACK")
                raise

    def update_progress(
        self,
        job_id: str,
        worker_id: str,
        lease_seconds: int,
        *,
        phase: str,
        message: str,
        stats: dict[str, int],
        warnings: list[str] | None = None,
    ) -> bool:
        now = utcnow()
        values = {
            "phase": phase,
            "message": message,
            "pages_found": stats.get("pages_found", 0),
            "pages_fetched": stats.get("pages_fetched", 0),
            "assets_found": stats.get("assets_found", 0),
            "assets_fetched": stats.get("assets_fetched", 0),
            "bytes_downloaded": stats.get("bytes_downloaded", 0),
            "files_written": stats.get("files_written", 0),
            "warnings_json": json.dumps(warnings or []),
            "worker_id": worker_id,
            "lease_expires_at": timestamp(now + timedelta(seconds=lease_seconds)),
            "updated_at": timestamp(now),
            "current_time": timestamp(now),
        }
        with self.connection() as connection:
            cursor = connection.execute(
                """
                UPDATE jobs SET
                    phase = :phase, message = :message,
                    pages_found = :pages_found, pages_fetched = :pages_fetched,
                    assets_found = :assets_found, assets_fetched = :assets_fetched,
                    bytes_downloaded = :bytes_downloaded, files_written = :files_written,
                    warnings_json = :warnings_json, worker_id = :worker_id,
                    lease_expires_at = :lease_expires_at, updated_at = :updated_at
                WHERE id = :job_id
                  AND state = 'running'
                  AND worker_id = :worker_id
                  AND lease_expires_at >= :current_time
                """,
                {**values, "job_id": job_id},
            )
            if cursor.rowcount:
                self._append_event(
                    connection,
                    job_id,
                    "progress",
                    {
                        "phase": phase,
                        "message": message,
                        "stats": stats,
                        "warnings": warnings or [],
                    },
                )
                return True
        return False

    def renew_lease(self, job_id: str, worker_id: str, lease_seconds: int) -> bool:
        now = utcnow()
        now_text = timestamp(now)
        with self.connection() as connection:
            cursor = connection.execute(
                """
                UPDATE jobs
                SET lease_expires_at = ?, updated_at = ?
                WHERE id = ?
                  AND state = 'running'
                  AND worker_id = ?
                  AND lease_expires_at >= ?
                """,
                (
                    timestamp(now + timedelta(seconds=lease_seconds)),
                    now_text,
                    job_id,
                    worker_id,
                    now_text,
                ),
            )
        return bool(cursor.rowcount)

    def is_cancel_requested(self, job_id: str) -> bool:
        with self.connection() as connection:
            row = connection.execute(
                "SELECT cancel_requested FROM jobs WHERE id = ?", (job_id,)
            ).fetchone()
        return bool(row and row["cancel_requested"])

    def owns_lease(self, job_id: str, worker_id: str) -> bool:
        now = timestamp()
        with self.connection() as connection:
            row = connection.execute(
                """
                SELECT state, worker_id, lease_expires_at
                FROM jobs
                WHERE id = ?
                """,
                (job_id,),
            ).fetchone()
        return bool(
            row
            and row["state"] == "running"
            and row["worker_id"] == worker_id
            and row["lease_expires_at"]
            and row["lease_expires_at"] >= now
        )

    def worker_should_stop(self, job_id: str, worker_id: str) -> bool:
        """Check cancellation and lease ownership without stalling a crawl on a busy database."""
        now = timestamp()
        try:
            with self.connection(timeout=0.5) as connection:
                row = connection.execute(
                    """
                    SELECT state, cancel_requested, worker_id, lease_expires_at
                    FROM jobs
                    WHERE id = ?
                    """,
                    (job_id,),
                ).fetchone()
        except sqlite3.OperationalError:
            return False
        return not bool(
            row
            and not row["cancel_requested"]
            and row["state"] == "running"
            and row["worker_id"] == worker_id
            and row["lease_expires_at"]
            and row["lease_expires_at"] >= now
        )

    def request_cancel(self, job_id: str) -> dict[str, Any] | None:
        now = timestamp()
        with self.connection() as connection:
            try:
                connection.execute("BEGIN IMMEDIATE")
                row = connection.execute(
                    "SELECT state FROM jobs WHERE id = ?", (job_id,)
                ).fetchone()
                if not row:
                    connection.execute("COMMIT")
                    return None
                if row["state"] in {"ready", "failed", "cancelled", "expired"}:
                    connection.execute("COMMIT")
                    return self.get_job(job_id)

                if row["state"] == "queued":
                    cursor = connection.execute(
                        """
                        UPDATE jobs
                        SET cancel_requested = 1, state = 'cancelled', phase = 'Cancelled',
                            message = 'Cancelled before starting', completed_at = ?, updated_at = ?
                        WHERE id = ? AND state = 'queued'
                        """,
                        (now, now, job_id),
                    )
                    payload = {
                        "state": "cancelled",
                        "phase": "Cancelled",
                        "message": "Cancelled before starting",
                    }
                else:
                    cursor = connection.execute(
                        """
                        UPDATE jobs
                        SET cancel_requested = 1, state = 'cancelled', phase = 'Cancelled',
                            message = 'Cancellation requested', worker_id = NULL,
                            lease_expires_at = NULL, completed_at = ?, updated_at = ?
                        WHERE id = ? AND state = 'running'
                        """,
                        (now, now, job_id),
                    )
                    payload = {
                        "state": "cancelled",
                        "phase": "Cancelled",
                        "message": "Cancellation requested",
                    }
                if cursor.rowcount:
                    self._append_event(connection, job_id, "state", payload)
                connection.execute("COMMIT")
            except Exception:
                connection.execute("ROLLBACK")
                raise
        return self.get_job(job_id)

    def mark_ready(
        self,
        job_id: str,
        worker_id: str,
        *,
        artifact_path: str,
        report_path: str,
        archive_size: int,
        entry_point: str,
        stats: dict[str, int],
        warnings: list[str],
        retention_days: int,
    ) -> bool:
        now_datetime = utcnow()
        now = timestamp(now_datetime)
        retention_days = min(retention_days, MAX_ARCHIVE_RETENTION_DAYS)
        with self.connection() as connection:
            cursor = connection.execute(
                """
                UPDATE jobs
                SET state = 'ready', phase = 'Ready', message = 'Archive is ready to download',
                    artifact_path = ?, report_path = ?, archive_size = ?, entry_point = ?,
                    pages_found = ?, pages_fetched = ?, assets_found = ?, assets_fetched = ?,
                    bytes_downloaded = ?, files_written = ?, warnings_json = ?,
                    worker_id = NULL, lease_expires_at = NULL, expires_at = ?,
                    completed_at = ?, updated_at = ?
                WHERE id = ?
                  AND state = 'running'
                  AND worker_id = ?
                  AND cancel_requested = 0
                  AND lease_expires_at >= ?
                """,
                (
                    artifact_path,
                    report_path,
                    archive_size,
                    entry_point,
                    stats.get("pages_found", 0),
                    stats.get("pages_fetched", 0),
                    stats.get("assets_found", 0),
                    stats.get("assets_fetched", 0),
                    stats.get("bytes_downloaded", 0),
                    stats.get("files_written", 0),
                    json.dumps(warnings),
                    timestamp(now_datetime + timedelta(days=retention_days)),
                    now,
                    now,
                    job_id,
                    worker_id,
                    now,
                ),
            )
            if cursor.rowcount:
                self._append_event(
                    connection,
                    job_id,
                    "state",
                    {"state": "ready", "phase": "Ready", "message": "Archive is ready to download"},
                )
                return True
        return False

    def mark_cancelled(
        self, job_id: str, worker_id: str, message: str = "Capture cancelled"
    ) -> bool:
        now = timestamp()
        with self.connection() as connection:
            cursor = connection.execute(
                """
                UPDATE jobs
                SET state = 'cancelled', phase = 'Cancelled', message = ?, worker_id = NULL,
                    lease_expires_at = NULL, completed_at = ?, updated_at = ?
                WHERE id = ?
                  AND state = 'running'
                  AND worker_id = ?
                  AND cancel_requested = 1
                  AND lease_expires_at >= ?
                """,
                (message, now, now, job_id, worker_id, now),
            )
            if cursor.rowcount:
                self._append_event(
                    connection,
                    job_id,
                    "state",
                    {"state": "cancelled", "phase": "Cancelled", "message": message},
                )
                return True
        return False

    def mark_failed(self, job_id: str, worker_id: str, code: str, message: str) -> bool:
        now = timestamp()
        with self.connection() as connection:
            cursor = connection.execute(
                """
                UPDATE jobs
                SET state = 'failed', phase = 'Failed', message = ?, error_code = ?,
                    error_message = ?, worker_id = NULL, lease_expires_at = NULL,
                    completed_at = ?, updated_at = ?
                WHERE id = ?
                  AND state = 'running'
                  AND worker_id = ?
                  AND cancel_requested = 0
                  AND lease_expires_at >= ?
                """,
                (message, code, message, now, now, job_id, worker_id, now),
            )
            if cursor.rowcount:
                self._append_event(
                    connection,
                    job_id,
                    "state",
                    {"state": "failed", "phase": "Failed", "message": message, "code": code},
                )
                return True
        return False

    def mark_expired(self, job_id: str) -> None:
        now = timestamp()
        with self.connection() as connection:
            connection.execute(
                """
                UPDATE jobs
                SET state = 'expired', phase = 'Expired',
                    message = 'Artifact retention period ended',
                    artifact_path = NULL, report_path = NULL, updated_at = ?
                WHERE id = ?
                """,
                (now, job_id),
            )
            self._append_event(
                connection,
                job_id,
                "state",
                {
                    "state": "expired",
                    "phase": "Expired",
                    "message": "Artifact retention period ended",
                },
            )

    def mark_deleted(self, job_id: str) -> None:
        now = timestamp()
        with self.connection() as connection:
            cursor = connection.execute(
                """
                UPDATE jobs
                SET state = 'expired', phase = 'Deleted', message = 'Artifact deleted by user',
                    artifact_path = NULL, report_path = NULL, updated_at = ?
                WHERE id = ? AND state = 'ready'
                """,
                (now, job_id),
            )
            if cursor.rowcount:
                self._append_event(
                    connection,
                    job_id,
                    "state",
                    {
                        "state": "expired",
                        "phase": "Deleted",
                        "message": "Artifact deleted by user",
                    },
                )

    def expired_ready_jobs(self) -> list[dict[str, Any]]:
        now_datetime = utcnow()
        now = timestamp(now_datetime)
        maximum_completion_age = timestamp(
            now_datetime - timedelta(days=MAX_ARCHIVE_RETENTION_DAYS)
        )
        with self.connection() as connection:
            rows = connection.execute(
                """
                SELECT * FROM jobs
                WHERE state = 'ready'
                  AND (
                    expires_at <= ?
                    OR COALESCE(completed_at, created_at) <= ?
                  )
                """,
                (now, maximum_completion_age),
            ).fetchall()
        return [self._row_to_job(row) for row in rows]

    def referenced_data_paths(self) -> set[str]:
        with self.connection() as connection:
            rows = connection.execute(
                """
                SELECT artifact_path, report_path
                FROM jobs
                WHERE state = 'ready'
                """
            ).fetchall()
        return {
            path
            for row in rows
            for path in (row["artifact_path"], row["report_path"])
            if isinstance(path, str) and path
        }

    def _append_event(
        self, connection: sqlite3.Connection, job_id: str, kind: str, payload: dict[str, Any]
    ) -> None:
        connection.execute(
            "INSERT INTO job_events (job_id, kind, payload_json, created_at) VALUES (?, ?, ?, ?)",
            (job_id, kind, json.dumps(payload, separators=(",", ":")), timestamp()),
        )

    @staticmethod
    def _row_to_job(row: sqlite3.Row) -> dict[str, Any]:
        data = dict(row)
        data["config"] = json.loads(data.pop("config_json"))
        data["warnings"] = json.loads(data.pop("warnings_json"))
        data["cancel_requested"] = bool(data["cancel_requested"])
        return data