from __future__ import annotations

import json
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any, Protocol

from .models import BehaviorSkeleton
from .skeletons import generate_skeletons


class JsonModel(Protocol):
    def request_json(
        self,
        *,
        system: str,
        user: str,
        schema_name: str,
        schema: dict[str, Any],
    ) -> dict[str, Any]: ...

    def request_text(self, *, system: str, user: str) -> str: ...


@dataclass(frozen=True)
class EvaluationResult:
    skeleton_id: str
    family: str
    complexity: int
    rendered_problem: str
    answer_options: tuple[str, ...]
    gold_answer: str
    gold_present_probability: float
    solver_output: str
    model_answer: str
    correct: bool


RENDER_SCHEMA = {
    "type": "object",
    "properties": {"problem": {"type": "string"}},
    "required": ["problem"],
    "additionalProperties": False,
}


def render_skeleton(client: JsonModel, skeleton: BehaviorSkeleton) -> str:
    prompt = json.dumps(skeleton.renderer_payload(), ensure_ascii=False, indent=2)
    response = client.request_json(
        system=(
            "You turn behavioral observations from an unknown stochastic system into one self-contained "
            "natural-language inductive reasoning problem. Invent a coherent concrete or fictional domain. "
            "Give each event code and attribute a stable surface name, and preserve every case's attributes, "
            "event order, trial count, observed outcome counts, query, and exact answer labels. The histories "
            "are empirical observations, not explicit system laws: do not state or invent the hidden rules, "
            "trigger probabilities, intermediate flags, exact query probability, or answer. Make clear that "
            "all trials use one unchanged system and that the question asks which final outcome is more likely, "
            "not for an exact probability. Present repeated histories in a compact table rather than verbose "
            "paragraphs. End by requiring exactly one supplied answer label."
        ),
        user=f"Render this intermediate representation:\n{prompt}",
        schema_name="rendered_problem",
        schema=RENDER_SCHEMA,
    )
    problem = response.get("problem")
    if not isinstance(problem, str) or not problem.strip():
        raise RuntimeError("renderer returned an empty problem")
    return problem.strip()


def parse_marked_answer(output: str, answer_options: tuple[str, ...]) -> str:
    lines = [line.strip() for line in output.splitlines() if line.strip()]
    if not lines or not lines[-1].startswith("####"):
        raise RuntimeError("solver output must end with '#### <answer>'")
    answer = lines[-1][4:].strip()
    if answer not in answer_options:
        raise RuntimeError(f"solver returned invalid answer label: {answer!r}")
    return answer


def solve_problem(client: JsonModel, problem: str, answer_options: tuple[str, ...]) -> tuple[str, str]:
    output = client.request_text(
        system=(
            "Infer the stable hidden mechanism from the repeated behavioral observations, then predict which "
            "query outcome is more likely. Event order, attributes, stochastic frequencies, interactions, and "
            "counterfactual controls may matter. Do not demand an exact probability. Give a concise explanation, then put "
            "the final answer on the last non-empty line in the exact form '#### <answer>', where <answer> "
            "is exactly one permitted answer label. Do not use unstated real-world assumptions."
        ),
        user=f"{problem}\n\nPermitted answer labels: {', '.join(answer_options)}",
    )
    return output, parse_marked_answer(output, answer_options)


def _write_jsonl(path: Path, records: list[dict[str, Any]]) -> None:
    with path.open("w", encoding="utf-8") as handle:
        for record in records:
            handle.write(json.dumps(record, ensure_ascii=False) + "\n")


def run_evaluation(
    *,
    renderer: JsonModel,
    solver: JsonModel,
    count: int,
    length: str | int,
    seed: int,
    output_dir: Path,
) -> dict[str, Any]:
    output_dir.mkdir(parents=True, exist_ok=True)
    skeletons = generate_skeletons(count=count, length=length, seed=seed)
    _write_jsonl(output_dir / "skeletons.jsonl", [item.to_dict() for item in skeletons])

    results: list[EvaluationResult] = []
    for skeleton in skeletons:
        problem = render_skeleton(renderer, skeleton)
        solver_output, answer = solve_problem(solver, problem, skeleton.answer_options)
        results.append(
            EvaluationResult(
                skeleton_id=skeleton.id,
                family=skeleton.family,
                complexity=skeleton.complexity,
                rendered_problem=problem,
                answer_options=skeleton.answer_options,
                gold_answer=skeleton.gold_answer,
                gold_present_probability=skeleton.gold_present_probability,
                solver_output=solver_output,
                model_answer=answer,
                correct=answer == skeleton.gold_answer,
            )
        )

    _write_jsonl(output_dir / "results.jsonl", [asdict(item) for item in results])
    correct = sum(item.correct for item in results)
    summary = {
        "count": len(results),
        "correct": correct,
        "accuracy": correct / len(results),
        "length": length,
        "complexity": skeletons[0].complexity,
        "seed": seed,
    }
    (output_dir / "summary.json").write_text(
        json.dumps(summary, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
    )
    return summary