from __future__ import annotations

import unittest
from unittest.mock import patch

from postgrest.exceptions import APIError

from backend.research_worker import ResearchWorker, _is_attempt_budget_exhausted_error


class _FakeResponse:
    def __init__(self, data=None, error: Exception | None = None) -> None:
        self.data = data
        self._error = error

    def execute(self):
        if self._error is not None:
            raise self._error
        return self


class _FakeClient:
    def __init__(self, responses: dict[str, list[_FakeResponse]]) -> None:
        self._responses = {name: list(items) for name, items in responses.items()}
        self.calls: list[tuple[str, dict]] = []

    def rpc(self, name: str, params: dict):
        self.calls.append((name, params))
        queue = self._responses.get(name, [])
        if not queue:
            raise AssertionError(f"Unexpected RPC call: {name}")
        return queue.pop(0)


class ResearchWorkerTests(unittest.TestCase):
    def test_detects_attempt_budget_exhaustion_api_error(self) -> None:
        exc = APIError(
            {
                "message": "Research episode deadbeef exhausted its attempt budget",
                "code": "P0001",
            }
        )
        self.assertTrue(_is_attempt_budget_exhausted_error(exc))

    @patch("backend.research_worker.build_repair_prompt", return_value="repair prompt")
    def test_process_episode_ignores_stale_exhaustion_error(self, _mock_prompt) -> None:
        client = _FakeClient(
            {
                "backend_enqueue_research_attempt": [
                    _FakeResponse(
                        error=APIError(
                            {
                                "message": "Research episode episode-1 exhausted its attempt budget",
                                "code": "P0001",
                            }
                        )
                    )
                ]
            }
        )
        job = {
            "id": "job-1",
            "total_budget": 2000,
            "credits_committed": 1000,
            "budget_per_attempt": 1000,
            "model": "test-model",
        }
        episode = {
            "id": "episode-1",
            "status": "active",
            "attempt_count": 1,
            "max_attempts": 2,
            "conjecture_text": "Prove that True implies True.",
            "formal_statement": "",
            "domain": "logic",
            "strategy_tag": "repair",
            "metadata": {"proof_hint": "Use trivial."},
            "interestingness_score": 1.0,
            "novelty_score": 1.0,
            "attempts": [
                {
                    "id": "attempt-1",
                    "verification_status": "malformed_output",
                    "generated_lean_code": "bad",
                    "lean_stderr": "oops",
                    "lean_stdout": "",
                }
            ],
        }
        worker = ResearchWorker()

        with patch("backend.research_worker.get_supabase_client", return_value=client):
            worker._process_episode(job, episode)

        self.assertEqual(client.calls[0][0], "backend_enqueue_research_attempt")

    def test_maybe_finish_job_completes_when_budget_exhausted_with_verified_episode(self) -> None:
        job_state = {
            "id": "job-1",
            "total_budget": 1000,
            "credits_committed": 1000,
            "budget_per_attempt": 1000,
            "episodes": [
                {"id": "episode-1", "status": "verified", "attempts": []},
                {"id": "episode-2", "status": "active", "attempts": []},
            ],
        }
        client = _FakeClient(
            {
                "backend_get_research_job_state": [_FakeResponse(data=job_state)],
                "backend_update_research_job_status": [_FakeResponse(data=True)],
            }
        )
        worker = ResearchWorker()

        with patch("backend.research_worker.get_supabase_client", return_value=client):
            worker._maybe_finish_job("job-1", enqueued_new_work=False)

        self.assertEqual(
            client.calls,
            [
                ("backend_get_research_job_state", {"p_job_id": "job-1"}),
                (
                    "backend_update_research_job_status",
                    {"p_job_id": "job-1", "p_status": "completed", "p_latest_error": ""},
                ),
            ],
        )


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