"""Cache frozen FP32 features by signal, preprocessing and encoder identity."""

from pathlib import Path
import hashlib
import json
import sqlite3

import torch

from .provenance import CACHE_CONTRACT, fingerprint
from .atomic import atomic_output


def tensor_digest(values):
    import hashlib
    digest = hashlib.sha256()
    digest.update(values.contiguous().numpy().tobytes())
    return digest.hexdigest()


class FeatureCache:
    def __init__(self, directory):
        self.directory = Path(directory)
        self.directory.mkdir(parents=True, exist_ok=True)

    def get_or_encode(self, model, prepared):
        binding = {"contract": CACHE_CONTRACT, "model": model.identity(),
                   "prepared": prepared.identity}
        path = self.directory / (fingerprint(binding) + ".pt")
        if path.exists():
            payload = torch.load(path, map_location="cpu", weights_only=True)
            if payload.get("binding") != binding:
                raise ValueError("Frozen feature cache contract mismatch")
            values = payload["values"]
            if values.shape != (1, 512):
                raise ValueError("Frozen feature cache shape mismatch")
            if values.dtype != torch.float32:
                raise ValueError("Frozen feature cache must preserve FP32 computation")
            if not torch.isfinite(values).all():
                raise ValueError("Nonfinite cached features")
            if tensor_digest(values) != payload.get("digest"):
                raise ValueError("Cached feature content changed")
            return values.to(model.device)
        encoded = model.encode(prepared)
        values = encoded.cpu().float()
        with atomic_output(path) as output:
            torch.save({"binding": binding, "values": values,
                        "digest": tensor_digest(values)}, output)
        return encoded


class SQLiteFeatureCache:
    """Batch native encodes and keep verified FP32 features in one durable store."""

    def __init__(self, path):
        path = Path(path)
        path.parent.mkdir(parents=True, exist_ok=True)
        self.db = sqlite3.connect(path)
        self.db.executescript('''PRAGMA journal_mode=WAL; PRAGMA synchronous=FULL;
            PRAGMA cache_size=-16384; PRAGMA temp_store=FILE;
            CREATE TABLE IF NOT EXISTS metadata(binding TEXT PRIMARY KEY);
            CREATE TABLE IF NOT EXISTS features(id TEXT PRIMARY KEY, values_blob BLOB, digest TEXT);''')

    @torch.no_grad()
    def get_many_or_encode(self, model, prepared):
        identity = dict(model.identity())
        identity.pop('head', None)
        current = {"contract":CACHE_CONTRACT,"model":identity}
        binding = json.dumps(current,sort_keys=True)
        existing = self.db.execute('SELECT binding FROM metadata').fetchone()
        if existing:
            previous = json.loads(existing[0])
            # Old stores bound head settings even though frozen features do not depend on them.
            if isinstance(previous.get('model'),dict):
                previous['model'].pop('head',None)
            if previous != current:
                raise ValueError("Frozen feature cache contract mismatch")
        if not existing:
            with self.db: self.db.execute('INSERT INTO metadata VALUES (?)',(binding,))
        unique = {value.identity:value for value in prepared}
        found = {}
        keys = list(unique)
        for start in range(0,len(keys),500):
            chunk=keys[start:start+500]
            for key,content,digest in self.db.execute(
                    'SELECT id,values_blob,digest FROM features WHERE id IN ('+','.join('?' for _ in chunk)+')',chunk):
                if len(content)!=512*4 or hashlib.sha256(content).hexdigest()!=digest:
                    raise ValueError("Cached feature content changed")
                values=torch.frombuffer(bytearray(content),dtype=torch.float32).clone()
                if not torch.isfinite(values).all():
                    raise ValueError("Nonfinite cached features")
                found[key]=values
        shapes = {}
        for key,value in unique.items():
            if key not in found:
                shape=tuple(model.preprocessor.format(value).samples.shape)
                shapes.setdefault(shape,[]).append((key,value))
        for entries in shapes.values():
            encoded=model.encode_batch([value for _,value in entries]).cpu().float()
            updates=[]
            for (key,_),values in zip(entries,encoded):
                content=values.contiguous().numpy().tobytes()
                updates.append((key,content,hashlib.sha256(content).hexdigest()))
                found[key]=values
            with self.db:
                self.db.executemany('INSERT INTO features VALUES (?,?,?)',updates)
        return torch.stack([found[value.identity] for value in prepared]).to(model.device)

    def get_or_encode(self, model, prepared):
        return self.get_many_or_encode(model,[prepared])


def open_feature_cache(config):
    if not config.get("feature_cache"):
        return None
    kind=config.get("feature_cache_format","files")
    if kind=="sqlite": return SQLiteFeatureCache(config["feature_cache"])
    if kind=="files": return FeatureCache(config["feature_cache"])
    raise ValueError("feature_cache_format must be files or sqlite")