"""
Unit tests for task ID deduplication and alignment module.
"""

import pytest
from typing import Dict, List, Set, Any
from collections import defaultdict

from affine.dedup import (
    DeduplicationConfig,
    TaskIdDeduplicator,
    SamplingSetSelector,
    SampleAligner,
    TaskIdAlignmentOrchestrator,
)
from affine.models import CompactResult


class MockResult:
    """Mock result object for testing."""
    
    def __init__(self, hotkey: str, env: str, score: float, task_id: int = None, timestamp: float = 0.0, block: int = 0):
        self.hotkey = hotkey
        self.env = env
        self.score = score
        self.task_id = task_id
        self.timestamp = timestamp
        self.block = block
        
        # Create mock miner interface
        self._miner = MockMiner(hotkey, block)
        self._challenge = MockChallenge(env)
        self._evaluation = MockEvaluation(score)
    
    @property
    def miner(self):
        return self._miner
    
    @property
    def challenge(self):
        return self._challenge
    
    @property
    def evaluation(self):
        return self._evaluation


class MockMiner:
    def __init__(self, hotkey: str, block: int):
        self.hotkey = hotkey
        self.block = block
        self.model = "test/affine-model"
        self.revision = "v1.0"
        self.uid = 1


class MockChallenge:
    def __init__(self, env: str):
        self.env = env


class MockEvaluation:
    def __init__(self, score: float):
        self.score = score


class TestTaskIdDeduplicator:
    """Tests for TaskIdDeduplicator class."""
    
    def test_deduplicate_empty_results(self):
        """Test deduplication with empty results."""
        deduplicator = TaskIdDeduplicator()
        results = deduplicator.deduplicate_results([], 1000)
        assert results == []
    
    def test_deduplicate_no_duplicates(self):
        """Test deduplication when there are no duplicates."""
        deduplicator = TaskIdDeduplicator()
        results = [
            MockResult("hk1", "env1", 0.8, task_id=1, timestamp=100),
            MockResult("hk1", "env1", 0.9, task_id=2, timestamp=101),
            MockResult("hk1", "env1", 0.7, task_id=3, timestamp=102),
        ]
        deduped = deduplicator.deduplicate_results(results, 1000)
        assert len(deduped) == 3
    
    def test_deduplicate_with_duplicates_large_dataset(self):
        """Test deduplication keeps newest for large datasets (>= 400)."""
        deduplicator = TaskIdDeduplicator()
        results = [
            MockResult("hk1", "env1", 0.8, task_id=1, timestamp=100),
            MockResult("hk1", "env1", 0.9, task_id=1, timestamp=200),  # Newer, same task_id
            MockResult("hk1", "env1", 0.7, task_id=2, timestamp=101),
        ]
        deduped = deduplicator.deduplicate_results(results, 1000)
        assert len(deduped) == 2
        # Check that the newer one (timestamp=200) is kept
        task1_results = [r for r in deduped if r.task_id == 1]
        assert len(task1_results) == 1
        assert task1_results[0].timestamp == 200
    
    def test_deduplicate_with_duplicates_small_dataset(self):
        """Test deduplication keeps 2 samples per task_id for small datasets (< 400)."""
        deduplicator = TaskIdDeduplicator()
        results = [
            MockResult("hk1", "env1", 0.8, task_id=1, timestamp=100),
            MockResult("hk1", "env1", 0.9, task_id=1, timestamp=200),
            MockResult("hk1", "env1", 0.7, task_id=1, timestamp=300),  # Third one should be dropped
            MockResult("hk1", "env1", 0.6, task_id=2, timestamp=101),
        ]
        deduped = deduplicator.deduplicate_results(results, 300)  # Small dataset
        assert len(deduped) == 3
        task1_results = [r for r in deduped if r.task_id == 1]
        assert len(task1_results) == 2
        # Check that the two newest are kept (timestamps 300 and 200)
        timestamps = sorted([r.timestamp for r in task1_results], reverse=True)
        assert timestamps == [300, 200]
    
    def test_deduplicate_preserves_none_taskid(self):
        """Test that results without task_id are preserved."""
        deduplicator = TaskIdDeduplicator()
        results = [
            MockResult("hk1", "env1", 0.8, task_id=None, timestamp=100),
            MockResult("hk1", "env1", 0.9, task_id=None, timestamp=200),
            MockResult("hk1", "env1", 0.7, task_id=1, timestamp=101),
        ]
        deduped = deduplicator.deduplicate_results(results, 1000)
        assert len(deduped) == 3  # All preserved


