import json
import sqlite3
from collections.abc import Iterable, Mapping
from dataclasses import dataclass
from datetime import datetime, timezone
from os import PathLike
from typing import Any

from .adapters import (
    NormalizedObject,
    ObjectReference,
    Observation,
)


_FORBIDDEN_PERSISTED_KEYS = frozenset(
    {
        "private_key",
        "certificate_key",
        "certificate",
        "dns_provider_credentials",
        "secret",
    }
)


@dataclass(frozen=True)
class ObservationStatus:
    endpoint_id: str
    object_type: str
    last_attempt_at: datetime
    last_success_at: datetime | None
    last_error: str | None

    @property
    def is_live(self) -> bool:
        return self.last_error is None


@dataclass(frozen=True)
class DesiredState:
    object: NormalizedObject
    selected_at: datetime


@dataclass(frozen=True)
class AppliedState:
    endpoint_id: str
    object: NormalizedObject
    applied_at: datetime
    verification_attempted_at: datetime | None
    verified_at: datetime | None
    verification_error: str | None

    @property
    def is_verified(self) -> bool:
        return (
            self.verified_at is not None
            and self.verification_error is None
        )


def _utc(value: datetime, label: str) -> datetime:
    if value.tzinfo is None:
        raise RuntimeError(f"{label} must be timezone-aware")
    return value.astimezone(timezone.utc)


def _timestamp(value: datetime, label: str) -> str:
    return _utc(value, label).isoformat()


def _parse_timestamp(value: str | None) -> datetime | None:
    if value is None:
        return None
    parsed = datetime.fromisoformat(value)
    if parsed.tzinfo is None:
        raise RuntimeError(
            f"Persisted timestamp is not timezone-aware: {value!r}"
        )
    return parsed.astimezone(timezone.utc)


def _validate_safe_payload(value: Any, path: str = "payload") -> None:
    if isinstance(value, Mapping):
        for raw_key, child in value.items():
            key = str(raw_key).strip().lower().replace("-", "_")
            if key in _FORBIDDEN_PERSISTED_KEYS:
                raise RuntimeError(
                    f"Refusing to persist sensitive field {path}.{key}"
                )
            _validate_safe_payload(child, f"{path}.{key}")
    elif isinstance(value, (list, tuple)):
        for index, child in enumerate(value):
            _validate_safe_payload(child, f"{path}[{index}]")


def _serialize_payload(payload: Mapping[str, Any]) -> str:
    _validate_safe_payload(payload)
    try:
        return json.dumps(
            payload,
            sort_keys=True,
            separators=(",", ":"),
        )
    except (TypeError, ValueError) as exc:
        raise RuntimeError(
            f"Object payload is not JSON serializable: {exc}"
        ) from exc


def _deserialize_payload(value: str) -> dict[str, Any]:
    try:
        payload = json.loads(value)
    except json.JSONDecodeError as exc:
        raise RuntimeError("Persisted object payload is invalid JSON") from exc

    if not isinstance(payload, dict):
        raise RuntimeError("Persisted object payload is not a mapping")
    return payload


def _serialize_references(
    references: Iterable[ObjectReference],
) -> str:
    return json.dumps(
        [
            {
                "relation": reference.relation,
                "object_type": reference.object_type,
                "local_key": reference.local_key,
                "logical_key": reference.logical_key,
            }
            for reference in references
        ],
        sort_keys=True,
        separators=(",", ":"),
    )


def _deserialize_references(
    value: str,
) -> tuple[ObjectReference, ...]:
    try:
        serialized = json.loads(value)
    except json.JSONDecodeError as exc:
        raise RuntimeError(
            "Persisted object references are invalid JSON"
        ) from exc

    if not isinstance(serialized, list):
        raise RuntimeError(
            "Persisted object references are not a list"
        )

    expected_fields = {
        "relation",
        "object_type",
        "local_key",
        "logical_key",
    }
    references: list[ObjectReference] = []
    for item in serialized:
        if (
            not isinstance(item, dict)
            or set(item) != expected_fields
        ):
            raise RuntimeError(
                "Persisted object reference has invalid fields"
            )
        try:
            references.append(
                ObjectReference(
                    relation=item["relation"],
                    object_type=item["object_type"],
                    local_key=item["local_key"],
                    logical_key=item["logical_key"],
                )
            )
        except (TypeError, ValueError) as exc:
            raise RuntimeError(
                "Persisted object reference is invalid"
            ) from exc

    return tuple(references)


