"""Surface-renaming generation using the verbatim original prompts."""

from __future__ import annotations

import re
from dataclasses import dataclass
from typing import Literal

from .clients import JsonLLM
from .models import CanonicalItem, ModelCallRecord
from .prompts import surface_system, surface_user
from .store import RunStore


SurfaceFamily = Literal[
    "descriptive_long",
    "descriptive_long_confusing",
    "descriptive_long_misleading",
    "garbled_string",
]

SURFACE_FAMILIES: tuple[SurfaceFamily, ...] = (
    "descriptive_long",
    "descriptive_long_confusing",
    "descriptive_long_misleading",
    "garbled_string",
)

IDENTIFIER_RE = re.compile(r"^[a-z]{8,}$")


@dataclass(frozen=True)
class SurfaceVariant:
    family: SurfaceFamily
    rename_map: dict[str, str]
    question: str
    solution: str

    def as_release_payload(self) -> dict[str, object]:
        return {
            "map": dict(self.rename_map),
            "question": self.question,
            "solution": self.solution,
        }


def validate_surface_variant(
    item: CanonicalItem,
    *,
    rename_map: dict[str, str],
    question: str,
    solution: str,
) -> None:
    expected = set(item.variables + item.parameters)
    if set(rename_map) != expected:
        raise ValueError(
            "rename map must contain every var and param exactly once; "
            f"expected {sorted(expected)}, received {sorted(rename_map)}"
        )
    if len(rename_map.values()) != len(set(rename_map.values())):
        raise ValueError("surface replacements must be one-to-one")
    forbidden = expected | set(item.scientific_constants)
    for replacement in rename_map.values():
        if replacement in forbidden:
            raise ValueError(f"replacement {replacement!r} collides with input symbols")
        if not IDENTIFIER_RE.fullmatch(replacement):
            raise ValueError(
                f"replacement {replacement!r} violates the original prompt contract"
            )
    if not question.strip() or not solution.strip():
        raise ValueError("surface question and solution must be non-empty")


class SurfacePipeline:
    def __init__(self, proposer: JsonLLM, store: RunStore) -> None:
        self.proposer = proposer
        self.store = store

    async def run_family(
        self,
        item: CanonicalItem,
        family: SurfaceFamily,
    ) -> SurfaceVariant:
        request_id = f"{item.item_id}.surface.{family}"
        system_prompt = surface_system(family)
        user_prompt = surface_user(item)
        response = await self.proposer.generate_json(
            system_prompt=system_prompt,
            user_prompt=user_prompt,
            request_id=request_id,
        )
        self.store.write_call(
            ModelCallRecord(
                request_id=request_id,
                model=response.model,
                system_prompt=system_prompt,
                user_prompt=user_prompt,
                response_data=response.data,
                raw_text=response.raw_text,
                provider_response_id=response.response_id,
                usage=response.usage,
            )
        )
        rename_map = {
            str(old): str(new)
            for old, new in dict(response.data["map"]).items()
        }
        question = str(response.data["question"])
        solution = str(response.data["solution"])
        validate_surface_variant(
            item,
            rename_map=rename_map,
            question=question,
            solution=solution,
        )
        variant = SurfaceVariant(
            family=family,
            rename_map=rename_map,
            question=question,
            solution=solution,
        )
        self.store.write_stage(
            f"surface_{family}_variant",
            variant.as_release_payload(),
            request_id=request_id,
        )
        return variant

    async def run_all(
        self,
        item: CanonicalItem,
    ) -> dict[SurfaceFamily, SurfaceVariant]:
        output: dict[SurfaceFamily, SurfaceVariant] = {}
        for family in SURFACE_FAMILIES:
            output[family] = await self.run_family(item, family)
        return output