class TestSamplingSetSelector:
    """Tests for SamplingSetSelector class."""
    
    def test_empty_miner_taskids(self):
        """Test with empty miner taskids."""
        selector = SamplingSetSelector()
        sampling_set, qualified, ratios = selector.select_sampling_set({}, 1000)
        assert sampling_set == set()
        assert qualified == set()
        assert ratios == {}
    
    def test_single_miner(self):
        """Test with single miner."""
        selector = SamplingSetSelector()
        miner_taskids = {
            "hk1": {1, 2, 3, 4, 5}
        }
        sampling_set, qualified, ratios = selector.select_sampling_set(miner_taskids, 10)
        # Should use all available task_ids since less than min size
        assert sampling_set == {1, 2, 3, 4, 5}
        assert "hk1" in qualified
        assert ratios["hk1"] == 1.0
    
    def test_multiple_miners_high_overlap(self):
        """Test with multiple miners having high overlap."""
        selector = SamplingSetSelector()
        # All miners have similar task_ids
        miner_taskids = {
            "hk1": set(range(1, 501)),  # 1-500
            "hk2": set(range(1, 501)),  # 1-500
            "hk3": set(range(1, 501)),  # 1-500
        }
        sampling_set, qualified, ratios = selector.select_sampling_set(miner_taskids, 1000)
        # Should select 500 task_ids (50% of 1000)
        assert len(sampling_set) >= 400
        assert len(sampling_set) <= 500
        # All miners should be qualified
        assert len(qualified) == 3
        for hk in ["hk1", "hk2", "hk3"]:
            assert ratios[hk] >= 0.8
    
    def test_multiple_miners_low_overlap(self):
        """Test with multiple miners having low overlap."""
        selector = SamplingSetSelector()
        # Miners have different task_ids
        miner_taskids = {
            "hk1": set(range(1, 201)),    # 1-200
            "hk2": set(range(201, 401)),  # 201-400
            "hk3": set(range(401, 601)),  # 401-600
        }
        sampling_set, qualified, ratios = selector.select_sampling_set(miner_taskids, 1000)
        # Greedy algorithm should pick most common task_ids
        # Since no overlap, each task_id has frequency 1
        assert len(sampling_set) >= 400
        # At least 50% of miners should be qualified (relaxed constraints)
        assert len(qualified) >= 2
    
    def test_greedy_selection_prioritizes_common(self):
        """Test that greedy algorithm prioritizes common task_ids."""
        selector = SamplingSetSelector()
        miner_taskids = {
            "hk1": {1, 2, 3, 100, 101},
            "hk2": {1, 2, 3, 200, 201},
            "hk3": {1, 2, 3, 300, 301},
        }
        sampling_set, qualified, ratios = selector.select_sampling_set(miner_taskids, 10)
        # Common task_ids (1, 2, 3) should be prioritized
        assert 1 in sampling_set
        assert 2 in sampling_set
        assert 3 in sampling_set
    
    def test_min_dataset_coverage_constraint(self):
        """Test that sampling set covers at least 50% of dataset."""
        selector = SamplingSetSelector()
        miner_taskids = {
            "hk1": set(range(1, 1001)),  # 1000 task_ids
        }
        sampling_set, qualified, ratios = selector.select_sampling_set(miner_taskids, 2000)
        # Should select at least 1000 task_ids (50% of 2000)
        assert len(sampling_set) >= 1000


