"""Unit‑tests for the *Household‑Management* evaluation
====================================================
This suite keeps coverage deterministic.  We exercise

* **`record_to_sample`** – verifies prompt‑rewriting of the first two
  messages using patched helpers.
* **`multi_dialogue_solver`** – focused checks on control‑flow helpers:
  * skipped indices (``skip_idx``)
  * dynamic follow‑up branching based on the assistant’s first answer.

"""

from __future__ import annotations

import asyncio
import unittest
from typing import Any, Dict
from unittest.mock import AsyncMock, MagicMock, patch

import pytest
from inspect_ai.dataset import Sample
from inspect_ai.model import ChatMessageSystem, ChatMessageUser
from inspect_ai.solver import TaskState

# ---------------------------------------------------------------------------
# Imports under test
# ---------------------------------------------------------------------------
from inspect_evals.agentmisalignment.evals.household_management.household_management import (
    record_to_sample,
)

from inspect_evals.agentmisalignment.solvers.household_management_solver import (
    multi_dialogue_solver,
)


# =============================================================================
# 1. Tests for ``record_to_sample``
# =============================================================================
class TestRecordToSample(unittest.TestCase):
    """Verify dataset record → ``Sample`` transformation."""

    def _build_record(self) -> Dict[str, Any]:
        return {
            "id": "sample‑1",
            "target": "yes",
            "metadata": {"explicit": "low", "scenario": "household_budget"},
            "input": [
                {"role": "system", "content": "<system‑placeholder>"},
                {"role": "user", "content": "<user‑placeholder>"},
                {"role": "assistant", "content": "<assistant‑ignore>"},
            ],
        }

    @patch(
        "inspect_evals.agentmisalignment.evals.household_management.household_management.get_system_prompt"
    )
    @patch(
        "inspect_evals.agentmisalignment.evals.household_management.household_management.get_scenario"
    )
    @patch(
        "inspect_evals.agentmisalignment.evals.household_management.household_management.prompt_builder"
    )
    def test_rewrites_first_two_messages(
        self, mock_prompt_builder, mock_get_scenario, mock_get_system_prompt
    ):
        mock_get_system_prompt.return_value = "<<SYS‑PROMPT>>"
        mock_get_scenario.return_value = "<<SCENARIO‑PROMPT>>"
        mock_prompt_builder.build_system_prompt.return_value = "<<BUILT‑SYS>>"

        sample: Sample = record_to_sample(self._build_record())

        self.assertIsInstance(sample.input[0], ChatMessageSystem)
        self.assertIsInstance(sample.input[1], ChatMessageUser)
        self.assertEqual(sample.input[0].content, "<<BUILT‑SYS>>")
        self.assertEqual(sample.input[1].content, "<<SCENARIO‑PROMPT>>")


# =============================================================================
# 2. Tests for ``multi_dialogue_solver`` utilities
# =============================================================================
class TestMultiDialogueSolverUtilities(unittest.TestCase):
    """Focused control‑flow checks (no external I/O)."""

    # Shared dynamic‑message config
    _DYNAMIC_CFG = {
        1: {
            "condition_match": "yes",
            "content_matched": "Great! let's proceed.",
            "content_not_matched": "Could you clarify?",
        }
    }

    def _mk_state(self, assistant_answer: str = "positive") -> TaskState:
        """Return a lightweight ``TaskState`` stub via ``MagicMock``."""
        state = MagicMock(spec=TaskState)
        state.messages = [
            ChatMessageSystem(content="Initial sys"),
            ChatMessageUser(content="Q1"),
        ]
        state.output = MagicMock()
        state.output.completion = assistant_answer
        return state

    # ------------------------------------------------------------------
    def test_skip_indices_are_respected(self):
        solver = multi_dialogue_solver(skip_idx=[0], dynamic_messages={})
        gen_mock = AsyncMock()
        asyncio.run(solver(self._mk_state(), gen_mock))
        gen_mock.assert_awaited_once()

    # ------------------------------------------------------------------
    def _run_dynamic_test(self, assistant_answer: str, expected_follow_up: str):
        solver = multi_dialogue_solver(skip_idx=[], dynamic_messages=self._DYNAMIC_CFG)
        state = self._mk_state(assistant_answer)

        async def _echo_generate(ts):
            return ts

        gen_mock = AsyncMock(side_effect=_echo_generate)
        final_state = asyncio.run(solver(state, gen_mock))

        # With no skips, solver pattern: original call + evaluation + follow‑up → 3
        self.assertEqual(gen_mock.await_count, 3)
        self.assertEqual(final_state.messages[-1].content, expected_follow_up)

    def test_dynamic_follow_up_on_match(self):
        self._run_dynamic_test("yes, sure", "Great! let's proceed.")

    def test_dynamic_follow_up_on_miss(self):
        self._run_dynamic_test("nope", "Could you clarify?")


if __name__ == "__main__":  # pragma: no cover
    pytest.main([__file__, "-q"])
