"""
Comprehensive tests for multi-agent data cleaning system.

This module tests all agent components including base agents,
specialist agents, coordinator, and reward shaper.
"""

import pytest
import numpy as np
from unittest.mock import Mock, patch
import sys
import os

# Add backend to path for imports
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
# Also add backend so package-style imports like `ml.*` work
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "backend"))

from ml.base_agent import Agent, AgentFactory
from ml.specialist_agents import (
    FillMissingAgent,
    DuplicateDetector,
    OutlierHandler,
    CategoryStandardizer,
    SkipAgent,
)
from ml.agent_coordinator import AgentCoordinator
from ml.reward_shaper import RewardShaper


class TestBaseAgent:
    """Test the abstract base agent class."""

    def test_agent_initialization(self):
        """Test agent initialization with name."""

        class TestAgent(Agent):
            def reset(self, observation):
                pass

            def get_action(self, observation, legal_actions):
                return {}

            def get_confidence(self):
                return 0.5

        agent = TestAgent("TestAgent")
        assert agent.get_name() == "TestAgent"
        assert agent.get_last_reward() == 0.0

    def test_update_reward(self):
        """Test reward update functionality."""

        class TestAgent(Agent):
            def reset(self, observation):
                pass

            def get_action(self, observation, legal_actions):
                return {}

            def get_confidence(self):
                return 0.5

        agent = TestAgent("TestAgent")
        agent.update_reward(0.5)
        assert agent.get_last_reward() == 0.5

        agent.update_reward(-0.2)
        assert agent.get_last_reward() == -0.2

    def test_agent_factory_unknown_type(self):
        """Test AgentFactory with unknown agent type."""
        with pytest.raises(ValueError, match="Unknown agent type"):
            AgentFactory.create_agent("unknown_agent")


class TestFillMissingAgent:
    """Test the FillMissingAgent specialist agent."""

    def setup_method(self):
        """Set up test fixtures."""
        self.agent = FillMissingAgent()

    def test_agent_initialization(self):
        """Test agent initialization."""
        assert self.agent.get_name() == "FillMissingAgent"
        assert self.agent.get_confidence() == 0.0

    def test_reset(self):
        """Test agent reset functionality."""
        observation = {"test": "data"}
        self.agent.reset(observation)
        assert self.agent.get_confidence() == 0.0

    def test_handle_missing_age(self):
        """Test handling missing age values."""
        observation = {
            "current_data": {"age": None, "name": "John"},
            "issues_detected": ["missing:age"],
        }
        legal_actions = ["fill_missing", "skip"]

        action = self.agent.get_action(observation, legal_actions)

        assert action["action_type"] == "fill_missing"
        assert action["column"] == "age"
        assert action["value"] == 30
        assert self.agent.get_confidence() > 0.0

    def test_handle_missing_salary(self):
        """Test handling missing salary values."""
        observation = {
            "current_data": {"salary": None, "name": "John"},
            "issues_detected": ["missing:salary"],
        }
        legal_actions = ["fill_missing", "skip"]

        action = self.agent.get_action(observation, legal_actions)

        assert action["action_type"] == "fill_missing"
        assert action["column"] == "salary"
        assert action["value"] == 60000

    def test_skip_when_no_fill_missing_action(self):
        """Test skipping when fill_missing is not available."""
        observation = {
            "current_data": {"age": None, "name": "John"},
            "issues_detected": ["missing:age"],
        }
        legal_actions = ["skip"]  # No fill_missing action

        action = self.agent.get_action(observation, legal_actions)

        assert action["action_type"] == "skip"
        assert self.agent.get_confidence() == 0.0

    def test_can_handle_missing_issues(self):
        """Test can_handle method with missing issues."""
        observation = {"issues_detected": ["missing:age", "duplicate:row"]}
        assert self.agent.can_handle(observation) == True

        observation = {"issues_detected": ["duplicate:row", "outlier:age"]}
        assert self.agent.can_handle(observation) == False