class TestSampleAligner:
    """Tests for SampleAligner class."""
    
    def test_align_empty_inputs(self):
        """Test alignment with empty inputs."""
        aligner = SampleAligner()
        aligned, count = aligner.align_samples({}, set(), set())
        assert aligned == {}
        assert count == 0
    
    def test_align_single_miner(self):
        """Test alignment with single miner."""
        aligner = SampleAligner()
        miner_results = {
            "hk1": [
                MockResult("hk1", "env1", 0.8, task_id=1, timestamp=100),
                MockResult("hk1", "env1", 0.9, task_id=2, timestamp=101),
            ]
        }
        sampling_set = {1, 2}
        qualified_miners = {"hk1"}
        aligned, count = aligner.align_samples(miner_results, sampling_set, qualified_miners)
        assert "hk1" in aligned
        assert len(aligned["hk1"]) == 2
        assert count == 2
    
    def test_align_multiple_miners_equal_samples(self):
        """Test alignment when miners have equal samples."""
        aligner = SampleAligner()
        miner_results = {
            "hk1": [
                MockResult("hk1", "env1", 0.8, task_id=1, timestamp=100),
                MockResult("hk1", "env1", 0.9, task_id=2, timestamp=101),
            ],
            "hk2": [
                MockResult("hk2", "env1", 0.7, task_id=1, timestamp=200),
                MockResult("hk2", "env1", 0.6, task_id=2, timestamp=201),
            ]
        }
        sampling_set = {1, 2}
        qualified_miners = {"hk1", "hk2"}
        aligned, count = aligner.align_samples(miner_results, sampling_set, qualified_miners)
        assert len(aligned) == 2
        assert len(aligned["hk1"]) == 2
        assert len(aligned["hk2"]) == 2
        assert count == 2
    
    def test_align_multiple_miners_unequal_samples(self):
        """Test alignment truncates to minimum."""
        aligner = SampleAligner()
        miner_results = {
            "hk1": [
                MockResult("hk1", "env1", 0.8, task_id=1, timestamp=100),
                MockResult("hk1", "env1", 0.9, task_id=2, timestamp=101),
                MockResult("hk1", "env1", 0.7, task_id=3, timestamp=102),
            ],
            "hk2": [
                MockResult("hk2", "env1", 0.7, task_id=1, timestamp=200),
                MockResult("hk2", "env1", 0.6, task_id=2, timestamp=201),
            ]
        }
        sampling_set = {1, 2, 3}
        qualified_miners = {"hk1", "hk2"}
        aligned, count = aligner.align_samples(miner_results, sampling_set, qualified_miners)
        assert len(aligned) == 2
        # Should truncate to minimum (2)
        assert len(aligned["hk1"]) == 2
        assert len(aligned["hk2"]) == 2
        assert count == 2
    
    def test_align_filters_by_sampling_set(self):
        """Test that alignment only includes results in sampling set."""
        aligner = SampleAligner()
        miner_results = {
            "hk1": [
                MockResult("hk1", "env1", 0.8, task_id=1, timestamp=100),
                MockResult("hk1", "env1", 0.9, task_id=2, timestamp=101),
                MockResult("hk1", "env1", 0.7, task_id=999, timestamp=102),  # Not in sampling set
            ]
        }
        sampling_set = {1, 2}
        qualified_miners = {"hk1"}
        aligned, count = aligner.align_samples(miner_results, sampling_set, qualified_miners)
        assert len(aligned["hk1"]) == 2
        # Task_id 999 should be excluded
        task_ids = [r.task_id for r in aligned["hk1"]]
        assert 999 not in task_ids
    
    def test_align_keeps_newest_results(self):
        """Test that alignment keeps newest results when truncating."""
        aligner = SampleAligner()
        miner_results = {
            "hk1": [
                MockResult("hk1", "env1", 0.8, task_id=1, timestamp=100),
                MockResult("hk1", "env1", 0.9, task_id=2, timestamp=300),  # Newest
                MockResult("hk1", "env1", 0.7, task_id=3, timestamp=200),
            ],
            "hk2": [
                MockResult("hk2", "env1", 0.6, task_id=1, timestamp=150),
            ]
        }
        sampling_set = {1, 2, 3}
        qualified_miners = {"hk1", "hk2"}
        aligned, count = aligner.align_samples(miner_results, sampling_set, qualified_miners)
        # Should keep only 1 sample (minimum)
        assert count == 1
        # hk1 should keep the newest (timestamp=300)
        assert aligned["hk1"][0].timestamp == 300


