from __future__ import annotations
import duckdb
import hashlib
import contextlib
import threading
import time
from dataclasses import dataclass
from typing import Optional, Iterable, Tuple, List, Dict
import pandas as pd
import logging
import os

log = logging.getLogger(__name__)


@dataclass
class ScoreCacheConfig:
    db_path: str
    version: str = "v2"  # Bumped due to Gmail dot normalization changes
    threads: int = min(os.cpu_count() or 4, 8)
    ttl_days: int = 30
    batch_size: int = 10_000  # tune for your workload


def _normalize(s: str) -> str:
    return (s or "").strip().lower()


def _hash16(s: str, namespace: str = "") -> bytes:
    """Hash with optional namespace to prevent collisions."""
    combined = f"{namespace}:{_normalize(s)}" if namespace else _normalize(s)
    return hashlib.blake2b(combined.encode("utf-8"), digest_size=16).digest()


def _ordered_hashes(a: str, b: str, namespace: str = "") -> Tuple[bytes, bytes]:
    """Create ordered hashes with namespace."""
    h1, h2 = _hash16(a, namespace), _hash16(b, namespace)
    return (h1, h2) if h1 <= h2 else (h2, h1)


class ScoreCache:
    _writer_lock = threading.Lock()
    _writer_initialized = False

    """
    Per-process cache. Safe to create in each worker's init.
    Uses batched upserts inside transactions; DuckDB handles read concurrency.
    """

    def __init__(
        self, cfg: ScoreCacheConfig, is_worker: bool = False, read_only: bool = False
    ):
        self.cfg = cfg
        self._is_worker = is_worker

        db_target = ":memory:" if is_worker else cfg.db_path

        effective_read_only = read_only and db_target != ":memory:"
        if not is_worker and not effective_read_only:
            with ScoreCache._writer_lock:
                if ScoreCache._writer_initialized:
                    log.warning(
                        "Additional ScoreCache writer requested; switching to read-only mode"
                    )
                    effective_read_only = True
                else:
                    ScoreCache._writer_initialized = True
        self._is_writer = (not is_worker) and (not effective_read_only)

        self._con = duckdb.connect(db_target, read_only=effective_read_only)
        threads = 1 if is_worker else cfg.threads
        self._con.execute(f"PRAGMA threads={threads}")
        self._init_schema()
        self._lock = threading.Lock()  # Serialize only our own batched writes
        self._pending: List[Tuple[str, bytes, bytes, float]] = []

    @classmethod
    def reset_writer_flag(cls):
        with cls._writer_lock:
            cls._writer_initialized = False

    def _init_schema(self):
        self._con.execute("""
            CREATE TABLE IF NOT EXISTS pair_scores (
                algo TEXT NOT NULL,
                key1 BLOB NOT NULL,
                key2 BLOB NOT NULL,
                version TEXT NOT NULL,
                score DOUBLE NOT NULL,
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                PRIMARY KEY (algo, key1, key2, version)
            );
        """)
        self._con.execute(
            "CREATE INDEX IF NOT EXISTS idx_pair_scores_lookup "
            "ON pair_scores(algo, key1, key2, version);"
        )
        self._con.execute(
            "CREATE INDEX IF NOT EXISTS idx_pair_scores_updated "
            "ON pair_scores(updated_at);"
        )

    def get(self, algo: str, s1: str, s2: str) -> Optional[float]:
        namespace = f"{algo}:{self.cfg.version}"
        k1, k2 = _ordered_hashes(s1, s2, namespace)
        row = self._con.execute(
            "SELECT score FROM pair_scores WHERE algo=? AND key1=? AND key2=? AND version=?",
            [algo, k1, k2, self.cfg.version],
        ).fetchone()
        return float(row[0]) if row else None

    def get_many(
        self, queries: Iterable[Tuple[str, str, str]]
    ) -> Dict[Tuple[str, bytes, bytes], float]:
        # Efficient batch lookup via temp table + join
        data = []
        for algo, s1, s2 in queries:
            namespace = f"{algo}:{self.cfg.version}"
            k1, k2 = _ordered_hashes(s1, s2, namespace)
            data.append((algo, k1, k2, self.cfg.version))
        if not data:
            return {}

        df = pd.DataFrame(data, columns=["algo", "key1", "key2", "version"])
        self._con.register("q", df)
        rows = self._con.execute("""
            SELECT q.algo, q.key1, q.key2, p.score
            FROM q LEFT JOIN pair_scores p
            USING(algo, key1, key2, version)
            WHERE p.score IS NOT NULL
        """).fetchall()
        self._con.unregister("q")
        return {(algo, k1, k2): float(score) for (algo, k1, k2, score) in rows}

    def set(self, algo: str, s1: str, s2: str, score: float):
        namespace = f"{algo}:{self.cfg.version}"
        k1, k2 = _ordered_hashes(s1, s2, namespace)
        self._pending.append((algo, k1, k2, float(score)))
        if len(self._pending) >= self.cfg.batch_size:
            self.flush()

    def flush(self):
        """Write pending changes to disk."""
        if self._is_worker:
            # In-memory DB for workers, no need to flush
            self._pending.clear()  # Just clear the pending list
            return

        if not self._pending:
            return
        with self._lock:
            df = pd.DataFrame(self._pending, columns=["algo", "key1", "key2", "score"])

            # Retry logic for write contention
            for attempt in range(3):
                try:
                    self._con.register("ins", df)
                    self._con.execute(
                        """
                        BEGIN TRANSACTION;
                        INSERT OR REPLACE INTO pair_scores (algo, key1, key2, version, score, updated_at)
                        SELECT algo, key1, key2, ?, score, CURRENT_TIMESTAMP FROM ins;
                        COMMIT;
                    """,
                        [self.cfg.version],
                    )
                    self._con.unregister("ins")
                    self._pending.clear()
                    return
                except (
                    duckdb.IOException,
                    duckdb.CatalogException,
                    duckdb.TransactionException,
                ) as e:
                    if attempt == 2:
                        log.error(f"Failed to flush cache after 3 attempts: {e}")
                        raise
                    time.sleep(0.1 * (2**attempt))  # Exponential backoff

    def cleanup_old(self):
        self._con.execute(
            f"DELETE FROM pair_scores WHERE updated_at < CURRENT_TIMESTAMP - INTERVAL '{self.cfg.ttl_days} days';"
        )
        self._con.execute("CHECKPOINT;")

    def close(self):
        try:
            self.flush()
        finally:
            with contextlib.suppress(Exception):
                self._con.close()
            if self._is_writer:
                self.reset_writer_flag()
