from __future__ import annotations

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

from .models import WorldSkeleton
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
    solver_output: str
    model_answer: str | None
    format_compliant: bool
    parse_error: str | None
    correct: bool


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


def render_skeleton(client: JsonModel, skeleton: WorldSkeleton) -> str:
    prompt = json.dumps(skeleton.renderer_payload(), ensure_ascii=False, indent=2)
    response = client.request_json(
        system=(
            "Turn a large formal history from one simulated world into a self-contained natural-language "
            "reasoning problem. Invent one coherent concrete or fictional setting and assign stable names to "
            "every symbol. Render every initial observation and every realized event exactly once, preserving "
            "time order, participants, relation direction, numeric values, and state changes. Render the query "
            "at the end. The events are facts that actually occurred, not repeated probability experiments. "
            "Never mention the generator, hidden program, probabilities, proof path, gold answer, or authoring "
            "instructions. Do not directly state the queried proposition. Begin each rendered fact or event line "
            "with its original I#### or V##### audit ID. Use compact chronological log lines rather than "
            "explanatory prose, but do not omit apparent distractors. State the answer labels; the "
            "evaluator separately supplies the solver response protocol."
        ),
        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")
    problem = problem.strip()
    expected_ids = [item["id"] for item in skeleton.initial_state["public_observations"]]
    expected_ids.extend(event.id for event in skeleton.events)
    missing = [identifier for identifier in expected_ids if identifier not in problem]
    if missing:
        preview = ", ".join(missing[:10])
        raise RuntimeError(
            f"renderer omitted {len(missing)} initial observations or events; first missing IDs: {preview}"
        )
    return problem


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 parse_solver_answer(output: str, answer_options: tuple[str, ...]) -> tuple[str, bool]:
    """Prefer the marker protocol, then recover one unambiguous label from the final line."""
    try:
        return parse_marked_answer(output, answer_options), True
    except RuntimeError as strict_error:
        lines = [line.strip() for line in output.splitlines() if line.strip()]
        if not lines:
            raise strict_error
        final_line = lines[-1]
        matches = [
            option
            for option in answer_options
            if re.search(rf"(?<![A-Za-z0-9_]){re.escape(option)}(?![A-Za-z0-9_])", final_line)
        ]
        if len(matches) == 1:
            return matches[0], False
        raise strict_error


def solve_problem(client: JsonModel, problem: str, answer_options: tuple[str, ...]) -> tuple[str, str]:
    output = request_solver_output(client, problem, answer_options)
    answer, _format_compliant = parse_solver_answer(output, answer_options)
    return output, answer


def request_solver_output(client: JsonModel, problem: str, answer_options: tuple[str, ...]) -> str:
    return client.request_text(
        system=(
            "Solve the question from the described world history. Track entity identity, relation direction, "
            "temporal updates, set inheritance, ordering, and relevant compositions while ignoring distractors. "
            "Use only the stated observations and ordinary meanings explicitly established by the story. Keep "
            "the explanation under 200 words, 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)}",
    )


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,
    entity_count: int = 128,
    simulation_steps: int = 1000,
) -> dict[str, Any]:
    output_dir.mkdir(parents=True, exist_ok=True)
    skeletons = generate_skeletons(
        count=count,
        length=length,
        seed=seed,
        entity_count=entity_count,
        simulation_steps=simulation_steps,
    )
    _write_jsonl(output_dir / "skeletons.jsonl", [item.to_dict() for item in skeletons])

    results: list[EvaluationResult] = []
    rendered_records: list[dict[str, str]] = []
    solver_records: list[dict[str, str]] = []
    for skeleton in skeletons:
        problem = render_skeleton(renderer, skeleton)
        rendered_records.append({"skeleton_id": skeleton.id, "rendered_problem": problem})
        _write_jsonl(output_dir / "rendered.jsonl", rendered_records)
        solver_output = request_solver_output(solver, problem, skeleton.answer_options)
        solver_records.append({"skeleton_id": skeleton.id, "solver_output": solver_output})
        _write_jsonl(output_dir / "solver-outputs.jsonl", solver_records)
        try:
            answer, format_compliant = parse_solver_answer(solver_output, skeleton.answer_options)
            parse_error = None
        except RuntimeError as exc:
            answer = None
            format_compliant = False
            parse_error = str(exc)
        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,
                solver_output=solver_output,
                model_answer=answer,
                format_compliant=format_compliant,
                parse_error=parse_error,
                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)
    parsed = sum(item.model_answer is not None for item in results)
    format_compliant = sum(item.format_compliant for item in results)
    summary = {
        "count": len(results),
        "correct": correct,
        "accuracy": correct / len(results),
        "parsed": parsed,
        "parse_rate": parsed / len(results),
        "format_compliant": format_compliant,
        "format_compliance": format_compliant / len(results),
        "length": length,
        "complexity": skeletons[0].complexity,
        "seed": seed,
        "entity_count": entity_count,
        "simulation_steps": simulation_steps,
        "event_counts": [len(item.events) for item in skeletons],
    }
    (output_dir / "summary.json").write_text(
        json.dumps(summary, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
    )
    return summary