import json
import sys
import tempfile
import types
import unittest
from pathlib import Path

try:
    import openrouter_wrapper.retry  # noqa: F401
except ModuleNotFoundError:
    package = types.ModuleType("openrouter_wrapper")
    retry = types.ModuleType("openrouter_wrapper.retry")
    retry.openrouter_request = None
    package.retry = retry
    sys.modules["openrouter_wrapper"] = package
    sys.modules["openrouter_wrapper.retry"] = retry

import moralmaps.read_api as read_api


ITEMS = [{"id": "x", "question": "Question?", "options": ["A", "B"], "n": 2}]
IDENTITY_ARGS = {
    "n_samples": 2,
    "temperature": 1.0,
    "max_tokens": 512,
    "concurrency": 1,
    "req_timeout": 90.0,
    "reasoning": None,
    "structured_output": False,
}


class TestDenseAuditReason(unittest.TestCase):
    def test_v1_prompt_and_identity_are_unchanged(self) -> None:
        prompt = read_api._rate_plan(ITEMS, 2)[0]["prompt"]
        self.assertNotIn('"reason"', prompt)
        protocol_id = read_api.rated_protocol_identity("test/model", ITEMS, **IDENTITY_ARGS)
        self.assertEqual(
            protocol_id,
            "b6c5ec06b6ded9b93011deaccb6e7cde5e028d989a448d59f9e3cde9e794b711",
        )

    def test_v2_reason_is_unscored_and_schema_required(self) -> None:
        prompt = read_api._rate_plan(ITEMS, 2, audit_reason=True)[0]["prompt"]
        self.assertLess(prompt.index("Choose every rating first"), prompt.index('"reason" field'))
        schema = read_api._rating_schema(2, audit_reason=True)["json_schema"]["schema"]
        self.assertEqual(schema["required"], ["0", "1", "reason"])
        parsed = read_api._parse_rating_response(
            '{"0": 2, "1": 5, "reason": "This reason has more than eight separate words in total"}',
            2,
            audit_reason=True,
        )
        self.assertEqual(parsed[0], {0: 2.0, 1: 5.0})
        self.assertEqual(parsed[2], "too_long")
        missing = read_api._parse_rating_response('{"0": 2, "1": 5}', 2, audit_reason=True)
        self.assertEqual(missing[0], parsed[0])
        self.assertEqual(missing[2], "missing")

    def test_v2_records_reason_without_changing_ratings(self) -> None:
        async def request(payload):
            self.assertIn("reason", payload["response_format"]["json_schema"]["schema"]["required"])
            return {
                "choices": [{"message": {"content": '{"0": 2, "1": 5, "reason": "Values differ by social context"}'}}],
                "usage": {"cost": 0},
            }

        original_request = read_api.openrouter_request
        read_api.openrouter_request = request
        try:
            with tempfile.TemporaryDirectory() as directory:
                records = Path(directory) / "records.jsonl"
                rows = read_api.read_items_rated(
                    "test/model",
                    ITEMS,
                    n_samples=1,
                    concurrency=1,
                    structured_output=True,
                    records_path=records,
                    eval_version="wvs-score-all-options-v2",
                    audit_reason=True,
                )
                self.assertEqual(rows[0]["p_samples"], [[2 / 7, 5 / 7]])
                self.assertEqual(
                    rows[0]["audit_reasons"],
                    [{"sample": 0, "reason": "Values differ by social context", "status": "valid"}],
                )
                events = [json.loads(line) for line in records.read_text().splitlines()]
                parsed = next(event for event in events if event["event"] == "answer_parsed")
                self.assertEqual(parsed["audit_reason_status"], "valid")
                self.assertEqual(parsed["eval_version"], "wvs-score-all-options-v2")
        finally:
            read_api.openrouter_request = original_request

    def test_v2_requires_new_eval_version(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            with self.assertRaisesRegex(ValueError, "wvs-score-all-options-v2"):
                read_api.read_items_rated(
                    "test/model",
                    ITEMS,
                    n_samples=1,
                    records_path=Path(directory) / "records.jsonl",
                    audit_reason=True,
                )


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