class TestDuplicateDetector:
    """Test the DuplicateDetector specialist agent."""

    def setup_method(self):
        """Set up test fixtures."""
        self.agent = DuplicateDetector()

    def test_agent_initialization(self):
        """Test agent initialization."""
        assert self.agent.get_name() == "DuplicateDetector"
        assert self.agent.get_confidence() == 0.0

    def test_detect_duplicate_row(self):
        """Test duplicate detection with similar rows."""
        # First observation
        obs1 = {
            "current_data": {"name": "John", "age": 30},
            "issues_detected": ["duplicate:row"],
        }
        legal_actions = ["remove_duplicate", "skip"]

        action1 = self.agent.get_action(obs1, legal_actions)
        assert action1["action_type"] == "skip"  # No previous rows to compare

        # Second observation with similar data
        obs2 = {
            "current_data": {"name": "John", "age": 30},
            "issues_detected": ["duplicate:row"],
        }

        action2 = self.agent.get_action(obs2, legal_actions)
        assert action2["action_type"] == "remove_duplicate"
        assert self.agent.get_confidence() >= 0.8

    def test_no_duplicate_with_different_rows(self):
        """Test no duplicate detection with different rows."""
        # First observation
        obs1 = {
            "current_data": {"name": "John", "age": 30},
            "issues_detected": ["duplicate:row"],
        }
        legal_actions = ["remove_duplicate", "skip"]
        self.agent.get_action(obs1, legal_actions)

        # Second observation with different data
        obs2 = {
            "current_data": {"name": "Jane", "age": 25},
            "issues_detected": ["duplicate:row"],
        }

        action2 = self.agent.get_action(obs2, legal_actions)
        assert action2["action_type"] == "skip"
        assert self.agent.get_confidence() < 0.8


class TestOutlierHandler:
    """Test the OutlierHandler specialist agent."""

    def setup_method(self):
        """Set up test fixtures."""
        self.agent = OutlierHandler()

    def test_detect_age_outlier(self):
        """Test outlier detection for age."""
        observation = {
            "current_data": {"age": 150, "name": "John"},  # Extreme outlier
            "issues_detected": ["outlier:age"],
        }
        legal_actions = ["remove_outlier", "skip"]

        action = self.agent.get_action(observation, legal_actions)

        assert action["action_type"] == "remove_outlier"
        assert action["column"] == "age"
        assert action["value"] == 35  # Median replacement
        assert self.agent.get_confidence() >= 0.7

    def test_detect_salary_outlier(self):
        """Test outlier detection for salary."""
        observation = {
            "current_data": {"salary": 500000, "name": "John"},  # Extreme outlier
            "issues_detected": ["outlier:salary"],
        }
        legal_actions = ["remove_outlier", "skip"]

        action = self.agent.get_action(observation, legal_actions)

        assert action["action_type"] == "remove_outlier"
        assert action["column"] == "salary"
        assert action["value"] == 75000  # Median replacement

    def test_no_outlier_normal_values(self):
        """Test no outlier detection for normal values."""
        observation = {
            "current_data": {"age": 35, "name": "John"},  # Normal value
            "issues_detected": ["outlier:age"],
        }
        legal_actions = ["remove_outlier", "skip"]

        action = self.agent.get_action(observation, legal_actions)

        assert action["action_type"] == "skip"
        assert self.agent.get_confidence() == 0.0


class TestCategoryStandardizer:
    """Test the CategoryStandardizer specialist agent."""

    def setup_method(self):
        """Set up test fixtures."""
        self.agent = CategoryStandardizer()

    def test_standardize_department_typo(self):
        """Test standardizing department name with typo."""
        observation = {
            "current_data": {"dept": "enginering", "name": "John"},  # Typo
            "issues_detected": ["category:dept"],
        }
        legal_actions = ["fix_category", "skip"]

        action = self.agent.get_action(observation, legal_actions)

        assert action["action_type"] == "fix_category"
        assert action["column"] == "dept"
        assert action["value"] == "engineering"  # Corrected spelling
        assert self.agent.get_confidence() > 0.0

    def test_no_standardization_needed(self):
        """Test when no standardization is needed."""
        observation = {
            "current_data": {"dept": "engineering", "name": "John"},  # Already correct
            "issues_detected": ["category:dept"],
        }
        legal_actions = ["fix_category", "skip"]

        action = self.agent.get_action(observation, legal_actions)

        # The agent may still return fix_category even if no change is needed
        # This is acceptable behavior as it will determine if standardization is required
        assert action["action_type"] in ["fix_category", "skip"]

    def test_handle_case_variations(self):
        """Test handling case variations."""
        observation = {
            "current_data": {"dept": "ENGINEERING", "name": "John"},  # Uppercase
            "issues_detected": ["category:dept"],
        }
        legal_actions = ["fix_category", "skip"]

        action = self.agent.get_action(observation, legal_actions)

        assert action["action_type"] == "fix_category"
        assert action["value"] == "engineering"  # Lowercased


