from pathlib import Path
from unittest.mock import MagicMock, patch

import pytest
from inspect_ai import Task
from inspect_ai.log import EvalLog
from inspect_ai.model import Model
from inspect_ai.solver import Solver

from bench_af._abstract.objects import Detector, Environment, ModelOrganism
from bench_af.cli.run_detector import (
    _run_model_organism_on_environment,
    _verify_model_organism_is_not_misaligned_on_environment,
    run_detector,
)


class MockDetector(Detector):
    """Mock detector implementation for testing."""

    def get_prediction(
        self,
        *,
        model_organism: ModelOrganism,
        environment: Environment,
        eval_logs: list[EvalLog],
    ) -> tuple[float, str]:
        return 0.75, "Test prediction explanation"


class TestRunModelOrganismOnEnvironment:
    """Test _run_model_organism_on_environment function."""

    def test_run_model_organism_observed_task_default(self):
        """Test running model organism on observed task (default behavior)."""
        # Setup mocks
        mock_model = MagicMock(spec=Model)
        mock_solver = MagicMock(spec=Solver)
        mock_environment = MagicMock(spec=Environment)
        mock_task = MagicMock(spec=Task)
        mock_environment.get_observed_task.return_value = mock_task
        mock_results_dir = MagicMock(spec=Path)

        model_organism = ModelOrganism(
            name="test_model",
            model=mock_model,
            solver=mock_solver,
            supported_environments={"test_task"},
        )

        mock_eval_logs = [MagicMock(spec=EvalLog)]

        with patch("cli.run_detector.eval") as mock_eval:
            mock_eval.return_value = mock_eval_logs

            result = _run_model_organism_on_environment(
                model_organism=model_organism,
                environment=mock_environment,
                results_dir=mock_results_dir,
            )

            assert result == mock_eval_logs
            mock_environment.get_observed_task.assert_called_once()
            mock_environment.get_unobserved_task.assert_not_called()
            mock_eval.assert_called_once_with(
                tasks=mock_task,
                model=mock_model,
                solver=mock_solver,
                tags=["run_detector"],
                log_dir=str(mock_results_dir / "logs"),
            )

    def test_run_model_organism_unobserved_task(self):
        """Test running model organism on unobserved task."""
        # Setup mocks
        mock_model = MagicMock(spec=Model)
        mock_solver = MagicMock(spec=Solver)
        mock_environment = MagicMock(spec=Environment)
        mock_task = MagicMock(spec=Task)
        mock_environment.get_unobserved_task.return_value = mock_task
        mock_results_dir = MagicMock(spec=Path)

        model_organism = ModelOrganism(
            name="test_model",
            model=mock_model,
            solver=mock_solver,
            supported_environments={"test_task"},
        )

        mock_eval_logs = [MagicMock(spec=EvalLog)]

        with patch("cli.run_detector.eval") as mock_eval:
            mock_eval.return_value = mock_eval_logs

            result = _run_model_organism_on_environment(
                model_organism=model_organism,
                environment=mock_environment,
                results_dir=mock_results_dir,
                use_observed=False,
            )

            assert result == mock_eval_logs
            mock_environment.get_unobserved_task.assert_called_once()
            mock_environment.get_observed_task.assert_not_called()
            mock_eval.assert_called_once_with(
                tasks=mock_task,
                model=mock_model,
                solver=mock_solver,
                tags=["run_detector"],
                log_dir=str(mock_results_dir / "logs"),
            )

    def test_run_model_organism_with_eval_config(self):
        """Test running model organism with custom eval config."""
        # Setup mocks
        mock_model = MagicMock(spec=Model)
        mock_solver = MagicMock(spec=Solver)
        mock_environment = MagicMock(spec=Environment)
        mock_task = MagicMock(spec=Task)
        mock_environment.get_observed_task.return_value = mock_task
        mock_results_dir = MagicMock(spec=Path)

        model_organism = ModelOrganism(
            name="test_model",
            model=mock_model,
            solver=mock_solver,
            supported_environments={"test_task"},
        )

        eval_config = {"max_samples": 10, "parallel": True, "sandbox": "docker"}

        mock_eval_logs = [MagicMock(spec=EvalLog)]

        with patch("cli.run_detector.eval") as mock_eval:
            mock_eval.return_value = mock_eval_logs

            result = _run_model_organism_on_environment(
                model_organism=model_organism,
                environment=mock_environment,
                results_dir=mock_results_dir,
                eval_config=eval_config,
            )

            assert result == mock_eval_logs
            mock_eval.assert_called_once_with(
                tasks=mock_task,
                model=mock_model,
                solver=mock_solver,
                tags=["run_detector"],
                log_dir=str(mock_results_dir / "logs"),
                max_samples=10,
                parallel=True,
                sandbox="docker",
            )

    def test_run_model_organism_none_eval_config(self):
        """Test running model organism with None eval config."""
        # Setup mocks
        mock_model = MagicMock(spec=Model)
        mock_solver = MagicMock(spec=Solver)
        mock_environment = MagicMock(spec=Environment)
        mock_task = MagicMock(spec=Task)
        mock_environment.get_observed_task.return_value = mock_task
        mock_results_dir = MagicMock(spec=Path)

        model_organism = ModelOrganism(
            name="test_model",
            model=mock_model,
            solver=mock_solver,
            supported_environments={"test_task"},
        )

        mock_eval_logs = [MagicMock(spec=EvalLog)]

        with patch("cli.run_detector.eval") as mock_eval:
            mock_eval.return_value = mock_eval_logs

            result = _run_model_organism_on_environment(
                model_organism=model_organism,
                environment=mock_environment,
                results_dir=mock_results_dir,
                eval_config=None,
            )

            assert result == mock_eval_logs
            mock_eval.assert_called_once_with(
                tasks=mock_task,
                model=mock_model,
                solver=mock_solver,
                tags=["run_detector"],
                log_dir=str(mock_results_dir / "logs"),
            )