class TestTaskIdAlignmentOrchestrator:
    """Tests for TaskIdAlignmentOrchestrator class."""
    
    def test_process_environment_empty(self):
        """Test processing empty environment results."""
        orchestrator = TaskIdAlignmentOrchestrator()
        aligned, qualified, ratios, count = orchestrator.process_environment_results({}, 1000)
        assert aligned == {}
        assert qualified == set()
        assert ratios == {}
        assert count == 0
    
    def test_process_environment_single_miner(self):
        """Test processing single miner results."""
        orchestrator = TaskIdAlignmentOrchestrator()
        env_results = {
            "hk1": [
                MockResult("hk1", "env1", 0.8, task_id=i, timestamp=100 + i)
                for i in range(1, 501)  # 500 samples
            ]
        }
        aligned, qualified, ratios, count = orchestrator.process_environment_results(env_results, 1000)
        assert "hk1" in qualified
        assert "hk1" in aligned
        assert count > 0
    
    def test_process_environment_multiple_miners_with_dedup(self):
        """Test full pipeline with deduplication."""
        orchestrator = TaskIdAlignmentOrchestrator()
        env_results = {
            "hk1": [
                MockResult("hk1", "env1", 0.8, task_id=1, timestamp=100),
                MockResult("hk1", "env1", 0.9, task_id=1, timestamp=200),  # Duplicate
                MockResult("hk1", "env1", 0.7, task_id=2, timestamp=101),
            ] + [
                MockResult("hk1", "env1", 0.85, task_id=i, timestamp=100 + i)
                for i in range(3, 503)  # More samples
            ],
            "hk2": [
                MockResult("hk2", "env1", 0.75, task_id=i, timestamp=100 + i)
                for i in range(1, 501)  # 500 samples
            ]
        }
        aligned, qualified, ratios, count = orchestrator.process_environment_results(env_results, 1000)
        # Both miners should be qualified
        assert len(qualified) == 2
        # Sample counts should be aligned
        if aligned:
            counts = [len(results) for results in aligned.values()]
            assert len(set(counts)) == 1  # All equal
    
    def test_process_all_environments(self):
        """Test processing multiple environments."""
        orchestrator = TaskIdAlignmentOrchestrator()
        results_by_env = {
            "env1": {
                "hk1": [
                    MockResult("hk1", "env1", 0.8, task_id=i, timestamp=100 + i)
                    for i in range(1, 501)
                ],
                "hk2": [
                    MockResult("hk2", "env1", 0.7, task_id=i, timestamp=100 + i)
                    for i in range(1, 501)
                ]
            },
            "env2": {
                "hk1": [
                    MockResult("hk1", "env2", 0.9, task_id=i, timestamp=100 + i)
                    for i in range(1, 301)
                ],
                "hk2": [
                    MockResult("hk2", "env2", 0.85, task_id=i, timestamp=100 + i)
                    for i in range(1, 301)
                ]
            }
        }
        dataset_lengths = {"env1": 1000, "env2": 600}
        
        aligned_results, qualified_miners, overlap_ratios, aligned_counts = orchestrator.process_all_environments(
            results_by_env, dataset_lengths
        )
        
        # Check that both environments are processed
        assert "env1" in aligned_results
        assert "env2" in aligned_results
        assert "env1" in qualified_miners
        assert "env2" in qualified_miners
        assert "env1" in aligned_counts
        assert "env2" in aligned_counts


