from __future__ import annotations

import json
import os
import sqlite3
import time
from dataclasses import dataclass
from hashlib import sha256
from pathlib import Path
from typing import Any, Iterable


def _ensure_parent(path: Path) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)


def _sha(text: str) -> str:
    return sha256(text.encode("utf-8")).hexdigest()


@dataclass(frozen=True)
class SqliteCacheConfig:
    path: str = "cache/ragnav.sqlite3"


class SqliteKV:
    def __init__(self, cfg: SqliteCacheConfig):
        self.path = Path(cfg.path)
        _ensure_parent(self.path)
        self._conn = sqlite3.connect(str(self.path))
        self._conn.execute("PRAGMA journal_mode=WAL;")
        self._conn.execute(
            """
            CREATE TABLE IF NOT EXISTS kv (
              namespace TEXT NOT NULL,
              key TEXT NOT NULL,
              value_json TEXT NOT NULL,
              created_at REAL NOT NULL,
              PRIMARY KEY (namespace, key)
            );
            """
        )
        self._conn.commit()

    def get(self, namespace: str, key: str) -> Any | None:
        cur = self._conn.execute("SELECT value_json FROM kv WHERE namespace=? AND key=?", (namespace, key))
        row = cur.fetchone()
        if not row:
            return None
        return json.loads(row[0])

    def set(self, namespace: str, key: str, value: Any) -> None:
        self._conn.execute(
            "INSERT OR REPLACE INTO kv(namespace,key,value_json,created_at) VALUES(?,?,?,?)",
            (namespace, key, json.dumps(value), time.time()),
        )
        self._conn.commit()

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


class EmbeddingCache:
    """
    Embedding cache keyed by sha256(model + '\\n' + text).
    """

    def __init__(self, kv: SqliteKV, *, namespace: str = "embeddings"):
        self._kv = kv
        self._ns = namespace

    def key(self, *, model: str | None, text: str) -> str:
        return _sha((model or "") + "\n" + (text or ""))

    def get_many(self, *, model: str | None, texts: list[str]) -> dict[int, list[float]]:
        out: dict[int, list[float]] = {}
        for i, t in enumerate(texts):
            v = self._kv.get(self._ns, self.key(model=model, text=t))
            if isinstance(v, list):
                out[i] = v
        return out

    def set_many(self, *, model: str | None, texts: list[str], embeddings: list[list[float]]) -> None:
        for t, e in zip(texts, embeddings):
            self._kv.set(self._ns, self.key(model=model, text=t), e)


class RetrievalCache:
    """
    Retrieval cache storing a list of block_ids for a query+params key.
    """

    def __init__(self, kv: SqliteKV, *, namespace: str = "retrieval"):
        self._kv = kv
        self._ns = namespace

    def key(self, payload: dict[str, Any]) -> str:
        # stable canonicalization
        raw = json.dumps(payload, sort_keys=True, separators=(",", ":"))
        return _sha(raw)

    def get(self, payload: dict[str, Any]) -> list[str] | None:
        v = self._kv.get(self._ns, self.key(payload))
        if isinstance(v, list) and all(isinstance(x, str) for x in v):
            return v
        return None

    def set(self, payload: dict[str, Any], block_ids: list[str]) -> None:
        self._kv.set(self._ns, self.key(payload), block_ids)

