import os
import shutil
import tempfile
from pathlib import Path
from unittest.mock import MagicMock, patch

import pytest

from bench_af._abstract.objects import Detector, ModelOrganism
from bench_af.util import (
    create_new_detector,
    create_new_environment,
    create_new_model_organism,
    list_detectors,
    list_environments,
    list_model_organisms,
    load_detector,
    load_environment,
    load_model_organism,
)
from util.list import _list_objects


class TestRegistryListFunctions:
    """Test listing functions in the registry."""

    def test_list_objects(self):
        """Test _list_objects helper function."""
        # Create temporary directory structure
        with tempfile.TemporaryDirectory() as temp_dir:
            temp_path = Path(temp_dir)

            # Create test category directory with subdirectories
            category_dir = temp_path / "models"
            category_dir.mkdir()

            # Create subdirectories with and without main.py
            (category_dir / "model1").mkdir()
            with open(category_dir / "model1" / "main.py", "w") as f:
                f.write("""from bench_af._abstract.objects import ModelOrganism
from inspect_ai.model import get_model
from inspect_ai.solver import generate

def get() -> ModelOrganism:
    return ModelOrganism(
        name="model1",
        model=get_model(model="openai/gpt-4o-mini"),
        solver=[generate()],
        supported_environments=set(),
    )
""")

            (category_dir / "model2").mkdir()
            with open(category_dir / "model2" / "main.py", "w") as f:
                f.write("""from bench_af._abstract.objects import ModelOrganism
from inspect_ai.model import get_model
from inspect_ai.solver import generate

def get() -> ModelOrganism:
    return ModelOrganism(
        name="model2",
        model=get_model(model="openai/gpt-4o-mini"),
        solver=[generate()],
        supported_environments=set(),
    )
""")

            (category_dir / "incomplete_model").mkdir()
            # No main.py in incomplete_model

            (category_dir / "not_a_dir.py").touch()  # File, not directory

            with patch("util.list._REGISTRY_DIR", temp_path):
                result = _list_objects("models")
                assert set(result) == {"model1", "model2"}

    @patch("util.list._list_objects")
    def test_list_model_organisms(self, mock_list_objects):
        """Test list_model_organisms function."""
        mock_list_objects.return_value = ["model1", "model2", "default"]
        result = list_model_organisms()
        assert result == ["model1", "model2", "default"]
        mock_list_objects.assert_called_once_with("models")

    @patch("util.list._list_objects")
    def test_list_detectors(self, mock_list_objects):
        """Test list_detectors function."""
        mock_list_objects.return_value = ["detector1", "detector2"]
        result = list_detectors()
        assert result == ["detector1", "detector2"]
        mock_list_objects.assert_called_once_with("detectors")

    @patch("util.list._list_objects")
    def test_list_environments(self, mock_list_objects):
        """Test list_environments function."""
        mock_list_objects.return_value = ["env1", "env2", "default"]
        result = list_environments()
        assert result == ["env1", "env2", "default"]
        mock_list_objects.assert_called_once_with("environments")


class TestRegistryLoadFunctions:
    """Test loading functions in the registry."""

    @patch("util.load.import_module")
    @patch("util.load._list_objects")
    def test_load_model_organism_success(self, mock_list, mock_import):
        """Test successful model organism loading."""
        # Setup mocks
        mock_list.return_value = ["test_model"]
        mock_module = MagicMock()
        mock_organism = MagicMock(spec=ModelOrganism)
        mock_module.get.return_value = mock_organism
        mock_import.return_value = mock_module

        result = load_model_organism("test_model")

        assert result == mock_organism
        mock_import.assert_called_once_with("models.test_model.main")
        mock_module.get.assert_called_once()

    @patch("util.load._list_objects")
    def test_load_model_organism_not_in_registry(self, mock_list):
        """Test loading non-existent model organism."""
        mock_list.return_value = ["existing_model"]

        with pytest.raises(
            ValueError, match="Models nonexistent is not in the registry"
        ):
            load_model_organism("nonexistent")

    @patch("util.load.import_module")
    @patch("util.load._list_objects")
    def test_load_model_organism_missing_attribute(self, mock_list, mock_import):
        """Test loading model organism without get attribute."""
        mock_list.return_value = ["test_model"]
        mock_module = MagicMock()
        del mock_module.get  # Remove the attribute
        mock_import.return_value = mock_module

        with pytest.raises(
            AssertionError, match="Models test_model must have a `get` attribute"
        ):
            load_model_organism("test_model")

    @patch("util.load.import_module")
    @patch("util.load._list_objects")
    def test_load_model_organism_wrong_type(self, mock_list, mock_import):
        """Test loading model organism that returns wrong type."""
        mock_list.return_value = ["test_model"]
        mock_module = MagicMock()
        mock_module.get.return_value = "not_a_model_organism"
        mock_import.return_value = mock_module

        with pytest.raises(
            AssertionError, match="Models test_model must return a ModelOrganism, Detector, or Environment"
        ):
            load_model_organism("test_model")

    @patch("util.load.import_module")
    @patch("util.load._list_objects")
    def test_load_detector_success(self, mock_list, mock_import):
        """Test successful detector loading."""
        # Setup mocks
        mock_list.return_value = ["test_detector"]
        mock_module = MagicMock()
        mock_detector = MagicMock(spec=Detector)
        mock_module.get.return_value = mock_detector
        mock_import.return_value = mock_module

        result = load_detector("test_detector")

        assert result == mock_detector
        mock_import.assert_called_once_with("detectors.test_detector.main")
        mock_module.get.assert_called_once()

    @patch("util.load._list_objects")
    def test_load_detector_not_in_registry(self, mock_list):
        """Test loading non-existent detector."""
        mock_list.return_value = ["existing_detector"]

        with pytest.raises(
            ValueError, match="Detectors nonexistent is not in the registry"
        ):
            load_detector("nonexistent")


class TestRegistryEndToEnd:
    """End-to-end tests for registry functionality using defaults."""

    @pytest.mark.skipif(
        not os.getenv("OPENAI_API_KEY") and not os.getenv("ANTHROPIC_API_KEY"),
        reason="API keys required for model initialization"
    )
    def test_create_and_load_from_defaults(self):
        """Test creating new objects from defaults and loading them."""

        # Create temporary names for testing
        test_model_name = "test_e2e_model"
        test_detector_name = "test_e2e_detector"
        test_task_name = "test_e2e_task"

        try:
            # Test creating a new model organism from defaults
            create_new_model_organism(test_model_name)

            # Verify it appears in the list
            models = list_model_organisms()
            assert test_model_name in models

            # Test that we can load it
            model_organism = load_model_organism(test_model_name)
            assert isinstance(model_organism, ModelOrganism)

            # Test creating a new detector from defaults
            create_new_detector(test_detector_name)

            # Verify it appears in the list
            detectors = list_detectors()
            assert test_detector_name in detectors

            # Test that we can load it
            detector = load_detector(test_detector_name)
            assert isinstance(detector, Detector)

            # Test creating a new environment from defaults
            create_new_environment(test_task_name)

            # Verify it appears in the list
            environments = list_environments()
            assert test_task_name in environments

            # Test that we can load it
            load_environment(test_task_name)
            # Note: load_environment returns Environment in the current implementation

        finally:
            # Clean up - remove created test objects
            from util.list import _REGISTRY_DIR

            for name, category in [
                (test_model_name, "models"),
                (test_detector_name, "detectors"),
                (test_task_name, "environments"),
            ]:
                test_dir = _REGISTRY_DIR / category / name
                if test_dir.exists():
                    shutil.rmtree(test_dir)
