"""Process-shared live snapshot state and refresh coordination."""

from __future__ import annotations

import json
import os
import secrets
import sqlite3
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Callable, Iterable

from src.private_state import connect_private_sqlite
from src.state_encryption import (
    PrivateStateCipher,
    PrivateStateEncryptionError,
    private_state_encryption_key,
)

SECTIONS = ("context", "events", "notifications")


class RefreshLeaseLost(RuntimeError):
    """The refresh is no longer authorized to publish shared state."""


@dataclass(frozen=True)
class LiveSnapshotState:
    value: dict | None
    created_at: dict[str, float | None]
    failure_count: dict[str, int]
    retry_at: dict[str, float | None]
    revisions: dict[str, int]
    generation: str
    refreshing_sections: set[str]
    lease_expires_at: float | None


@dataclass(frozen=True)
class LiveSnapshotMetadata:
    """Shared snapshot coordination fields without the potentially large value."""

    created_at: dict[str, float | None]
    failure_count: dict[str, int]
    retry_at: dict[str, float | None]
    revisions: dict[str, int]
    generation: str
    refreshing_sections: set[str]
    lease_expires_at: float | None


class LiveSnapshotStore:
    """A private SQLite snapshot with an expiring, cross-process refresh lease."""

    def __init__(
        self,
        path: str | os.PathLike[str],
        *,
        clock: Callable[[], float] | None = None,
        encryption_key: bytes | None = None,
    ):
        self.path = Path(path)
        self.clock = clock or time.time
        self._cipher = PrivateStateCipher(
            encryption_key if encryption_key is not None else private_state_encryption_key(),
            store="live-snapshot",
        )
        self._initialize()

    def _connect(self) -> sqlite3.Connection:
        connection = connect_private_sqlite(self.path, timeout=1.0, isolation_level=None)
        connection.row_factory = sqlite3.Row
        connection.execute("PRAGMA busy_timeout = 1000")
        return connection

    def _initialize(self) -> None:
        with self._connect() as connection:
            connection.executescript(
                """
                CREATE TABLE IF NOT EXISTS live_snapshot (
                    singleton INTEGER PRIMARY KEY CHECK (singleton = 1),
                    value_json TEXT,
                    created_at_json TEXT NOT NULL,
                    failure_count_json TEXT NOT NULL,
                    retry_at_json TEXT NOT NULL,
                    revisions_json TEXT NOT NULL,
                    generation TEXT NOT NULL
                );
                CREATE TABLE IF NOT EXISTS live_refresh_lease (
                    singleton INTEGER PRIMARY KEY CHECK (singleton = 1),
                    owner TEXT NOT NULL,
                    sections_json TEXT NOT NULL,
                    expires_at REAL NOT NULL
                );
                CREATE TABLE IF NOT EXISTS live_read_notification (
                    notification_id INTEGER PRIMARY KEY
                );
                """
            )
            connection.execute(
                """INSERT OR IGNORE INTO live_snapshot VALUES
                (1, NULL, ?, ?, ?, ?, ?)""",
                (
                    json.dumps({section: None for section in SECTIONS}),
                    json.dumps({section: 0 for section in SECTIONS}),
                    json.dumps({section: None for section in SECTIONS}),
                    json.dumps({section: 0 for section in SECTIONS}),
                    secrets.token_hex(8),
                ),
            )

    def try_acquire_refresh(
        self, sections: Iterable[str], *, lease_seconds: float
    ) -> str | None:
        requested = sorted(set(sections))
        if not requested or any(section not in SECTIONS for section in requested):
            raise ValueError("refresh lease requires known sections")
        if lease_seconds <= 0:
            raise ValueError("refresh lease duration must be positive")
        now = self.clock()
        owner = secrets.token_hex(16)
        with self._connect() as connection:
            connection.execute("BEGIN IMMEDIATE")
            active = connection.execute(
                "SELECT expires_at FROM live_refresh_lease WHERE singleton = 1"
            ).fetchone()
            if active is not None and active["expires_at"] > now:
                connection.rollback()
                return None
            connection.execute("DELETE FROM live_refresh_lease WHERE singleton = 1")
            connection.execute(
                "INSERT INTO live_refresh_lease VALUES (1, ?, ?, ?)",
                (owner, json.dumps(requested, separators=(",", ":")), now + lease_seconds),
            )
            connection.commit()
        return owner

    def load(self) -> LiveSnapshotState:
        now = self.clock()
        with self._connect() as connection:
            row = connection.execute(
                "SELECT * FROM live_snapshot WHERE singleton = 1"
            ).fetchone()
            lease = connection.execute(
                "SELECT sections_json, expires_at FROM live_refresh_lease "
                "WHERE singleton = 1 AND expires_at > ?",
                (now,),
            ).fetchone()
        assert row is not None
        value = None
        if row["value_json"] is not None:
            value, legacy = self._cipher.open(
                row["value_json"], binding=row["generation"]
            )
            if not isinstance(value, dict):
                raise PrivateStateEncryptionError("private state could not be decrypted")
            if legacy:
                migrated = self._cipher.seal(value, binding=row["generation"])
                with self._connect() as connection:
                    connection.execute(
                        "UPDATE live_snapshot SET value_json = ? "
                        "WHERE singleton = 1 AND value_json = ?",
                        (migrated, row["value_json"]),
                    )
        return LiveSnapshotState(
            value=value,
            created_at=json.loads(row["created_at_json"]),
            failure_count=json.loads(row["failure_count_json"]),
            retry_at=json.loads(row["retry_at_json"]),
            revisions=json.loads(row["revisions_json"]),
            generation=row["generation"],
            refreshing_sections=set(json.loads(lease["sections_json"])) if lease else set(),
            lease_expires_at=lease["expires_at"] if lease else None,
        )

    def load_metadata(self) -> LiveSnapshotMetadata:
        """Load freshness and coordination state without reading the snapshot value."""
        now = self.clock()
        with self._connect() as connection:
            row = connection.execute(
                """SELECT created_at_json, failure_count_json, retry_at_json,
                revisions_json, generation FROM live_snapshot WHERE singleton = 1"""
            ).fetchone()
            lease = connection.execute(
                "SELECT sections_json, expires_at FROM live_refresh_lease "
                "WHERE singleton = 1 AND expires_at > ?",
                (now,),
            ).fetchone()
        assert row is not None
        return LiveSnapshotMetadata(
            created_at=json.loads(row["created_at_json"]),
            failure_count=json.loads(row["failure_count_json"]),
            retry_at=json.loads(row["retry_at_json"]),
            revisions=json.loads(row["revisions_json"]),
            generation=row["generation"],
            refreshing_sections=set(json.loads(lease["sections_json"])) if lease else set(),
            lease_expires_at=lease["expires_at"] if lease else None,
        )

    def publish_refresh(
        self,
        owner: str,
        *,
        value: dict,
        created_at: dict[str, float | None],
        failure_count: dict[str, int],
        retry_at: dict[str, float | None],
        changed_sections: Iterable[str],
    ) -> LiveSnapshotState:
        changed = set(changed_sections)
        if changed - set(SECTIONS):
            raise ValueError("publication contains unknown sections")
        if any(set(mapping) != set(SECTIONS) for mapping in (created_at, failure_count, retry_at)):
            raise ValueError("publication metadata must contain every section")
        now = self.clock()
        with self._connect() as connection:
            connection.execute("BEGIN IMMEDIATE")
            lease = connection.execute(
                "SELECT owner, expires_at FROM live_refresh_lease WHERE singleton = 1"
            ).fetchone()
            if lease is None or lease["owner"] != owner or lease["expires_at"] <= now:
                connection.rollback()
                raise RefreshLeaseLost("live refresh lease expired or changed owner")
            row = connection.execute(
                "SELECT revisions_json, generation FROM live_snapshot WHERE singleton = 1"
            ).fetchone()
            revisions = json.loads(row["revisions_json"])
            notifications = value.get("notifications")
            if isinstance(notifications, list):
                returned_ids = {
                    item.get("id") for item in notifications if isinstance(item, dict)
                }
                read_ids = {
                    item["notification_id"]
                    for item in connection.execute(
                        "SELECT notification_id FROM live_read_notification"
                    )
                }
                value = dict(value)
                value["notifications"] = [
                    item for item in notifications
                    if not isinstance(item, dict) or item.get("id") not in read_ids
                ]
                if returned_ids:
                    placeholders = ",".join("?" for _ in returned_ids)
                    connection.execute(
                        f"DELETE FROM live_read_notification WHERE notification_id NOT IN ({placeholders})",
                        tuple(returned_ids),
                    )
                else:
                    connection.execute("DELETE FROM live_read_notification")
            for section in changed:
                revisions[section] += 1
            connection.execute(
                """UPDATE live_snapshot SET value_json = ?, created_at_json = ?,
                failure_count_json = ?, retry_at_json = ?, revisions_json = ?
                WHERE singleton = 1""",
                (
                    self._cipher.seal(value, binding=row["generation"]),
                    json.dumps(created_at, separators=(",", ":")),
                    json.dumps(failure_count, separators=(",", ":")),
                    json.dumps(retry_at, separators=(",", ":")),
                    json.dumps(revisions, separators=(",", ":")),
                ),
            )
            connection.execute(
                "DELETE FROM live_refresh_lease WHERE singleton = 1 AND owner = ?", (owner,)
            )
            connection.commit()
        return self.load()

    def remove_notifications(self, notification_ids: Iterable[int]) -> LiveSnapshotState:
        """Filter confirmed reads from every worker and subsequent stale refreshes."""
        read_ids = set(notification_ids)
        if not read_ids:
            return self.load()
        with self._connect() as connection:
            connection.execute("BEGIN IMMEDIATE")
            connection.executemany(
                "INSERT OR IGNORE INTO live_read_notification VALUES (?)",
                ((notification_id,) for notification_id in read_ids),
            )
            row = connection.execute(
                "SELECT value_json, revisions_json, generation FROM live_snapshot WHERE singleton = 1"
            ).fetchone()
            value = self._cipher.open(
                row["value_json"], binding=row["generation"]
            )[0] if row["value_json"] else None
            revisions = json.loads(row["revisions_json"])
            if value is not None and isinstance(value.get("notifications"), list):
                previous = value["notifications"]
                retained = [
                    item
                    for item in previous
                    if not isinstance(item, dict) or item.get("id") not in read_ids
                ]
                if retained != previous:
                    value["notifications"] = retained
                    revisions["notifications"] += 1
                    connection.execute(
                        "UPDATE live_snapshot SET value_json = ?, revisions_json = ? WHERE singleton = 1",
                        (
                            self._cipher.seal(value, binding=row["generation"]),
                            json.dumps(revisions, separators=(",", ":")),
                        ),
                    )
            connection.commit()
        return self.load()

    def remove_notification(self, notification_id: int) -> LiveSnapshotState:
        """Filter one confirmed read from every worker and subsequent stale refresh."""
        return self.remove_notifications([notification_id])

    def release_refresh(self, owner: str) -> None:
        """Relinquish a lease that became unnecessary after a coherent recheck."""
        with self._connect() as connection:
            connection.execute(
                "DELETE FROM live_refresh_lease WHERE singleton = 1 AND owner = ?",
                (owner,),
            )

    def fail_refresh(
        self,
        owner: str,
        *,
        failure_count: dict[str, int],
        retry_at: dict[str, float | None],
    ) -> LiveSnapshotState:
        """Publish backoff metadata and relinquish a failed refresh atomically."""
        now = self.clock()
        with self._connect() as connection:
            connection.execute("BEGIN IMMEDIATE")
            lease = connection.execute(
                "SELECT owner, expires_at FROM live_refresh_lease WHERE singleton = 1"
            ).fetchone()
            if lease is None or lease["owner"] != owner or lease["expires_at"] <= now:
                connection.rollback()
                raise RefreshLeaseLost("live refresh lease expired or changed owner")
            connection.execute(
                """UPDATE live_snapshot SET failure_count_json = ?, retry_at_json = ?
                WHERE singleton = 1""",
                (
                    json.dumps(failure_count, separators=(",", ":")),
                    json.dumps(retry_at, separators=(",", ":")),
                ),
            )
            connection.execute("DELETE FROM live_refresh_lease WHERE singleton = 1")
            connection.commit()
        return self.load()
