"""
Test agents functionality.
"""

import pytest
from unittest.mock import AsyncMock, MagicMock

from app.agents.planner import PlannerAgent, PlannerInput
from app.agents.compiler import CompilerAgent, CompilerInput
from app.agents.validator import ValidatorAgent, ValidatorInput
from app.domain.models.ir import DatasetSchema, ColumnSchema, OperationType


@pytest.fixture
def sample_schema():
    """Sample dataset schema."""
    return DatasetSchema(
        columns=[
            ColumnSchema(name="customer_id", dtype="int64", nullable=False),
            ColumnSchema(name="age", dtype="int64", nullable=True),
            ColumnSchema(name="email", dtype="string", nullable=True),
        ]
    )


class TestPlannerAgent:
    """Tests for Planner agent."""

    @pytest.mark.asyncio
    async def test_planner_execution(self, sample_schema, monkeypatch):
        """Test planner can parse natural language."""
        # Mock LLM response
        mock_response = """```json
[
    {"operation": "drop_duplicates", "params": {}, "rationale": "Remove duplicate rows"},
    {"operation": "fill_missing", "params": {"column": "age", "strategy": "median"}, "rationale": "Fill missing ages"}
]
```"""

        async def mock_call_llm(self, prompt):
            return mock_response

        monkeypatch.setattr(PlannerAgent, "_call_llm", mock_call_llm)

        planner = PlannerAgent()
        input_data = PlannerInput(
            natural_language="Remove duplicates and fill missing ages with median",
            dataset_schema=sample_schema,
        )

        result = await planner.run(input_data)

        assert result.success
        assert result.output is not None
        assert len(result.output.operations) == 2
        assert result.output.operations[0].operation == OperationType.DROP_DUPLICATES
        assert result.output.operations[1].operation == OperationType.FILL_MISSING


class TestCompilerAgent:
    """Tests for Compiler agent."""

    @pytest.mark.asyncio
    async def test_compiler_builds_dag(self, sample_schema):
        """Test compiler builds valid DAG from operations."""
        from app.agents.planner import OperationPlan
        from uuid import uuid4

        compiler = CompilerAgent()
        
        operations = [
            OperationPlan(
                operation=OperationType.DROP_DUPLICATES,
                params={},
                rationale="Clean data",
            ),
            OperationPlan(
                operation=OperationType.FILL_MISSING,
                params={"column": "age", "strategy": "median"},
                rationale="Fill missing values",
            ),
        ]

        input_data = CompilerInput(
            operations=operations,
            dataset_id=uuid4(),
            dataset_schema=sample_schema,
        )

        result = await compiler.run(input_data)

        assert result.success
        assert result.output is not None
        assert len(result.output.ir.dag.nodes) == 2
        assert len(result.output.ir.dag.execution_order) == 2
        
        # Validate DAG
        is_valid, message = result.output.ir.validate()
        assert is_valid


class TestValidatorAgent:
    """Tests for Validator agent."""

    @pytest.mark.asyncio
    async def test_validator_checks_nulls(self, tmp_path):
        """Test validator detects null values."""
        import pandas as pd
        from app.domain.models.ir import ValidationRule

        # Create test data with nulls
        df = pd.DataFrame({
            "customer_id": [1, 2, 3],
            "age": [25, None, 30],
        })

        # Save to file
        test_file = tmp_path / "test.csv"
        df.to_csv(test_file, index=False)

        validator = ValidatorAgent()
        input_data = ValidatorInput(
            data_path=test_file,
            validation_rules=[
                ValidationRule(rule="no_nulls", columns=["customer_id", "age"])
            ],
            engine="pandas",
        )

        result = await validator.run(input_data)

        assert result.success
        assert result.output is not None
        # Should fail validation due to null in age column
        assert not result.output.is_valid
        assert len(result.output.issues) > 0


def test_operation_types():
    """Test all operation types are defined."""
    operations = [
        "drop_duplicates",
        "fill_missing",
        "normalize",
        "encode_categorical",
        "filter_rows",
    ]

    for op in operations:
        assert hasattr(OperationType, op.upper())