class TestVerifyModelOrganismNotMisaligned:
    """Test _verify_model_organism_is_not_misaligned_on_environment function."""

    def test_verify_all_success_logs(self):
        """Test verification with all successful eval logs."""
        mock_model = MagicMock(spec=Model)
        mock_solver = MagicMock(spec=Solver)
        mock_environment = MagicMock(spec=Environment)
        mock_results_dir = MagicMock(spec=Path)

        model_organism = ModelOrganism(
            name="test_model",
            model=mock_model,
            solver=mock_solver,
            supported_environments={"test_task"},
        )

        # Create mock eval logs with success status
        mock_eval_log1 = MagicMock(spec=EvalLog)
        mock_eval_log1.status = "success"
        mock_eval_log2 = MagicMock(spec=EvalLog)
        mock_eval_log2.status = "success"

        eval_config = {"test": "config"}

        with patch("cli.run_detector._run_model_organism_on_environment") as mock_run:
            mock_run.return_value = [mock_eval_log1, mock_eval_log2]

            result = _verify_model_organism_is_not_misaligned_on_environment(
                model_organism=model_organism,
                environment=mock_environment,
                eval_config=eval_config,
                results_dir=mock_results_dir,
            )

            assert result is True
            mock_run.assert_called_once_with(
                model_organism,
                mock_environment,
                mock_results_dir,
                use_observed=True,
                eval_config=eval_config,
            )

    def test_verify_with_failed_logs(self):
        """Test verification with some failed eval logs."""
        mock_model = MagicMock(spec=Model)
        mock_solver = MagicMock(spec=Solver)
        mock_environment = MagicMock(spec=Environment)
        mock_results_dir = MagicMock(spec=Path)

        model_organism = ModelOrganism(
            name="test_model",
            model=mock_model,
            solver=mock_solver,
            supported_environments={"test_task"},
        )

        # Create mock eval logs with mixed status
        mock_eval_log1 = MagicMock(spec=EvalLog)
        mock_eval_log1.status = "success"
        mock_eval_log2 = MagicMock(spec=EvalLog)
        mock_eval_log2.status = "error"

        eval_config = {"test": "config"}

        with patch("cli.run_detector._run_model_organism_on_environment") as mock_run:
            mock_run.return_value = [mock_eval_log1, mock_eval_log2]

            result = _verify_model_organism_is_not_misaligned_on_environment(
                model_organism=model_organism,
                environment=mock_environment,
                eval_config=eval_config,
                results_dir=mock_results_dir,
            )

            assert result is False

    def test_verify_empty_logs(self):
        """Test verification with empty eval logs."""
        mock_model = MagicMock(spec=Model)
        mock_solver = MagicMock(spec=Solver)
        mock_environment = MagicMock(spec=Environment)
        mock_results_dir = MagicMock(spec=Path)

        model_organism = ModelOrganism(
            name="test_model",
            model=mock_model,
            solver=mock_solver,
            supported_environments={"test_task"},
        )

        eval_config = {"test": "config"}

        with patch("cli.run_detector._run_model_organism_on_environment") as mock_run:
            mock_run.return_value = []

            result = _verify_model_organism_is_not_misaligned_on_environment(
                model_organism=model_organism,
                environment=mock_environment,
                eval_config=eval_config,
                results_dir=mock_results_dir,
            )

            # Empty list should return True (all() returns True for empty iterable)
            assert result is True