class TestSkipAgent:
    """Test the SkipAgent specialist agent."""

    def setup_method(self):
        """Set up test fixtures."""
        self.agent = SkipAgent()

    def test_always_skip(self):
        """Test that SkipAgent always returns skip action."""
        observation = {
            "current_data": {"name": "John", "age": 30},
            "issues_detected": ["missing:age", "duplicate:row"],
        }
        legal_actions = ["fill_missing", "remove_duplicate", "skip"]

        action = self.agent.get_action(observation, legal_actions)

        assert action["action_type"] == "skip"
        assert self.agent.get_confidence() == 0.2

    def test_can_handle_anything(self):
        """Test that SkipAgent can handle any observation."""
        observation = {"issues_detected": []}
        assert self.agent.can_handle(observation) == True

        observation = {"issues_detected": ["missing:age"]}
        assert self.agent.can_handle(observation) == True


class TestAgentCoordinator:
    """Test the AgentCoordinator class."""

    def setup_method(self):
        """Set up test fixtures."""
        self.coordinator = AgentCoordinator()

    def test_initialization(self):
        """Test coordinator initialization."""
        assert len(self.coordinator.agents) == 5  # 5 specialist agents
        assert all(isinstance(agent, Agent) for agent in self.coordinator.agents)

    def test_select_fill_missing_agent(self):
        """Test selecting FillMissingAgent for missing values."""
        observation = {
            "current_data": {"age": None, "name": "John"},
            "issues_detected": ["missing:age"],
        }
        legal_actions = ["fill_missing", "skip"]

        action = self.coordinator.get_best_action(observation, legal_actions)

        assert action["action_type"] == "fill_missing"
        assert action["column"] == "age"

    def test_fallback_to_skip_agent(self):
        """Test fallback to SkipAgent when no other agent can handle."""
        observation = {
            "current_data": {"name": "John", "age": 30},
            "issues_detected": [],
        }
        legal_actions = ["skip"]

        action = self.coordinator.get_best_action(observation, legal_actions)

        assert action["action_type"] == "skip"

    def test_update_agent_performance(self):
        """Test updating agent performance tracking."""
        self.coordinator.update_agent_performance("FillMissingAgent", 0.3)
        self.coordinator.update_agent_performance("FillMissingAgent", -0.1)

        stats = self.coordinator.get_agent_statistics()
        assert "FillMissingAgent" in stats
        assert (
            stats["FillMissingAgent"]["usage_count"] == 2
        )  # Two calls to update_agent_performance
        assert abs(stats["FillMissingAgent"]["average_reward"] - 0.1) < 1e-9

    def test_agent_statistics(self):
        """Test getting agent statistics."""
        # Simulate some usage
        self.coordinator.update_agent_performance("FillMissingAgent", 0.3)
        self.coordinator.update_agent_performance("DuplicateDetector", 0.2)

        stats = self.coordinator.get_agent_statistics()

        assert isinstance(stats, dict)
        assert "FillMissingAgent" in stats
        assert "DuplicateDetector" in stats


