"""Spike D — VerdictStore + ExpeL retrieval (B2).

De-risks the assumption that retrieval surfaces a relevant *prior* verdict for a similar
new proposal — the substrate for the framework's learning loop.

**Similarity is structural, not textual (reviewer refinement #2):** a weighted score over
*structured* fields — Jaccard on the affected cost-code set + a `measure_type` match + a
magnitude-bucket match on the claimed saving — **never** raw description text. This makes
the retrieval non-tautological: a true match with different wording beats decoys that
merely share surface text.

`retrieve()` is the guaranteed SC-D deliverable (always tested). A thin
`ExpeLContextProvider` wraps it for few-shot injection via the real `ContextProvider`
hook. The embedding-based similarity path is intentionally **out of scope** for a
throwaway spike (a Fase 2 option — see findings-d).
"""

from __future__ import annotations

from dataclasses import dataclass

from agent_framework import ContextProvider

# Weights: the affected cost-code overlap dominates, then measure type, then magnitude.
_W_CODES, _W_MEASURE, _W_MAGNITUDE = 0.60, 0.25, 0.15
_MAGNITUDE_BUCKETS = [(0.0, 1e5), (1e5, 5e5), (5e5, 1e6), (1e6, float("inf"))]


@dataclass(frozen=True)
class ProposalFeatures:
    """The *structured* features retrieval ranks on. ``description`` is surface text and
    is deliberately NOT part of the similarity score."""

    affected_codes: frozenset[str]
    measure_type: str
    claimed_saving_nok: float
    description: str = ""


@dataclass(frozen=True)
class Verdict:
    """One historical expert verdict in the store."""

    id: str
    proposal_features: ProposalFeatures
    decision: str  # "approved" | "rejected"
    rationale: str


def _magnitude_bucket(value: float) -> int:
    for i, (low, high) in enumerate(_MAGNITUDE_BUCKETS):
        if low <= value < high:
            return i
    return len(_MAGNITUDE_BUCKETS) - 1


def _jaccard(a: frozenset[str], b: frozenset[str]) -> float:
    union = a | b
    return len(a & b) / len(union) if union else 1.0


def similarity(query: ProposalFeatures, candidate: ProposalFeatures) -> float:
    """Weighted structural similarity in [0, 1] — text is ignored by design."""
    codes = _jaccard(query.affected_codes, candidate.affected_codes)
    measure = 1.0 if query.measure_type == candidate.measure_type else 0.0
    magnitude = (
        1.0
        if _magnitude_bucket(query.claimed_saving_nok)
        == _magnitude_bucket(candidate.claimed_saving_nok)
        else 0.0
    )
    return _W_CODES * codes + _W_MEASURE * measure + _W_MAGNITUDE * magnitude


@dataclass
class VerdictStore:
    """Minimal in-memory store of historical verdicts."""

    verdicts: list[Verdict]

    def retrieve(self, query: ProposalFeatures, k: int) -> list[Verdict]:
        """Return the top-``k`` verdicts by structural similarity. Deterministic: ties
        break by verdict id, so ordering is stable across runs."""
        if k <= 0:
            raise ValueError(f"k must be positive, got {k}")
        ranked = sorted(
            self.verdicts,
            key=lambda v: (-similarity(query, v.proposal_features), v.id),
        )
        return ranked[:k]


class ExpeLContextProvider(ContextProvider):
    """Wraps `VerdictStore.retrieve` for ExpeL few-shot injection.

    Asserted only against the introspected `ContextProvider` interface
    (`before_run` + `SessionContext.extend_instructions`); the `retrieve` ranking remains
    the guaranteed SC-D deliverable regardless of the MAF session surface.
    """

    def __init__(self, store: VerdictStore, query: ProposalFeatures, *, k: int = 3) -> None:
        super().__init__(source_id="expel-verdictstore")
        self._store = store
        self._query = query
        self._k = k

    def format_fewshot(self) -> str:
        hits = self._store.retrieve(self._query, self._k)
        body = "\n".join(f"- [{v.id}] {v.decision}: {v.rationale}" for v in hits)
        return f"Relevant prior verdicts (ExpeL few-shot):\n{body}"

    async def before_run(
        self, *, agent: object, session: object, context: object, state: dict
    ) -> None:
        context.extend_instructions([self.format_fewshot()])  # type: ignore[attr-defined]


def seed_store() -> VerdictStore:
    """Seed 12 synthetic verdicts spanning the reference domain's cost codes, measure
    types, magnitudes, and decisions (B2 store of 10–20)."""
    rows = [
        (
            "V01",
            {"05.2", "03.1"},
            "scope_reduction",
            180_000,
            "approved",
            "asphalt + base course trimmed within feasible range",
        ),
        (
            "V02",
            {"05.2"},
            "rate_renegotiation",
            60_000,
            "approved",
            "renegotiated asphalt unit rate",
        ),
        (
            "V03",
            {"07.4"},
            "material_substitution",
            120_000,
            "rejected",
            "granite kerb substitution unsafe",
        ),
        (
            "V04",
            {"09.1"},
            "scope_reduction",
            240_000,
            "approved",
            "fewer LED masts on low-traffic stretch",
        ),
        (
            "V05",
            {"02.3", "03.1"},
            "scope_reduction",
            350_000,
            "rejected",
            "soil replacement is load-bearing, cannot cut",
        ),
        ("V06", {"21.2"}, "rate_renegotiation", 90_000, "approved", "blasting rate renegotiated"),
        (
            "V07",
            {"22.4"},
            "material_substitution",
            700_000,
            "rejected",
            "fiber shotcrete spec is mandated",
        ),
        (
            "V08",
            {"88.2"},
            "scope_reduction",
            150_000,
            "approved",
            "concrete repair area re-measured smaller",
        ),
        (
            "V09",
            {"87.3"},
            "material_substitution",
            130_000,
            "approved",
            "alternative membrane qualified",
        ),
        (
            "V10",
            {"05.2", "03.1"},
            "rate_renegotiation",
            200_000,
            "approved",
            "combined paving rate discount",
        ),
        (
            "V11",
            {"01.1"},
            "scope_reduction",
            95_000,
            "rejected",
            "rigging is fixed cost, no scope to cut",
        ),
        (
            "V12",
            {"31.3"},
            "scope_reduction",
            110_000,
            "approved",
            "drainage length reduced after survey",
        ),
    ]
    return VerdictStore(
        verdicts=[
            Verdict(
                id=vid,
                proposal_features=ProposalFeatures(
                    affected_codes=frozenset(codes),
                    measure_type=mtype,
                    claimed_saving_nok=saving,
                    description=desc,
                ),
                decision=decision,
                rationale=desc,
            )
            for vid, codes, mtype, saving, decision, desc in rows
        ]
    )