import json
import sqlite3
from dataclasses import dataclass
from datetime import UTC, datetime
from pathlib import Path
from typing import cast

from home_agent.contracts.types import (
    EmbeddingRecord,
    Item,
    ItemKind,
    LLMUsage,
    MemoryCandidateStatus,
    MemoryReviewAction,
    MemoryReviewRecord,
    MemoryRule,
    MemoryRuleCandidate,
    RenderedItemText,
    RetrievalDecisionTrace,
    RetrievalStage,
    StoredEmbeddingDocument,
    TodoItem,
    UrgencyDecision,
    UrgencyLabel,
)


@dataclass
class SqliteStore:
    path: Path

    def __post_init__(self) -> None:
        self.path.parent.mkdir(parents=True, exist_ok=True)
        with self._connect() as conn:
            conn.executescript(
                """
                CREATE TABLE IF NOT EXISTS runs (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    started_at TEXT NOT NULL,
                    finished_at TEXT,
                    status TEXT NOT NULL,
                    errors_json TEXT NOT NULL,
                    cost_usd REAL NOT NULL DEFAULT 0.0
                );
                CREATE TABLE IF NOT EXISTS items (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    provider TEXT NOT NULL,
                    external_id TEXT NOT NULL,
                    kind TEXT NOT NULL,
                    received_at TEXT NOT NULL,
                    subject TEXT NOT NULL,
                    sender TEXT NOT NULL,
                    metadata_json TEXT NOT NULL DEFAULT '{}',
                    urgency_score REAL NOT NULL,
                    urgency_label TEXT NOT NULL,
                    reasons TEXT NOT NULL,
                    used_llm INTEGER NOT NULL,
                    UNIQUE(provider, external_id)
                );
                CREATE TABLE IF NOT EXISTS todos (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    title TEXT NOT NULL,
                    priority TEXT NOT NULL,
                    due_at TEXT,
                    source_provider TEXT NOT NULL,
                    source_external_id TEXT NOT NULL,
                    status TEXT NOT NULL DEFAULT 'open',
                    created_at TEXT NOT NULL,
                    updated_at TEXT NOT NULL
                );
                CREATE TABLE IF NOT EXISTS memory_candidates (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    rule_key TEXT NOT NULL,
                    description TEXT NOT NULL,
                    evidence_count INTEGER NOT NULL,
                    confidence REAL NOT NULL,
                    status TEXT NOT NULL DEFAULT 'pending',
                    rejected_reason TEXT,
                    reviewed_at TEXT,
                    created_at TEXT NOT NULL
                );
                CREATE TABLE IF NOT EXISTS memory_rules (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    rule_key TEXT NOT NULL UNIQUE,
                    description TEXT NOT NULL,
                    active INTEGER NOT NULL,
                    created_at TEXT NOT NULL
                );
                CREATE TABLE IF NOT EXISTS memory_reviews (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    rule_key TEXT NOT NULL,
                    action TEXT NOT NULL,
                    reason TEXT,
                    reviewed_at TEXT NOT NULL
                );
                CREATE TABLE IF NOT EXISTS budget (
                    month_key TEXT PRIMARY KEY,
                    spend_usd REAL NOT NULL,
                    hard_cap_usd REAL NOT NULL,
                    blocked INTEGER NOT NULL
                );
                CREATE TABLE IF NOT EXISTS llm_usage (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    run_id INTEGER NOT NULL,
                    provider TEXT NOT NULL,
                    model TEXT,
                    total_cost_usd REAL,
                    input_tokens INTEGER,
                    output_tokens INTEGER,
                    raw_json TEXT NOT NULL,
                    created_at TEXT NOT NULL,
                    FOREIGN KEY(run_id) REFERENCES runs(id)
                );
                CREATE TABLE IF NOT EXISTS item_text_representations (
                    provider TEXT NOT NULL,
                    external_id TEXT NOT NULL,
                    kind TEXT NOT NULL,
                    embedding_text TEXT NOT NULL,
                    summary_text TEXT NOT NULL,
                    source_version TEXT NOT NULL,
                    created_at TEXT NOT NULL,
                    PRIMARY KEY(provider, external_id)
                );
                CREATE TABLE IF NOT EXISTS item_embeddings (
                    provider TEXT NOT NULL,
                    external_id TEXT NOT NULL,
                    embedding_provider TEXT NOT NULL,
                    embedding_model TEXT NOT NULL,
                    vector_json TEXT NOT NULL,
                    dimension INTEGER NOT NULL,
                    created_at TEXT NOT NULL,
                    PRIMARY KEY(provider, external_id, embedding_provider, embedding_model)
                );
                CREATE TABLE IF NOT EXISTS retrieval_events (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    run_id INTEGER NOT NULL,
                    source_provider TEXT NOT NULL,
                    source_external_id TEXT NOT NULL,
                    matched_provider TEXT NOT NULL,
                    matched_external_id TEXT NOT NULL,
                    similarity REAL NOT NULL,
                    stage TEXT NOT NULL,
                    included_in_prompt INTEGER NOT NULL,
                    created_at TEXT NOT NULL
                );
                """
            )
            self._ensure_column(conn, "items", "metadata_json", "TEXT NOT NULL DEFAULT '{}'")
            self._ensure_column(conn, "memory_candidates", "rejected_reason", "TEXT")
            self._ensure_column(conn, "memory_candidates", "reviewed_at", "TEXT")

    def _connect(self) -> sqlite3.Connection:
        conn = sqlite3.connect(self.path)
        conn.row_factory = sqlite3.Row
        return conn

    @staticmethod
    def _ensure_column(conn: sqlite3.Connection, table: str, column: str, column_type: str) -> None:
        rows = conn.execute(f"PRAGMA table_info({table})").fetchall()
        if any(str(row["name"]) == column for row in rows):
            return
        conn.execute(f"ALTER TABLE {table} ADD COLUMN {column} {column_type}")

    def create_run(self, started_at: datetime) -> int:
        with self._connect() as conn:
            cur = conn.execute(
                "INSERT INTO runs(started_at, status, errors_json) VALUES(?, 'running', '[]')",
                (started_at.isoformat(),),
            )
            row_id = cur.lastrowid
            if row_id is None:
                msg = "sqlite did not return run id"
                raise RuntimeError(msg)
            return int(row_id)

    def finish_run(self, run_id: int, status: str, errors_json: str) -> None:
        finished_at = datetime.now(tz=UTC).isoformat()
        with self._connect() as conn:
            conn.execute(
                "UPDATE runs SET finished_at=?, status=?, errors_json=? WHERE id=?",
                (finished_at, status, errors_json, run_id),
            )

    def save_item(self, item: Item, decision: UrgencyDecision) -> None:
        with self._connect() as conn:
            conn.execute(
                """
                INSERT OR REPLACE INTO items(
                    provider, external_id, kind, received_at, subject, sender,
                    metadata_json, urgency_score, urgency_label, reasons, used_llm
                ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                """,
                (
                    item.provider,
                    item.external_id,
                    item.kind.value,
                    item.received_at.isoformat(),
                    item.subject,
                    item.sender,
                    json.dumps(item.metadata),
                    decision.score,
                    decision.label.value,
                    ",".join(decision.reasons),
                    int(decision.used_llm),
                ),
            )

    def save_llm_usage(self, run_id: int, usage: LLMUsage) -> None:
        with self._connect() as conn:
            conn.execute(
                """
                INSERT INTO llm_usage(
                    run_id, provider, model, total_cost_usd,
                    input_tokens, output_tokens, raw_json, created_at
                ) VALUES (?, ?, ?, ?, ?, ?, ?, ?)
                """,
                (
                    run_id,
                    usage.provider,
                    usage.model,
                    usage.total_cost_usd,
                    usage.input_tokens,
                    usage.output_tokens,
                    json.dumps(usage.raw),
                    datetime.now(tz=UTC).isoformat(),
                ),
            )

    def save_item_text(self, rendered: RenderedItemText) -> None:
        now = datetime.now(tz=UTC).isoformat()
        with self._connect() as conn:
            conn.execute(
                """
                INSERT OR REPLACE INTO item_text_representations(
                    provider, external_id, kind, embedding_text, summary_text, source_version, created_at
                ) VALUES (?, ?, ?, ?, ?, ?, ?)
                """,
                (
                    rendered.provider,
                    rendered.external_id,
                    rendered.kind.value,
                    rendered.embedding_text,
                    rendered.summary_text,
                    rendered.source_version,
                    now,
                ),
            )

    def get_item_text(self, provider: str, external_id: str) -> RenderedItemText | None:
        with self._connect() as conn:
            row = conn.execute(
                """
                SELECT provider, external_id, kind, embedding_text, summary_text, source_version
                FROM item_text_representations
                WHERE provider=? AND external_id=?
                """,
                (provider, external_id),
            ).fetchone()
        if row is None:
            return None
        return RenderedItemText(
            provider=str(row["provider"]),
            external_id=str(row["external_id"]),
            kind=ItemKind(str(row["kind"])),
            embedding_text=str(row["embedding_text"]),
            summary_text=str(row["summary_text"]),
            source_version=str(row["source_version"]),
        )

    def save_item_embedding(self, record: EmbeddingRecord) -> None:
        with self._connect() as conn:
            conn.execute(
                """
                INSERT OR REPLACE INTO item_embeddings(
                    provider, external_id, embedding_provider, embedding_model,
                    vector_json, dimension, created_at
                ) VALUES (?, ?, ?, ?, ?, ?, ?)
                """,
                (
                    record.provider,
                    record.external_id,
                    record.embedding_provider,
                    record.embedding_model,
                    json.dumps(record.vector),
                    record.dimension,
                    record.created_at.isoformat(),
                ),
            )

    def get_item_embedding(
        self,
        provider: str,
        external_id: str,
        embedding_provider: str,
        embedding_model: str,
    ) -> EmbeddingRecord | None:
        with self._connect() as conn:
            row = conn.execute(
                """
                SELECT provider, external_id, embedding_provider, embedding_model,
                       vector_json, dimension, created_at
                FROM item_embeddings
                WHERE provider=? AND external_id=? AND embedding_provider=? AND embedding_model=?
                """,
                (provider, external_id, embedding_provider, embedding_model),
            ).fetchone()
        if row is None:
            return None
        raw_vector = json.loads(str(row["vector_json"]))
        vector = [float(value) for value in raw_vector] if isinstance(raw_vector, list) else []
        return EmbeddingRecord(
            provider=str(row["provider"]),
            external_id=str(row["external_id"]),
            embedding_provider=str(row["embedding_provider"]),
            embedding_model=str(row["embedding_model"]),
            vector=vector,
            dimension=int(row["dimension"]),
            created_at=datetime.fromisoformat(str(row["created_at"])),
        )

    def save_retrieval_event(self, trace: RetrievalDecisionTrace) -> None:
        with self._connect() as conn:
            conn.execute(
                """
                INSERT INTO retrieval_events(
                    run_id, source_provider, source_external_id, matched_provider, matched_external_id,
                    similarity, stage, included_in_prompt, created_at
                ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
                """,
                (
                    trace.run_id,
                    trace.source_provider,
                    trace.source_external_id,
                    trace.matched_provider,
                    trace.matched_external_id,
                    trace.similarity,
                    trace.stage,
                    int(trace.included_in_prompt),
                    trace.created_at.isoformat(),
                ),
            )

    def list_retrieval_events(self) -> list[RetrievalDecisionTrace]:
        with self._connect() as conn:
            rows = conn.execute(
                """
                SELECT run_id, source_provider, source_external_id, matched_provider, matched_external_id,
                       similarity, stage, included_in_prompt, created_at
                FROM retrieval_events
                ORDER BY id ASC
                """
            ).fetchall()
        return [
            RetrievalDecisionTrace(
                run_id=int(row["run_id"]),
                source_provider=str(row["source_provider"]),
                source_external_id=str(row["source_external_id"]),
                matched_provider=str(row["matched_provider"]),
                matched_external_id=str(row["matched_external_id"]),
                similarity=float(row["similarity"]),
                stage=cast("RetrievalStage", str(row["stage"])),
                included_in_prompt=bool(int(row["included_in_prompt"])),
                created_at=datetime.fromisoformat(str(row["created_at"])),
            )
            for row in rows
        ]

    def list_embedding_documents(
        self,
        embedding_provider: str,
        embedding_model: str,
        include_item_kinds: set[str] | None = None,
    ) -> list[StoredEmbeddingDocument]:
        query = """
            SELECT e.provider, e.external_id, t.kind, t.summary_text, i.received_at, i.urgency_label,
                   e.embedding_provider, e.embedding_model, e.vector_json
            FROM item_embeddings AS e
            JOIN item_text_representations AS t
              ON t.provider = e.provider AND t.external_id = e.external_id
            JOIN items AS i
              ON i.provider = e.provider AND i.external_id = e.external_id
            WHERE e.embedding_provider=? AND e.embedding_model=?
        """
        params: list[object] = [embedding_provider, embedding_model]
        if include_item_kinds:
            placeholders = ", ".join("?" for _ in include_item_kinds)
            query += f" AND t.kind IN ({placeholders})"
            params.extend(sorted(include_item_kinds))

        with self._connect() as conn:
            rows = conn.execute(query, tuple(params)).fetchall()

        documents: list[StoredEmbeddingDocument] = []
        for row in rows:
            raw_vector = json.loads(str(row["vector_json"]))
            vector = [float(value) for value in raw_vector] if isinstance(raw_vector, list) else []
            label_raw = row["urgency_label"]
            urgency_label = UrgencyLabel(str(label_raw)) if label_raw is not None else None
            documents.append(
                StoredEmbeddingDocument(
                    provider=str(row["provider"]),
                    external_id=str(row["external_id"]),
                    kind=ItemKind(str(row["kind"])),
                    summary_text=str(row["summary_text"]),
                    received_at=datetime.fromisoformat(str(row["received_at"])),
                    urgency_label=urgency_label,
                    embedding_provider=str(row["embedding_provider"]),
                    embedding_model=str(row["embedding_model"]),
                    vector=vector,
                )
            )
        return documents

    def list_items_for_embedding(
        self,
        embedding_provider: str,
        embedding_model: str,
        limit: int | None = None,
        since: datetime | None = None,
        kind: ItemKind | None = None,
        rebuild: bool = False,
    ) -> list[Item]:
        query = """
            SELECT i.provider, i.external_id, i.kind, i.received_at, i.subject, i.sender, i.metadata_json
            FROM items AS i
        """
        clauses: list[str] = []
        params: list[object] = []
        if not rebuild:
            query += """
                LEFT JOIN item_embeddings AS e
                  ON e.provider = i.provider
                 AND e.external_id = i.external_id
                 AND e.embedding_provider = ?
                 AND e.embedding_model = ?
            """
            params.extend([embedding_provider, embedding_model])
            clauses.append("e.provider IS NULL")
        if since is not None:
            clauses.append("i.received_at >= ?")
            params.append(since.isoformat())
        if kind is not None:
            clauses.append("i.kind = ?")
            params.append(kind.value)
        if clauses:
            query += " WHERE " + " AND ".join(clauses)
        query += " ORDER BY i.received_at ASC"
        if limit is not None:
            query += " LIMIT ?"
            params.append(limit)

        with self._connect() as conn:
            rows = conn.execute(query, tuple(params)).fetchall()
        items: list[Item] = []
        for row in rows:
            raw_metadata = json.loads(str(row["metadata_json"]))
            metadata = raw_metadata if isinstance(raw_metadata, dict) else {}
            items.append(
                Item(
                    provider=str(row["provider"]),
                    external_id=str(row["external_id"]),
                    kind=ItemKind(str(row["kind"])),
                    received_at=datetime.fromisoformat(str(row["received_at"])),
                    subject=str(row["subject"]),
                    sender=str(row["sender"]),
                    metadata=metadata,
                )
            )
        return items

    def upsert_todo(self, todo: TodoItem) -> None:
        now = datetime.now(tz=UTC).isoformat()
        with self._connect() as conn:
            conn.execute(
                """
                INSERT INTO todos(
                    title, priority, due_at, source_provider,
                    source_external_id, created_at, updated_at
                ) VALUES (?, ?, ?, ?, ?, ?, ?)
                """,
                (
                    todo.title,
                    todo.priority.value,
                    todo.due_at.isoformat() if todo.due_at else None,
                    todo.source_provider,
                    todo.source_external_id,
                    now,
                    now,
                ),
            )

    def save_memory_candidates(self, candidates: list[MemoryRuleCandidate]) -> None:
        now = datetime.now(tz=UTC).isoformat()
        with self._connect() as conn:
            for candidate in candidates:
                conn.execute(
                    """
                    INSERT INTO memory_candidates(
                        rule_key, description, evidence_count,
                        confidence, status, rejected_reason,
                        reviewed_at, created_at
                    ) VALUES (?, ?, ?, ?, ?, ?, ?, ?)
                    """,
                    (
                        candidate.rule_key,
                        candidate.description,
                        candidate.evidence_count,
                        candidate.confidence,
                        candidate.status,
                        candidate.rejected_reason,
                        None,
                        now,
                    ),
                )

    def list_memory_candidates(
        self,
        status: MemoryCandidateStatus | None = None,
        min_confidence: float | None = None,
    ) -> list[MemoryRuleCandidate]:
        query = (
            "SELECT rule_key, description, evidence_count, confidence, "
            "status, rejected_reason "
            "FROM memory_candidates"
        )
        params: list[object] = []
        clauses: list[str] = []
        if status is not None:
            clauses.append("status=?")
            params.append(status)
        if min_confidence is not None:
            clauses.append("confidence>=?")
            params.append(min_confidence)
        if clauses:
            query += " WHERE " + " AND ".join(clauses)
        query += " ORDER BY id DESC"

        with self._connect() as conn:
            rows = conn.execute(query, tuple(params)).fetchall()
        return [
            MemoryRuleCandidate(
                rule_key=str(row["rule_key"]),
                description=str(row["description"]),
                evidence_count=int(row["evidence_count"]),
                confidence=float(row["confidence"]),
                status=cast("MemoryCandidateStatus", str(row["status"])),
                rejected_reason=row["rejected_reason"],
            )
            for row in rows
        ]

    def approve_memory_candidate(self, rule_key: str, reason: str | None = None) -> None:
        now = datetime.now(tz=UTC).isoformat()
        with self._connect() as conn:
            row = conn.execute(
                (
                    "SELECT rule_key, description, evidence_count, confidence "
                    "FROM memory_candidates "
                    "WHERE rule_key=? "
                    "ORDER BY id DESC LIMIT 1"
                ),
                (rule_key,),
            ).fetchone()
            if row is None:
                msg = f"unknown memory candidate: {rule_key}"
                raise ValueError(msg)
            conn.execute(
                (
                    "INSERT OR REPLACE INTO memory_rules("
                    "rule_key, description, active, created_at"
                    ") VALUES (?, ?, 1, ?)"
                ),
                (str(row["rule_key"]), str(row["description"]), now),
            )
            conn.execute(
                (
                    "UPDATE memory_candidates "
                    "SET status='approved', reviewed_at=?, rejected_reason=NULL "
                    "WHERE rule_key=?"
                ),
                (now, rule_key),
            )
            self._save_memory_review(conn, rule_key, "approved", reason, now)

    def reject_memory_candidate(self, rule_key: str, reason: str | None = None) -> None:
        now = datetime.now(tz=UTC).isoformat()
        with self._connect() as conn:
            exists = conn.execute(
                "SELECT 1 FROM memory_candidates WHERE rule_key=? LIMIT 1",
                (rule_key,),
            ).fetchone()
            if exists is None:
                msg = f"unknown memory candidate: {rule_key}"
                raise ValueError(msg)
            conn.execute(
                (
                    "UPDATE memory_candidates "
                    "SET status='rejected', rejected_reason=?, reviewed_at=? "
                    "WHERE rule_key=?"
                ),
                (reason, now, rule_key),
            )
            self._save_memory_review(conn, rule_key, "rejected", reason, now)

    @staticmethod
    def _save_memory_review(
        conn: sqlite3.Connection,
        rule_key: str,
        action: MemoryReviewAction,
        reason: str | None,
        reviewed_at: str,
    ) -> None:
        conn.execute(
            "INSERT INTO memory_reviews(rule_key, action, reason, reviewed_at) VALUES (?, ?, ?, ?)",
            (rule_key, action, reason, reviewed_at),
        )

    def list_memory_reviews(self) -> list[MemoryReviewRecord]:
        with self._connect() as conn:
            rows = conn.execute(
                "SELECT rule_key, action, reason, reviewed_at FROM memory_reviews ORDER BY id DESC"
            ).fetchall()
        return [
            MemoryReviewRecord(
                rule_key=str(row["rule_key"]),
                action=cast("MemoryReviewAction", str(row["action"])),
                reason=row["reason"],
                reviewed_at=datetime.fromisoformat(str(row["reviewed_at"])),
            )
            for row in rows
        ]

    def list_active_memory_rules(self) -> list[MemoryRule]:
        with self._connect() as conn:
            rows = conn.execute(
                "SELECT id, rule_key, description, active, created_at FROM memory_rules WHERE active=1"
            ).fetchall()
        return [
            MemoryRule(
                id=int(row["id"]),
                rule_key=str(row["rule_key"]),
                description=str(row["description"]),
                active=bool(int(row["active"])),
                created_at=datetime.fromisoformat(str(row["created_at"])),
            )
            for row in rows
        ]

    def get_month_spend(self, month_key: str, hard_cap_usd: float) -> tuple[float, bool]:
        with self._connect() as conn:
            row = conn.execute(
                "SELECT spend_usd, blocked FROM budget WHERE month_key=?",
                (month_key,),
            ).fetchone()
            if row is None:
                conn.execute(
                    (
                        "INSERT INTO budget("
                        "month_key, spend_usd, hard_cap_usd, blocked"
                        ") VALUES (?, 0.0, ?, 0)"
                    ),
                    (month_key, hard_cap_usd),
                )
                return 0.0, False
            return float(row["spend_usd"]), bool(int(row["blocked"]))

    def add_month_spend(self, month_key: str, delta_usd: float, hard_cap_usd: float) -> None:
        spend, _ = self.get_month_spend(month_key, hard_cap_usd)
        new_spend = spend + delta_usd
        blocked = 1 if new_spend >= hard_cap_usd else 0
        with self._connect() as conn:
            conn.execute(
                "UPDATE budget SET spend_usd=?, hard_cap_usd=?, blocked=? WHERE month_key=?",
                (new_spend, hard_cap_usd, blocked, month_key),
            )