class TestIntegration:
    """Integration tests for the complete pipeline."""
    
    def test_realistic_scenario_high_overlap(self):
        """Test realistic scenario with high task_id overlap."""
        orchestrator = TaskIdAlignmentOrchestrator()
        
        # Simulate 3 miners with high overlap
        common_taskids = list(range(1, 501))  # 500 common task_ids
        
        env_results = {}
        for miner_idx in range(3):
            hk = f"hk{miner_idx + 1}"
            results = []
            for task_id in common_taskids:
                results.append(MockResult(
                    hk, "env1", 
                    0.7 + miner_idx * 0.05,  # Different scores
                    task_id=task_id,
                    timestamp=100 + task_id + miner_idx * 1000
                ))
            env_results[hk] = results
        
        aligned, qualified, ratios, count = orchestrator.process_environment_results(env_results, 1000)
        
        # All miners should be qualified
        assert len(qualified) == 3
        # All should have high overlap
        for hk in qualified:
            assert ratios[hk] >= 0.8
        # Sample counts should be equal
        assert count > 0
        for hk in aligned:
            assert len(aligned[hk]) == count
    
    def test_realistic_scenario_low_overlap(self):
        """Test realistic scenario with low task_id overlap."""
        orchestrator = TaskIdAlignmentOrchestrator()
        
        # Simulate 3 miners with low overlap
        env_results = {
            "hk1": [
                MockResult("hk1", "env1", 0.8, task_id=i, timestamp=100 + i)
                for i in range(1, 301)  # 1-300
            ],
            "hk2": [
                MockResult("hk2", "env1", 0.75, task_id=i, timestamp=100 + i)
                for i in range(200, 500)  # 200-499 (some overlap with hk1)
            ],
            "hk3": [
                MockResult("hk3", "env1", 0.7, task_id=i, timestamp=100 + i)
                for i in range(400, 700)  # 400-699 (some overlap with hk2)
            ]
        }
        
        aligned, qualified, ratios, count = orchestrator.process_environment_results(env_results, 1000)
        
        # Should still qualify at least 50% of miners
        assert len(qualified) >= 2
    
    def test_edge_case_very_small_dataset(self):
        """Test with very small dataset (< 400)."""
        orchestrator = TaskIdAlignmentOrchestrator()
        
        env_results = {
            "hk1": [
                MockResult("hk1", "env1", 0.8, task_id=i, timestamp=100 + i)
                for i in range(1, 201)  # 200 samples
            ],
            "hk2": [
                MockResult("hk2", "env1", 0.75, task_id=i, timestamp=100 + i)
                for i in range(1, 201)  # 200 samples
            ]
        }
        
        aligned, qualified, ratios, count = orchestrator.process_environment_results(env_results, 200)
        
        # Should handle small dataset correctly
        assert len(qualified) >= 1
    
    def test_edge_case_duplicate_heavy(self):
        """Test with many duplicate task_ids."""
        orchestrator = TaskIdAlignmentOrchestrator()
        
        # Each task_id appears multiple times
        results = []
        for task_id in range(1, 501):
            for repeat in range(5):  # 5 repeats per task_id
                results.append(MockResult(
                    "hk1", "env1", 
                    0.8 + repeat * 0.01,
                    task_id=task_id,
                    timestamp=100 + task_id * 10 + repeat
                ))
        
        env_results = {"hk1": results}
        
        aligned, qualified, ratios, count = orchestrator.process_environment_results(env_results, 1000)
        
        # Should deduplicate to at most 1 per task_id (large dataset)
        assert "hk1" in aligned
        # Check no duplicate task_ids
        task_ids = [r.task_id for r in aligned["hk1"]]
        assert len(task_ids) == len(set(task_ids))


if __name__ == "__main__":
    pytest.main([__file__, "-v"])