"""Test-suite for the Community‑Curator evaluation
===================================================================

Run with::

    pytest tests/agentmisalignment/test_community_curator.py -q
"""

from __future__ import annotations

import asyncio
import math
import random
import unittest
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

# --- EVAL & SOLVER IMPORTS -------------------------------------------------
from inspect_evals.agentmisalignment.evals.community_curator.community_curator import (
    _load_comments,
    _process_stage_steps,
    create_record_to_sample_fn,
)

from inspect_evals.agentmisalignment.solvers.community_curator_solver import (
    _fill_agent_names,
    _stage_steps_to_sequence,
    multi_step_solver,
)


# =============================================================================
# 1.  Tests for Dataset / Comment helpers
# =============================================================================
class TestCommunityCurator(unittest.TestCase):
    """Unit‑tests for comment loading and stage‑step processing utilities."""

    def setUp(self):
        # Deterministic RNG for the whole suite
        random.seed(42)

        # Mock comments pool consumed by helper functions.
        self.mock_comments = {
            "toxic_comments": [f"toxic_comment_{i}" for i in range(100)],
            "neutral_comments": [f"neutral_comment_{i}" for i in range(100)],
        }

        # Typical multi‑stage configuration mirroring JSON metadata.
        self.sample_stage_steps = [
            {
                "name": "early_stage",
                "total": 5,
                "pct_toxic": 0,
                "order": "pattern:nt",
            },
            {
                "name": "mid_stage",
                "total": 6,
                "pct_toxic": 0.5,
                "order": "pattern:tn",
            },
            {
                "name": "late_stage",
                "total": 4,
                "pct_toxic": 0.75,
                "order": "block:nt",
            },
        ]

    # ------------------------------------------------------------------
    # _load_comments
    # ------------------------------------------------------------------
    @patch(
        "inspect_evals.agentmisalignment.evals.community_curator.community_curator.load_dataset"
    )
    def test_load_comments(self, mock_load_dataset):
        """_load_comments correctly splits and prefixes toxic / neutral rows."""
        mock_dataset = [
            {"id": "1", "toxic": 0, "comment_text": "neutral_text_1"},
            {"id": "2", "toxic": 1, "comment_text": "toxic_text_1"},
            {"id": "3", "toxic": 0, "comment_text": "neutral_text_2"},
            {"id": "4", "toxic": 1, "comment_text": "toxic_text_2"},
        ]
        mock_load_dataset.return_value = mock_dataset

        out = _load_comments(math.inf)
        self.assertEqual(len(out["toxic_comments"]), 2)
        self.assertEqual(len(out["neutral_comments"]), 2)
        self.assertIn("^ toxic_text_1", out["toxic_comments"])
        self.assertIn("^ toxic_text_2", out["toxic_comments"])
        self.assertIn("~ neutral_text_1", out["neutral_comments"])
        self.assertIn("~ neutral_text_2", out["neutral_comments"])

    # ------------------------------------------------------------------
    # _process_stage_steps – canonical patterns
    # ------------------------------------------------------------------
    def _assert_pattern(self, stage_steps, expected):
        msgs = _process_stage_steps(stage_steps, self.mock_comments)
        self.assertEqual([m.content for m in msgs], expected)

    def test_process_stage_steps_pattern_nt(self):
        stage_steps = [
            {
                "name": "test_stage",
                "total": 6,
                "pct_toxic": 0.5,  # 3 toxic / 3 neutral
                "order": "pattern:nt",
            }
        ]
        expected = [
            "neutral_comment_0",
            "toxic_comment_0",
            "neutral_comment_1",
            "toxic_comment_1",
            "neutral_comment_2",
            "toxic_comment_2",
        ]
        self._assert_pattern(stage_steps, expected)

    def test_process_stage_steps_pattern_tn(self):
        stage_steps = [
            {
                "name": "test_stage",
                "total": 6,
                "pct_toxic": 0.5,
                "order": "pattern:tn",
            }
        ]
        expected = [
            "toxic_comment_0",
            "neutral_comment_0",
            "toxic_comment_1",
            "neutral_comment_1",
            "toxic_comment_2",
            "neutral_comment_2",
        ]
        self._assert_pattern(stage_steps, expected)

    def test_process_stage_steps_block_orders(self):
        # block:nt
        stage_steps_nt = [
            {
                "name": "test_stage",
                "total": 6,
                "pct_toxic": 0.5,
                "order": "block:nt",
            }
        ]
        expected_nt = [
            *[f"neutral_comment_{i}" for i in range(3)],
            *[f"toxic_comment_{i}" for i in range(3)],
        ]
        self._assert_pattern(stage_steps_nt, expected_nt)

        # block:tn
        stage_steps_tn = [
            {
                "name": "test_stage",
                "total": 6,
                "pct_toxic": 0.5,
                "order": "block:tn",
            }
        ]
        expected_tn = [
            *[f"toxic_comment_{i}" for i in range(3)],
            *[f"neutral_comment_{i}" for i in range(3)],
        ]
        self._assert_pattern(stage_steps_tn, expected_tn)

    def test_create_record_to_sample_fn(self):
        fn = create_record_to_sample_fn(self.mock_comments)
        record = {
            "id": "id-1",
            "target": "safe",
            "metadata": {
                "stage_steps": self.sample_stage_steps,
            },
        }
        sample: Sample = fn(record)
        self.assertEqual(sample.id, "id-1")
        self.assertEqual(sample.target, "safe")
        # 5 + 6 + 4 = 15
        self.assertEqual(len(sample.input), 15)
        assert all(isinstance(m, ChatMessageUser) for m in sample.input)


# =============================================================================
# 2.  Tests for utilities in the solver module
# =============================================================================
class TestSolverUtilities(unittest.TestCase):
    """Deterministic tests for helper utilities."""

    def test_stage_steps_to_sequence(self):
        stage_cfg = [
            {
                "name": "stage-1",
                "total": 4,
                "pct_toxic": 0.25,  # 1 toxic / 3 neutral
                "order": "pattern:nt",
            },
            {
                "name": "stage-2",
                "total": 2,
                "pct_toxic": 1.0,  # 2 toxic
                "order": "block:tn",
            },
        ]
        seq = _stage_steps_to_sequence(stage_cfg)
        # Overall length should match totals
        self.assertEqual(len(seq), 6)

        # First 4 entries come from stage‑1: expected pattern n‑t‑n‑n
        stage1_seq = [label for label, _ in seq[:4]]
        self.assertEqual(stage1_seq, ["n", "t", "n", "n"])

        # Remaining two are toxic from stage‑2 (block:tn with no neutrals)
        self.assertTrue(all(label == "t" for label, _ in seq[4:]))

    def test_fill_agent_names(self):
        names = _fill_agent_names(5)
        self.assertEqual(len(names), 5)
        self.assertEqual(len(set(names)), 5)
        self.assertTrue(all(n.startswith("agent_b") for n in names))


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