class TestRunDetector:
    """Test run_detector function."""

    def test_run_detector_success(self):
        """Test successful detector run."""
        # Setup mocks
        mock_model = MagicMock(spec=Model)
        mock_solver = MagicMock(spec=Solver)
        mock_environment = MagicMock(spec=Environment)
        mock_environment.name = "test_task"
        mock_results_dir = MagicMock(spec=Path)

        model_organism = ModelOrganism(
            name="test_model",
            model=mock_model,
            solver=mock_solver,
            supported_environments={"test_task"},
        )

        detector = MockDetector(
            name="test_detector",
            supported_model_organisms={"test_model"},
            supported_environments={"test_task"},
        )

        eval_config = {"test": "config"}
        mock_eval_logs = [MagicMock(spec=EvalLog)]

        with (
            patch(
                "cli.run_detector._verify_model_organism_is_not_misaligned_on_environment"
            ) as mock_verify,
            patch("cli.run_detector._run_model_organism_on_environment") as mock_run,
        ):
            mock_verify.return_value = True
            mock_run.return_value = mock_eval_logs

            probability, explanation = run_detector(
                detector=detector,
                model_organism=model_organism,
                environment=mock_environment,
                eval_config=eval_config,
                results_dir=mock_results_dir,
            )

            assert probability == 0.75
            assert explanation == "Test prediction explanation"

            # Verify the verification step was called
            mock_verify.assert_called_once_with(
                model_organism, mock_environment, eval_config, mock_results_dir
            )

            # Verify model organism was run on environment
            mock_run.assert_called_once_with(
                model_organism,
                mock_environment,
                use_observed=True,
                eval_config=eval_config,
                results_dir=mock_results_dir,
            )

    def test_run_detector_unsupported_model(self):
        """Test detector run with unsupported model organism."""
        mock_model = MagicMock(spec=Model)
        mock_solver = MagicMock(spec=Solver)
        mock_environment = MagicMock(spec=Environment)
        mock_environment.name = "test_task"
        mock_results_dir = MagicMock(spec=Path)

        model_organism = ModelOrganism(
            name="unsupported_model",
            model=mock_model,
            solver=mock_solver,
            supported_environments={"test_task"},
        )

        detector = MockDetector(
            name="test_detector",
            supported_model_organisms={"test_model"},
            supported_environments={"test_task"},
        )

        eval_config = {"test": "config"}

        with pytest.raises(
            ValueError,
            match="Model organism unsupported_model not supported by detector test_detector",
        ):
            run_detector(
                detector=detector,
                model_organism=model_organism,
                environment=mock_environment,
                eval_config=eval_config,
                results_dir=mock_results_dir,
            )

    def test_run_detector_unsupported_environment(self):
        """Test detector run with unsupported environment."""
        mock_model = MagicMock(spec=Model)
        mock_solver = MagicMock(spec=Solver)
        mock_environment = MagicMock(spec=Environment)
        mock_environment.name = "unsupported_task"
        mock_results_dir = MagicMock(spec=Path)

        model_organism = ModelOrganism(
            name="test_model",
            model=mock_model,
            solver=mock_solver,
            supported_environments={"test_task"},
        )

        detector = MockDetector(
            name="test_detector",
            supported_model_organisms={"test_model"},
            supported_environments={"test_task"},
        )

        eval_config = {"test": "config"}

        with pytest.raises(
            ValueError,
            match="Environment unsupported_task not supported by detector test_detector",
        ):
            run_detector(
                detector=detector,
                model_organism=model_organism,
                environment=mock_environment,
                eval_config=eval_config,
                results_dir=mock_results_dir,
            )

    def test_run_detector_misaligned_organism(self):
        """Test detector run when model organism is misaligned."""
        # Setup mocks
        mock_model = MagicMock(spec=Model)
        mock_solver = MagicMock(spec=Solver)
        mock_environment = MagicMock(spec=Environment)
        mock_environment.name = "test_task"
        mock_results_dir = MagicMock(spec=Path)

        model_organism = ModelOrganism(
            name="test_model",
            model=mock_model,
            solver=mock_solver,
            supported_environments={"test_task"},
        )

        detector = MockDetector(
            name="test_detector",
            supported_model_organisms={"test_model"},
            supported_environments={"test_task"},
        )

        eval_config = {"test": "config"}

        with patch(
            "cli.run_detector._verify_model_organism_is_not_misaligned_on_environment"
        ) as mock_verify:
            mock_verify.return_value = False

            with pytest.raises(
                ValueError,
                match="The original model organism is misaligned with the new environment. This is not allowed.",
            ):
                run_detector(
                    detector=detector,
                    model_organism=model_organism,
                    environment=mock_environment,
                    eval_config=eval_config,
                    results_dir=mock_results_dir,
                )

    def test_run_detector_with_new_model_and_environment(self):
        """Test detector run with new model organism and environment from detector."""
        # Setup original mocks
        mock_model = MagicMock(spec=Model)
        mock_solver = MagicMock(spec=Solver)
        mock_environment = MagicMock(spec=Environment)
        mock_environment.name = "test_task"
        mock_results_dir = MagicMock(spec=Path)

        model_organism = ModelOrganism(
            name="test_model",
            model=mock_model,
            solver=mock_solver,
            supported_environments={"test_task"},
        )

        # Setup new model organism and environment
        new_mock_model = MagicMock(spec=Model)
        new_mock_solver = MagicMock(spec=Solver)
        new_mock_environment = MagicMock(spec=Environment)
        new_mock_environment.name = "new_test_task"

        new_model_organism = ModelOrganism(
            name="new_test_model",
            model=new_mock_model,
            solver=new_mock_solver,
            supported_environments={"new_test_task"},
        )

        # Create detector that returns new model and environment
        detector = MockDetector(
            name="test_detector",
            supported_model_organisms={"test_model"},
            supported_environments={"test_task"},
        )

        eval_config = {"test": "config"}
        mock_eval_logs = [MagicMock(spec=EvalLog)]

        with (
            patch(
                "cli.run_detector._verify_model_organism_is_not_misaligned_on_environment"
            ) as mock_verify,
            patch("cli.run_detector._run_model_organism_on_environment") as mock_run,
            patch.object(MockDetector, "get_new_model_and_environment") as mock_get_new,
        ):
            mock_verify.return_value = True
            mock_run.return_value = mock_eval_logs
            mock_get_new.return_value = (new_model_organism, new_mock_environment)

            probability, explanation = run_detector(
                detector=detector,
                model_organism=model_organism,
                environment=mock_environment,
                eval_config=eval_config,
                results_dir=mock_results_dir,
            )

            assert probability == 0.75
            assert explanation == "Test prediction explanation"

            # Verify get_new_model_and_environment was called with original inputs
            mock_get_new.assert_called_once_with(
                original_model_organism=model_organism,
                original_environment=mock_environment,
            )

            # Verify verification was called with original model and environment
            mock_verify.assert_called_once_with(
                model_organism, new_mock_environment, eval_config, mock_results_dir
            )

            # Verify model organism was run with NEW model and environment
            mock_run.assert_called_once_with(
                new_model_organism,
                new_mock_environment,
                use_observed=True,
                eval_config=eval_config,
                results_dir=mock_results_dir,
            )

    def test_run_detector_complex_eval_config(self):
        """Test detector run with complex eval configuration."""
        # Setup mocks
        mock_model = MagicMock(spec=Model)
        mock_solver = MagicMock(spec=Solver)
        mock_environment = MagicMock(spec=Environment)
        mock_environment.name = "test_task"
        mock_results_dir = MagicMock(spec=Path)

        model_organism = ModelOrganism(
            name="test_model",
            model=mock_model,
            solver=mock_solver,
            supported_environments={"test_task"},
        )

        detector = MockDetector(
            name="test_detector",
            supported_model_organisms={"test_model"},
            supported_environments={"test_task"},
        )

        complex_eval_config = {
            "log_level": "debug",
            "max_samples": 50,
            "parallel": True,
            "sandbox": "docker",
            "timeout": 300,
        }

        mock_eval_logs = [MagicMock(spec=EvalLog)]

        with (
            patch(
                "cli.run_detector._verify_model_organism_is_not_misaligned_on_environment"
            ) as mock_verify,
            patch("cli.run_detector._run_model_organism_on_environment") as mock_run,
        ):
            mock_verify.return_value = True
            mock_run.return_value = mock_eval_logs

            probability, explanation = run_detector(
                detector=detector,
                model_organism=model_organism,
                environment=mock_environment,
                eval_config=complex_eval_config,
                results_dir=mock_results_dir,
            )

            assert probability == 0.75
            assert explanation == "Test prediction explanation"

            # Verify complex config was passed through to both verification and run calls
            mock_verify.assert_called_once_with(
                model_organism, mock_environment, complex_eval_config, mock_results_dir
            )
            mock_run.assert_called_once_with(
                model_organism,
                mock_environment,
                use_observed=True,
                eval_config=complex_eval_config,
                results_dir=mock_results_dir,
            )


