"""Production runtime wrapper with graph/model/version-aware caching."""

from __future__ import annotations

import hashlib
from collections.abc import Mapping
from dataclasses import dataclass, replace

import torch

from graph.graph import GraphStore
from graph.model import GraphModel
from graph.training.calibration import TemperatureCalibrator
from graph.types import GraphQuery, ReasoningResult


@dataclass(frozen=True)
class CacheKey:
    tenant_id: str
    graph_version: str
    model_version: str
    query_hash: str


class ReasoningRuntime:
    def __init__(
        self,
        model: GraphModel,
        graphs: Mapping[str, GraphStore] | None = None,
        calibrator: TemperatureCalibrator | None = None,
    ) -> None:
        self.model = model
        self.graphs = dict(graphs or {})
        self.calibrator = calibrator
        self.cache: dict[CacheKey, ReasoningResult] = {}

    def register_graph(self, tenant_id: str, graph: GraphStore) -> None:
        self.graphs[tenant_id] = graph
        for key in tuple(self.cache):
            if key.tenant_id == tenant_id:
                del self.cache[key]

    def reason(
        self, tenant_id: str, query: GraphQuery, *, use_cache: bool = True
    ) -> ReasoningResult:
        graph = self.graphs.get(tenant_id)
        if graph is None:
            raise KeyError(f"unknown tenant graph {tenant_id!r}")
        key = CacheKey(
            tenant_id,
            graph.graph_version,
            self.model.config.model_version,
            hashlib.sha256(query.canonical_text().encode()).hexdigest(),
        )
        if use_cache and key in self.cache:
            return self.cache[key]
        result = self.model.reason(graph, query)
        if self.calibrator is not None:
            probability = torch.tensor(result.confidence).clamp(1e-6, 1 - 1e-6)
            logit = torch.log(probability / (1 - probability))
            calibrated = float(self.calibrator.transform(logit).item())
            result = replace(result, confidence=calibrated, calibrated=True)
        if use_cache:
            self.cache[key] = result
        return result