import threading
import unittest
from unittest.mock import patch

from dashboard import plugin_api as api


class ValidationHarnessTests(unittest.TestCase):
    def test_legacy_multiple_models_map_to_primary_and_validators(self):
        body = api.ChatRequest(models=["primary", "validator-a", "validator-b"])
        with patch.object(api, "_require_installed_model", side_effect=lambda name: name):
            primary, validators, harness = api._harness_models(body)
        self.assertEqual(primary, "primary")
        self.assertEqual(validators, ["validator-a", "validator-b"])
        self.assertTrue(harness)

    def test_harness_requires_two_distinct_validators(self):
        body = api.ChatRequest(primary_model="primary", validator_models=["validator-a"], harness=True)
        with patch.object(api, "_require_installed_model", side_effect=lambda name: name):
            with self.assertRaises(api.HTTPException) as context:
                api._harness_models(body)
        self.assertEqual(context.exception.status_code, 400)

    def test_chat_route_returns_one_compiled_answer(self):
        body = api.ChatRequest(
            primary_model="primary",
            validator_models=["validator-a", "validator-b"],
            harness=True,
            message="Answer this once and validate it.",
        )
        persisted_messages = []

        def fake_stream(payload, request_id, cancel_event=None, parent_id=None):
            prompt = payload["messages"][-1]["content"]
            if payload["model"] == "primary" and "VALIDATION REPORTS:" not in prompt:
                content = "draft"
            elif payload["model"].startswith("validator"):
                content = "No material issues found."
            else:
                content = "single compiled answer"
            return {"message": {"role": "assistant", "content": content}, "done": True}

        def fake_metric(conversation_id, request_id, model, state, status=None, error=""):
            return {"request_id": request_id, "model": model, "status": status}

        with patch.object(api, "_harness_models", return_value=("primary", ["validator-a", "validator-b"], True)), patch.object(
            api, "_require_installed_model", side_effect=lambda name: name
        ), patch.object(api, "_ensure_conversation"), patch.object(api, "_persist_message", side_effect=lambda *args, **kwargs: persisted_messages.append(args)), patch.object(
            api, "_persist_metric", side_effect=fake_metric
        ), patch.object(api, "_stream_chat_request", side_effect=fake_stream), patch.object(
            api, "_runtime_snapshot", return_value={}
        ):
            response = api.chat(body)

        self.assertEqual(response["mode"], "harness")
        self.assertEqual(response["message"]["content"], "single compiled answer")
        self.assertEqual(response["primary_model"], "primary")
        self.assertEqual(response["validator_models"], ["validator-a", "validator-b"])
        assistant_messages = [args for args in persisted_messages if len(args) >= 3 and args[2] == "assistant"]
        self.assertEqual(len(assistant_messages), 1)
        self.assertEqual(assistant_messages[0][3], "single compiled answer")
        self.assertEqual(len(response["validation_reports"]), 2)

    def test_primary_draft_validators_and_primary_compilation_produce_one_answer(self):
        body = api.ChatRequest(
            primary_model="primary",
            validator_models=["validator-a", "validator-b"],
            harness=True,
            message="What is the verified answer?",
        )
        calls = []

        def fake_stream(payload, request_id, cancel_event=None, parent_id=None):
            model = payload["model"]
            prompt = payload["messages"][-1]["content"]
            calls.append((model, prompt, parent_id))
            if model == "primary" and "VALIDATION REPORTS:" not in prompt:
                content = "primary draft"
            elif model.startswith("validator"):
                content = f"{model} found no material issue"
            else:
                content = "one compiled final answer"
            return {"message": {"role": "assistant", "content": content}, "done": True}

        def fake_metric(conversation_id, request_id, model, state, status=None, error=""):
            return {"request_id": request_id, "model": model, "status": status}

        with patch.object(api, "_require_installed_model", side_effect=lambda name: name), patch.object(
            api, "_stream_chat_request", side_effect=fake_stream
        ), patch.object(api, "_persist_metric", side_effect=fake_metric):
            final, metrics, reports = api._run_validation_harness(
                body,
                "root-request",
                "conversation",
                "primary",
                ["validator-a", "validator-b"],
                threading.Event(),
                [],
            )

        self.assertEqual(final, "one compiled final answer")
        self.assertEqual([item["model"] for item in reports], ["validator-a", "validator-b"])
        self.assertEqual(len(metrics), 4)
        self.assertEqual(len(calls), 4)
        self.assertEqual(calls[0][0], "primary")
        self.assertTrue(all(call[2] == "root-request" for call in calls))
        with api._chat_requests_lock:
            child_states = [state for key, state in api._chat_requests.items() if key != "root-request"]
        self.assertGreaterEqual(len(child_states), 4)
        self.assertTrue(all(state.get("parent_id") == "root-request" for state in child_states[-4:]))
        compiler_prompt = calls[-1][1]
        self.assertIn("PRIMARY DRAFT:", compiler_prompt)
        self.assertIn("VALIDATOR 1", compiler_prompt)
        self.assertIn("VALIDATOR 2", compiler_prompt)
        self.assertNotIn("validator-a found no material issue\n\nvalidator-b found no material issue", final)


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