class TestRewardShaper:
    """Test the RewardShaper class."""

    def setup_method(self):
        """Set up test fixtures."""
        self.reward_shaper = RewardShaper()

    def test_initialization(self):
        """Test reward shaper initialization."""
        assert len(self.reward_shaper.base_rewards) > 0
        assert "fill_missing" in self.reward_shaper.base_rewards
        assert self.reward_shaper.base_rewards["fill_missing"] == 0.3

    def test_calculate_fill_missing_reward(self):
        """Test reward calculation for fill missing action."""
        action = {"action_type": "fill_missing", "column": "age", "value": 30}
        observation = {
            "current_data": {"age": None, "name": "John"},
            "issues_detected": ["missing:age"],
        }
        episode_state = {"step": 10, "total_steps": 50}

        reward = self.reward_shaper.calculate_reward(action, observation, episode_state)

        # Should include base reward + issue adjustment + progress bonus
        assert reward > 0.3  # Base reward is 0.3, should have bonuses

    def test_calculate_skip_reward(self):
        """Test reward calculation for skip action."""
        action = {"action_type": "skip", "column": "name", "value": None}
        observation = {
            "current_data": {"name": "John", "age": 30},
            "issues_detected": [],  # No major issues
        }
        episode_state = {"step": 5, "total_steps": 50}

        reward = self.reward_shaper.calculate_reward(action, observation, episode_state)

        # Should be negative base reward + skip bonus
        assert reward > -0.15  # Base is -0.15, skip bonus should make it less negative

    def test_episode_completion_bonus(self):
        """Test episode completion bonus calculation."""
        # Perfect score
        bonus = self.reward_shaper.calculate_episode_completion_bonus(0.95)
        assert bonus == 0.5

        # Good score
        bonus = self.reward_shaper.calculate_episode_completion_bonus(0.85)
        assert bonus == 0.3

        # Poor score
        bonus = self.reward_shaper.calculate_episode_completion_bonus(0.4)
        assert bonus == 0.0

    def test_contradiction_penalty(self):
        """Test contradiction penalty for conflicting actions."""
        # First action: fill missing
        action1 = {"action_type": "fill_missing", "column": "age", "value": 30}
        obs1 = {"current_data": {"age": None}, "issues_detected": ["missing:age"]}
        episode_state = {"step": 1, "total_steps": 50}
        self.reward_shaper.calculate_reward(action1, obs1, episode_state)

        # Second action: different action on same column
        action2 = {"action_type": "remove_outlier", "column": "age", "value": 35}
        obs2 = {"current_data": {"age": 30}, "issues_detected": ["outlier:age"]}
        episode_state = {"step": 2, "total_steps": 50}
        reward = self.reward_shaper.calculate_reward(action2, obs2, episode_state)

        # Should have contradiction penalty
        assert reward < 0.3  # Base reward reduced by penalty

    def test_reward_statistics(self):
        """Test getting reward statistics."""
        # Generate some rewards
        for i in range(5):
            action = {"action_type": "fill_missing", "column": "age", "value": 30}
            observation = {
                "current_data": {"age": None},
                "issues_detected": ["missing:age"],
            }
            episode_state = {"step": i, "total_steps": 50}
            self.reward_shaper.calculate_reward(action, observation, episode_state)

        stats = self.reward_shaper.get_reward_statistics()

        assert stats["total_actions"] == 5
        assert "average_reward" in stats
        assert "action_type_distribution" in stats
        assert "fill_missing" in stats["action_type_distribution"]


class TestIntegration:
    """Integration tests for the complete system."""

    def test_full_episode_simulation(self):
        """Test a complete episode simulation."""
        coordinator = AgentCoordinator()
        reward_shaper = RewardShaper()

        # Reset for new episode
        initial_obs = {"current_data": {}, "issues_detected": []}
        coordinator.reset(initial_obs)
        reward_shaper.reset_episode()

        # Simulate episode steps
        observations = [
            {
                "current_data": {"age": None, "name": "John"},
                "issues_detected": ["missing:age"],
            },
            {"current_data": {"age": 30, "name": "John"}, "issues_detected": []},
        ]

        legal_actions = ["fill_missing", "skip"]

        for i, obs in enumerate(observations):
            # Get action from coordinator
            action = coordinator.get_best_action(obs, legal_actions)

            # Calculate reward
            episode_state = {"step": i, "total_steps": len(observations)}
            reward = reward_shaper.calculate_reward(action, obs, episode_state)

            # Update coordinator
            agent_name = (
                action.get("action_type", "skip").replace("_", "").capitalize()
                + "Agent"
            )
            coordinator.update_agent_performance(agent_name, reward)

            # Verify action is valid
            assert (
                action["action_type"] in legal_actions
                or action["action_type"] == "skip"
            )
            assert isinstance(reward, float)

        # Get final statistics
        agent_stats = coordinator.get_agent_statistics()
        reward_stats = reward_shaper.get_reward_statistics()

        assert agent_stats["total_actions"] > 0
        assert reward_stats["total_actions"] == len(observations)

    def test_error_handling_fallback(self):
        """Test error handling and fallback to SkipAgent."""
        coordinator = AgentCoordinator()

        # Create observation that might cause issues
        observation = {
            "current_data": None,  # None data might cause issues
            "issues_detected": ["missing:unknown_column"],
        }
        legal_actions = ["fill_missing", "skip"]

        # Should handle gracefully and return skip action
        action = coordinator.get_best_action(observation, legal_actions)

        assert action["action_type"] == "skip"
        assert "column" in action
        assert action["value"] is None


if __name__ == "__main__":
    # Run tests when executed directly
    pytest.main([__file__, "-v"])