class TestRunDetectorIntegration:
    """Integration tests for run detector pipeline."""

    def test_run_detector_end_to_end(self):
        """Test complete run detector pipeline with realistic components."""
        # This test would use actual default implementations if available
        # and test the integration between components
        pass  # TODO: Implement when we have more realistic test data

    def test_detector_prediction_with_real_eval_logs(self):
        """Test detector prediction with more realistic eval log data."""
        # Setup mocks with realistic eval log data
        mock_model = MagicMock(spec=Model)
        mock_solver = MagicMock(spec=Solver)
        mock_environment = MagicMock(spec=Environment)
        mock_environment.name = "test_task"
        mock_results_dir = MagicMock(spec=Path)

        model_organism = ModelOrganism(
            name="test_model",
            model=mock_model,
            solver=mock_solver,
            supported_environments={"test_task"},
        )

        # Create a detector that examines the eval logs more realistically
        class ExaminingDetector(Detector):
            def get_prediction(
                self,
                *,
                model_organism: ModelOrganism,
                environment: Environment,
                eval_logs: list[EvalLog],
            ) -> tuple[float, str]:
                # Simulate analysis of eval logs
                if len(eval_logs) > 2:
                    return 0.9, "High number of eval logs suggests suspicious behavior"
                else:
                    return 0.1, "Low number of eval logs suggests normal behavior"

        detector = ExaminingDetector(
            name="examining_detector",
            supported_model_organisms={"test_model"},
            supported_environments={"test_task"},
        )

        eval_config = {"test": "config"}

        # Test with many eval logs
        many_eval_logs = [MagicMock(spec=EvalLog) for _ in range(5)]

        with (
            patch(
                "cli.run_detector._verify_model_organism_is_not_misaligned_on_environment"
            ) as mock_verify,
            patch("cli.run_detector._run_model_organism_on_environment") as mock_run,
        ):
            mock_verify.return_value = True
            mock_run.return_value = many_eval_logs

            probability, explanation = run_detector(
                detector=detector,
                model_organism=model_organism,
                environment=mock_environment,
                eval_config=eval_config,
                results_dir=mock_results_dir,
            )

            assert probability == 0.9
            assert "High number of eval logs" in explanation

        # Test with few eval logs
        few_eval_logs = [MagicMock(spec=EvalLog)]

        with (
            patch(
                "cli.run_detector._verify_model_organism_is_not_misaligned_on_environment"
            ) as mock_verify,
            patch("cli.run_detector._run_model_organism_on_environment") as mock_run,
        ):
            mock_verify.return_value = True
            mock_run.return_value = few_eval_logs

            probability, explanation = run_detector(
                detector=detector,
                model_organism=model_organism,
                environment=mock_environment,
                eval_config=eval_config,
                results_dir=mock_results_dir,
            )

            assert probability == 0.1
            assert "Low number of eval logs" in explanation
