"""
LLM response cache backed by SQLite.

Avoids redundant API calls by caching responses keyed on
SHA-256(model_name || system_prompt || user_prompt).
"""

from __future__ import annotations

import sqlite3
import time
from typing import Optional

from tracescope.utils.hashing import sha256_key


class LLMResponseCache:
    """SQLite-backed LLM response cache.

    Args:
        db_path: Path to the SQLite database file.
        enabled: If False, all lookups miss and stores are skipped.
    """

    def __init__(self, db_path: str, enabled: bool = True):
        self.db_path = db_path
        self.enabled = enabled
        if enabled:
            self._conn = sqlite3.connect(db_path, check_same_thread=False)
            self._init_table()
        else:
            self._conn = None

    def _init_table(self):
        self._conn.execute(
            """
            CREATE TABLE IF NOT EXISTS llm_cache (
                hash        TEXT PRIMARY KEY,
                model_name  TEXT NOT NULL,
                response    TEXT NOT NULL,
                created_at  REAL NOT NULL
            )
            """
        )
        self._conn.commit()

    def get(self, model_name: str, system_prompt: str, user_prompt: str) -> Optional[str]:
        """Look up a cached response. Returns None on miss."""
        if not self.enabled:
            return None
        key = sha256_key(model_name, system_prompt, user_prompt)
        row = self._conn.execute(
            "SELECT response FROM llm_cache WHERE hash = ?", (key,)
        ).fetchone()
        return row[0] if row else None

    def put(self, model_name: str, system_prompt: str, user_prompt: str, response: str):
        """Store a response in the cache."""
        if not self.enabled:
            return
        key = sha256_key(model_name, system_prompt, user_prompt)
        self._conn.execute(
            """
            INSERT OR REPLACE INTO llm_cache (hash, model_name, response, created_at)
            VALUES (?, ?, ?, ?)
            """,
            (key, model_name, response, time.time()),
        )
        self._conn.commit()

    def clear(self):
        """Clear all cached responses."""
        if self._conn:
            self._conn.execute("DELETE FROM llm_cache")
            self._conn.commit()

    def close(self):
        if self._conn:
            self._conn.close()


class ResultCache:
    """SQLite-backed cache for expensive ML computation results.

    Caches clustering, dimension reduction, and velocity grid results
    keyed on embeddings fingerprint (matching Android's DimensionReducerAdapter).

    Uses the same SQLite database as LLMResponseCache but a separate table.

    Args:
        db_path: Path to the SQLite database file.
        enabled: If False, all lookups miss and stores are skipped.
    """

    def __init__(self, db_path: str, enabled: bool = True):
        self.db_path = db_path
        self.enabled = enabled
        if enabled:
            self._conn = sqlite3.connect(db_path, check_same_thread=False)
            self._init_table()
        else:
            self._conn = None

    def _init_table(self):
        self._conn.execute(
            """
            CREATE TABLE IF NOT EXISTS result_cache (
                key         TEXT PRIMARY KEY,
                step        TEXT NOT NULL,
                data        BLOB NOT NULL,
                created_at  REAL NOT NULL
            )
            """
        )
        self._conn.commit()

    def get_result(self, step: str, fingerprint: str) -> Optional[bytes]:
        """Look up a cached computation result.

        Args:
            step: Pipeline step name (e.g. "clustering", "dim_reduction", "velocity_grid").
            fingerprint: Embeddings fingerprint (SHA-256 hex digest).

        Returns:
            Cached data as bytes, or None on miss.
        """
        if not self.enabled:
            return None
        key = sha256_key(step, fingerprint)
        row = self._conn.execute(
            "SELECT data FROM result_cache WHERE key = ?", (key,)
        ).fetchone()
        return row[0] if row else None

    def put_result(self, step: str, fingerprint: str, data: bytes):
        """Store a computation result in the cache.

        Args:
            step: Pipeline step name.
            fingerprint: Embeddings fingerprint.
            data: Result data as bytes (JSON-encoded or raw numpy).
        """
        if not self.enabled:
            return
        key = sha256_key(step, fingerprint)
        self._conn.execute(
            """
            INSERT OR REPLACE INTO result_cache (key, step, data, created_at)
            VALUES (?, ?, ?, ?)
            """,
            (key, step, data, time.time()),
        )
        self._conn.commit()

    def clear(self):
        """Clear all cached results."""
        if self._conn:
            self._conn.execute("DELETE FROM result_cache")
            self._conn.commit()

    def close(self):
        if self._conn:
            self._conn.close()