import json
import tempfile
import unittest
from pathlib import Path

from logicbench.pipeline import parse_marked_answer, parse_solver_answer, run_evaluation


class QueueClient:
    def __init__(self, responses):
        self.responses = iter(responses)

    def request_json(self, **_kwargs):
        return next(self.responses)

    def request_text(self, **_kwargs):
        return next(self.responses)


class FailingSolver:
    def request_text(self, **_kwargs):
        raise RuntimeError("solver failed")


class AuditRenderer:
    def request_json(self, **kwargs):
        payload = json.loads(kwargs["user"][kwargs["user"].find("{") :])
        identifiers = [item["id"] for item in payload["initial_observations"]]
        identifiers.extend(item["id"] for item in payload["events"])
        return {"problem": "A compact behavioral problem.\n" + "\n".join(identifiers)}


class PipelineTests(unittest.TestCase):
    def test_end_to_end_writes_skeletons_before_scored_results(self):
        renderer = AuditRenderer()
        solver = QueueClient([
            "The relation follows through the event chain.\n#### YES",
            "The target set is disjoint from the inherited membership.\n#### NO",
            "The reverse ordering is false.\n#### NO",
            "Membership propagates through the nested sets.\n#### YES",
            "The relation follows through the event chain.\n#### YES",
        ])
        with tempfile.TemporaryDirectory() as directory:
            output = Path(directory)
            summary = run_evaluation(
                renderer=renderer,
                solver=solver,
                count=5,
                length="short",
                seed=20260922,
                output_dir=output,
            )
            self.assertTrue((output / "skeletons.jsonl").is_file())
            self.assertTrue((output / "rendered.jsonl").is_file())
            self.assertTrue((output / "solver-outputs.jsonl").is_file())
            self.assertTrue((output / "results.jsonl").is_file())
            self.assertEqual(summary["count"], 5)
            self.assertEqual(summary["correct"], 5)
            self.assertEqual(summary["format_compliant"], 5)

    def test_answer_parser_requires_final_marker_and_permitted_label(self):
        self.assertEqual(parse_marked_answer("Reason.\n#### YES", ("YES", "NO")), "YES")
        with self.assertRaisesRegex(RuntimeError, "must end"):
            parse_marked_answer("#### YES\nextra", ("YES", "NO"))
        with self.assertRaisesRegex(RuntimeError, "invalid answer"):
            parse_marked_answer("#### MAYBE", ("YES", "NO"))

    def test_answer_parser_recovers_unique_label_on_final_line(self):
        self.assertEqual(
            parse_solver_answer("Reasoning.\nFinal answer: YES", ("YES", "NO")),
            ("YES", False),
        )
        self.assertEqual(parse_solver_answer("Reasoning.\n#### YES", ("YES", "NO")), ("YES", True))

    def test_rendered_problem_is_checkpointed_before_solver_failure(self):
        renderer = AuditRenderer()
        with tempfile.TemporaryDirectory() as directory:
            output = Path(directory)
            with self.assertRaisesRegex(RuntimeError, "solver failed"):
                run_evaluation(
                    renderer=renderer,
                    solver=FailingSolver(),
                    count=1,
                    length="short",
                    seed=1,
                    output_dir=output,
                )
            self.assertIn("A compact behavioral problem.", (output / "rendered.jsonl").read_text())

    def test_unparseable_solver_output_is_scored_without_stopping_batch(self):
        renderer = AuditRenderer()
        solver = QueueClient(["The response ended before its final label."])
        with tempfile.TemporaryDirectory() as directory:
            output = Path(directory)
            summary = run_evaluation(
                renderer=renderer,
                solver=solver,
                count=1,
                length="short",
                seed=1,
                output_dir=output,
            )
            self.assertEqual(summary["parsed"], 0)
            self.assertEqual(summary["correct"], 0)
            self.assertIn('"parse_error":', (output / "results.jsonl").read_text())

    def test_renderer_omission_is_rejected_before_solving(self):
        renderer = QueueClient([{"problem": "This rendering omitted its audit IDs."}])
        with tempfile.TemporaryDirectory() as directory:
            with self.assertRaisesRegex(RuntimeError, "renderer omitted"):
                run_evaluation(
                    renderer=renderer,
                    solver=FailingSolver(),
                    count=1,
                    length="short",
                    seed=1,
                    output_dir=Path(directory),
                )


if __name__ == "__main__":
    unittest.main()