class SQLiteStateStore:
    def __init__(self, path: str | PathLike[str]):
        self.path = str(path)
        self.connection = sqlite3.connect(self.path)
        self.connection.row_factory = sqlite3.Row
        self.connection.execute("PRAGMA foreign_keys = ON")
        self.connection.execute("PRAGMA busy_timeout = 5000")
        if self.path != ":memory:":
            self.connection.execute("PRAGMA journal_mode = WAL")
        self._initialize()

    def _initialize(self) -> None:
        with self.connection:
            self.connection.executescript(
                """
                CREATE TABLE IF NOT EXISTS schema_meta (
                    version INTEGER NOT NULL
                );

                INSERT INTO schema_meta(version)
                SELECT 1
                WHERE NOT EXISTS (
                    SELECT 1 FROM schema_meta
                );

                CREATE TABLE IF NOT EXISTS observations (
                    endpoint_id TEXT NOT NULL,
                    object_type TEXT NOT NULL,
                    logical_key TEXT NOT NULL,
                    local_key TEXT NOT NULL,
                    payload_json TEXT NOT NULL,
                    observed_at TEXT NOT NULL,
                    PRIMARY KEY (
                        endpoint_id,
                        object_type,
                        local_key
                    )
                );

                CREATE INDEX IF NOT EXISTS
                    observations_by_logical_identity
                ON observations (
                    object_type,
                    logical_key
                );

                CREATE TABLE IF NOT EXISTS observation_status (
                    endpoint_id TEXT NOT NULL,
                    object_type TEXT NOT NULL,
                    last_attempt_at TEXT NOT NULL,
                    last_success_at TEXT,
                    last_error TEXT,
                    PRIMARY KEY (
                        endpoint_id,
                        object_type
                    )
                );
                """
            )

            versions = self.connection.execute(
                "SELECT version FROM schema_meta"
            ).fetchall()
            if len(versions) != 1:
                raise RuntimeError(
                    "SQLite state schema has invalid version metadata"
                )

            version = versions[0]["version"]
            if version not in (1, 2, 3):
                raise RuntimeError(
                    "Unsupported SQLite state schema version: "
                    f"{version!r}"
                )

            if version == 1:
                self.connection.executescript(
                    """
                    CREATE TABLE desired_state (
                        object_type TEXT NOT NULL,
                        logical_key TEXT NOT NULL,
                        payload_json TEXT NOT NULL,
                        selected_at TEXT NOT NULL,
                        PRIMARY KEY (
                            object_type,
                            logical_key
                        )
                    );

                    CREATE TABLE applied_state (
                        endpoint_id TEXT NOT NULL,
                        object_type TEXT NOT NULL,
                        logical_key TEXT NOT NULL,
                        local_key TEXT,
                        payload_json TEXT NOT NULL,
                        applied_at TEXT NOT NULL,
                        verification_attempted_at TEXT,
                        verified_at TEXT,
                        verification_error TEXT,
                        PRIMARY KEY (
                            endpoint_id,
                            object_type,
                            logical_key
                        )
                    );

                    CREATE INDEX applied_state_by_object
                    ON applied_state (
                        object_type,
                        logical_key
                    );

                    UPDATE schema_meta
                    SET version = 2;
                    """
                )
                version = 2

            if version == 2:
                migrations = (
                    ("observations", "references_json"),
                    ("desired_state", "references_json"),
                    ("applied_state", "references_json"),
                )
                for table, column in migrations:
                    existing_columns = {
                        row["name"]
                        for row in self.connection.execute(
                            f"PRAGMA table_info({table})"
                        )
                    }
                    if column not in existing_columns:
                        self.connection.execute(
                            f"ALTER TABLE {table} "
                            f"ADD COLUMN {column} TEXT NOT NULL "
                            "DEFAULT '[]'"
                        )

                self.connection.execute(
                    "UPDATE schema_meta SET version = 3"
                )

    def close(self) -> None:
        self.connection.close()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc, traceback):
        self.close()

    def replace_observations(
        self,
        endpoint_id: str,
        object_type: str,
        observations: Iterable[Observation],
        snapshot_at: datetime,
    ) -> None:
        rows = list(observations)
        snapshot_timestamp = _timestamp(
            snapshot_at,
            "Observation snapshot timestamp",
        )
        prepared: list[
            tuple[str, str, str, str, str, str, str]
        ] = []
        local_keys: set[str] = set()

        for observation in rows:
            normalized = observation.object
            if observation.endpoint_id != endpoint_id:
                raise RuntimeError(
                    "Observation endpoint does not match snapshot endpoint"
                )
            if normalized.object_type != object_type:
                raise RuntimeError(
                    "Observation object type does not match snapshot type"
                )
            if normalized.local_key is None:
                raise RuntimeError(
                    "Observed objects require an endpoint-local key"
                )
            if normalized.local_key in local_keys:
                raise RuntimeError(
                    f"Duplicate endpoint-local key: {normalized.local_key}"
                )

            local_keys.add(normalized.local_key)
            prepared.append(
                (
                    endpoint_id,
                    object_type,
                    normalized.logical_key,
                    normalized.local_key,
                    _serialize_payload(normalized.payload),
                    _serialize_references(normalized.references),
                    _timestamp(
                        observation.observed_at,
                        "Observation timestamp",
                    ),
                )
            )

        with self.connection:
            self.connection.execute(
                """
                DELETE FROM observations
                WHERE endpoint_id = ? AND object_type = ?
                """,
                (endpoint_id, object_type),
            )
            self.connection.executemany(
                """
                INSERT INTO observations (
                    endpoint_id,
                    object_type,
                    logical_key,
                    local_key,
                    payload_json,
                    references_json,
                    observed_at
                ) VALUES (?, ?, ?, ?, ?, ?, ?)
                """,
                prepared,
            )
            self.connection.execute(
                """
                INSERT INTO observation_status (
                    endpoint_id,
                    object_type,
                    last_attempt_at,
                    last_success_at,
                    last_error
                ) VALUES (?, ?, ?, ?, NULL)
                ON CONFLICT(endpoint_id, object_type) DO UPDATE SET
                    last_attempt_at = excluded.last_attempt_at,
                    last_success_at = excluded.last_success_at,
                    last_error = NULL
                """,
                (
                    endpoint_id,
                    object_type,
                    snapshot_timestamp,
                    snapshot_timestamp,
                ),
            )

    def record_discovery_failure(
        self,
        endpoint_id: str,
        object_type: str,
        error: str,
        attempted_at: datetime,
    ) -> None:
        attempted_timestamp = _timestamp(
            attempted_at,
            "Discovery attempt timestamp",
        )
        with self.connection:
            self.connection.execute(
                """
                INSERT INTO observation_status (
                    endpoint_id,
                    object_type,
                    last_attempt_at,
                    last_success_at,
                    last_error
                ) VALUES (?, ?, ?, NULL, ?)
                ON CONFLICT(endpoint_id, object_type) DO UPDATE SET
                    last_attempt_at = excluded.last_attempt_at,
                    last_error = excluded.last_error
                """,
                (
                    endpoint_id,
                    object_type,
                    attempted_timestamp,
                    str(error),
                ),
            )

    def load_observations(
        self,
        object_type: str | None = None,
    ) -> list[Observation]:
        sql = """
            SELECT
                endpoint_id,
                object_type,
                logical_key,
                local_key,
                payload_json,
                references_json,
                observed_at
            FROM observations
        """
        parameters: tuple[str, ...] = ()
        if object_type is not None:
            sql += " WHERE object_type = ?"
            parameters = (object_type,)
        sql += """
            ORDER BY
                endpoint_id,
                object_type,
                logical_key,
                local_key
        """

        rows = self.connection.execute(sql, parameters).fetchall()
        return [
            Observation(
                endpoint_id=row["endpoint_id"],
                object=NormalizedObject(
                    object_type=row["object_type"],
                    logical_key=row["logical_key"],
                    local_key=row["local_key"],
                    payload=_deserialize_payload(row["payload_json"]),
                    references=_deserialize_references(
                        row["references_json"]
                    ),
                ),
                observed_at=_parse_timestamp(row["observed_at"]),
            )
            for row in rows
        ]

    def replace_desired_state(
        self,
        object_type: str,
        objects: Iterable[NormalizedObject],
        selected_at: datetime,
    ) -> None:
        rows = list(objects)
        selected_timestamp = _timestamp(
            selected_at,
            "Desired-state selection timestamp",
        )
        prepared: list[
            tuple[str, str, str, str, str]
        ] = []
        logical_keys: set[str] = set()

        for normalized in rows:
            if normalized.object_type != object_type:
                raise RuntimeError(
                    "Desired object type does not match snapshot type"
                )
            if normalized.local_key is not None:
                raise RuntimeError(
                    "Canonical desired objects cannot contain "
                    "an endpoint-local key"
                )
            if normalized.logical_key in logical_keys:
                raise RuntimeError(
                    "Duplicate desired logical key: "
                    f"{normalized.logical_key}"
                )

            logical_keys.add(normalized.logical_key)
            prepared.append(
                (
                    object_type,
                    normalized.logical_key,
                    _serialize_payload(normalized.payload),
                    _serialize_references(normalized.references),
                    selected_timestamp,
                )
            )

        with self.connection:
            self.connection.execute(
                """
                DELETE FROM desired_state
                WHERE object_type = ?
                """,
                (object_type,),
            )
            self.connection.executemany(
                """
                INSERT INTO desired_state (
                    object_type,
                    logical_key,
                    payload_json,
                    references_json,
                    selected_at
                ) VALUES (?, ?, ?, ?, ?)
                """,
                prepared,
            )

    def load_desired_state(
        self,
        object_type: str | None = None,
    ) -> list[DesiredState]:
        sql = """
            SELECT
                object_type,
                logical_key,
                payload_json,
                references_json,
                selected_at
            FROM desired_state
        """
        parameters: tuple[str, ...] = ()
        if object_type is not None:
            sql += " WHERE object_type = ?"
            parameters = (object_type,)
        sql += " ORDER BY object_type, logical_key"

        rows = self.connection.execute(sql, parameters).fetchall()
        return [
            DesiredState(
                object=NormalizedObject(
                    object_type=row["object_type"],
                    logical_key=row["logical_key"],
                    payload=_deserialize_payload(row["payload_json"]),
                    references=_deserialize_references(
                        row["references_json"]
                    ),
                ),
                selected_at=_parse_timestamp(row["selected_at"]),
            )
            for row in rows
        ]

    def record_applied(
        self,
        endpoint_id: str,
        object: NormalizedObject,
        applied_at: datetime,
    ) -> None:
        if not endpoint_id:
            raise RuntimeError("Applied-state endpoint_id cannot be empty")

        payload_json = _serialize_payload(object.payload)
        references_json = _serialize_references(
            object.references
        )
        applied_timestamp = _timestamp(
            applied_at,
            "Application timestamp",
        )

        with self.connection:
            self.connection.execute(
                """
                INSERT INTO applied_state (
                    endpoint_id,
                    object_type,
                    logical_key,
                    local_key,
                    payload_json,
                    references_json,
                    applied_at,
                    verification_attempted_at,
                    verified_at,
                    verification_error
                ) VALUES (?, ?, ?, ?, ?, ?, ?, NULL, NULL, NULL)
                ON CONFLICT(
                    endpoint_id,
                    object_type,
                    logical_key
                ) DO UPDATE SET
                    local_key = excluded.local_key,
                    payload_json = excluded.payload_json,
                    references_json = excluded.references_json,
                    applied_at = excluded.applied_at,
                    verification_attempted_at = NULL,
                    verified_at = NULL,
                    verification_error = NULL
                """,
                (
                    endpoint_id,
                    object.object_type,
                    object.logical_key,
                    object.local_key,
                    payload_json,
                    references_json,
                    applied_timestamp,
                ),
            )

    def _record_verification(
        self,
        endpoint_id: str,
        object: NormalizedObject,
        attempted_at: datetime,
        *,
        succeeded: bool,
        error: str | None,
    ) -> None:
        attempted_timestamp = _timestamp(
            attempted_at,
            "Verification timestamp",
        )
        payload_json = _serialize_payload(object.payload)
        references_json = _serialize_references(
            object.references
        )

        if succeeded:
            verified_at = attempted_timestamp
            verification_error = None
        else:
            verified_at = None
            verification_error = str(
                error or "verification failed"
            )

        with self.connection:
            cursor = self.connection.execute(
                """
                UPDATE applied_state
                SET
                    verification_attempted_at = ?,
                    verified_at = COALESCE(?, verified_at),
                    verification_error = ?
                WHERE
                    endpoint_id = ?
                    AND object_type = ?
                    AND logical_key = ?
                    AND payload_json = ?
                    AND references_json = ?
                """,
                (
                    attempted_timestamp,
                    verified_at,
                    verification_error,
                    endpoint_id,
                    object.object_type,
                    object.logical_key,
                    payload_json,
                    references_json,
                ),
            )
            if cursor.rowcount != 1:
                raise RuntimeError(
                    "Verification does not match the currently "
                    "applied object"
                )

    def record_verification_success(
        self,
        endpoint_id: str,
        object: NormalizedObject,
        verified_at: datetime,
    ) -> None:
        self._record_verification(
            endpoint_id,
            object,
            verified_at,
            succeeded=True,
            error=None,
        )

    def record_verification_failure(
        self,
        endpoint_id: str,
        object: NormalizedObject,
        error: str,
        attempted_at: datetime,
    ) -> None:
        self._record_verification(
            endpoint_id,
            object,
            attempted_at,
            succeeded=False,
            error=error,
        )

    def load_applied_state(
        self,
        endpoint_id: str | None = None,
        object_type: str | None = None,
    ) -> list[AppliedState]:
        clauses: list[str] = []
        parameters: list[str] = []

        if endpoint_id is not None:
            clauses.append("endpoint_id = ?")
            parameters.append(endpoint_id)
        if object_type is not None:
            clauses.append("object_type = ?")
            parameters.append(object_type)

        sql = """
            SELECT
                endpoint_id,
                object_type,
                logical_key,
                local_key,
                payload_json,
                references_json,
                applied_at,
                verification_attempted_at,
                verified_at,
                verification_error
            FROM applied_state
        """
        if clauses:
            sql += " WHERE " + " AND ".join(clauses)
        sql += """
            ORDER BY
                endpoint_id,
                object_type,
                logical_key
        """

        rows = self.connection.execute(
            sql,
            tuple(parameters),
        ).fetchall()
        return [
            AppliedState(
                endpoint_id=row["endpoint_id"],
                object=NormalizedObject(
                    object_type=row["object_type"],
                    logical_key=row["logical_key"],
                    local_key=row["local_key"],
                    payload=_deserialize_payload(row["payload_json"]),
                    references=_deserialize_references(
                        row["references_json"]
                    ),
                ),
                applied_at=_parse_timestamp(row["applied_at"]),
                verification_attempted_at=_parse_timestamp(
                    row["verification_attempted_at"]
                ),
                verified_at=_parse_timestamp(row["verified_at"]),
                verification_error=row["verification_error"],
            )
            for row in rows
        ]

    def observation_status(
        self,
        endpoint_id: str,
        object_type: str,
    ) -> ObservationStatus | None:
        row = self.connection.execute(
            """
            SELECT
                endpoint_id,
                object_type,
                last_attempt_at,
                last_success_at,
                last_error
            FROM observation_status
            WHERE endpoint_id = ? AND object_type = ?
            """,
            (endpoint_id, object_type),
        ).fetchone()
        if row is None:
            return None

        return ObservationStatus(
            endpoint_id=row["endpoint_id"],
            object_type=row["object_type"],
            last_attempt_at=_parse_timestamp(row["last_attempt_at"]),
            last_success_at=_parse_timestamp(row["last_success_at"]),
            last_error=row["last_error"],
        )