"""
Tests for HOLD - Inference-Level Intervention System

Comprehensive tests for:
- HoldPoint creation and merkle chaining
- HoldResolution lifecycle
- Hold singleton and yield_point flow
- Override vs accept paths
- Non-blocking mode
- Listener callbacks
- Serialization
"""

import time
import pytest
import numpy as np
from threading import Thread
import sys

# Import HOLD components
from cascade.hold import (
    Hold,
    HoldPoint,
    HoldResolution,
    HoldState,
    HoldAwareMixin,
)


class TestHoldPoint:
    """Test HoldPoint creation and properties."""
    
    def test_hold_point_creation(self):
        """Test HoldPoint can be created with all fields."""
        probs = np.array([0.1, 0.6, 0.2, 0.1])
        point = HoldPoint(
            action_probs=probs,
            value=0.85,
            observation={"state": [1, 2, 3]},
            brain_id="test_brain",
        )
        
        assert point.brain_id == "test_brain"
        assert point.value == 0.85
        assert point.state == HoldState.PENDING
        assert len(point.id) == 16  # Hash prefix length
    
    def test_hold_point_ai_choice(self):
        """Test HoldPoint computes AI's preferred action."""
        # Highest prob at index 2
        probs = np.array([0.1, 0.2, 0.5, 0.2])
        point = HoldPoint(
            action_probs=probs,
            value=0.0,
            observation={},
            brain_id="test",
        )
        
        assert point.ai_choice == 2
        assert point.ai_confidence == 0.5
    
    def test_hold_point_merkle_chain(self):
        """Test merkle roots chain properly."""
        probs = np.array([0.5, 0.5])
        
        # First point - no parent
        point1 = HoldPoint(
            action_probs=probs,
            value=0.0,
            observation={},
            brain_id="test",
        )
        
        assert point1.merkle_root is not None
        assert point1.parent_merkle is None
        
        # Second point - chains to first
        point2 = HoldPoint(
            action_probs=probs,
            value=0.0,
            observation={},
            brain_id="test",
            parent_merkle=point1.merkle_root,
        )
        
        assert point2.parent_merkle == point1.merkle_root
        assert point2.merkle_root != point1.merkle_root  # Different hashes
    
    def test_hold_point_unique_ids(self):
        """Test each HoldPoint gets a unique ID (based on timestamp)."""
        probs = np.array([0.5, 0.5])
        ids = set()
        
        for i in range(10):
            # Add small delay to ensure unique timestamps
            time.sleep(0.01)
            point = HoldPoint(action_probs=probs, value=0.0, observation={"i": i}, brain_id="test")
            ids.add(point.id)
        
        # Should have unique IDs (at least some variety)
        # Note: With very fast execution, some might collide
        assert len(ids) >= 5  # At least half should be unique
    
    def test_hold_point_to_dict(self):
        """Test HoldPoint serializes correctly."""
        probs = np.array([0.25, 0.25, 0.25, 0.25])
        point = HoldPoint(
            action_probs=probs,
            value=1.5,
            observation={"key": "value"},
            brain_id="brain_001",
        )
        
        d = point.to_dict()
        
        assert d['brain_id'] == "brain_001"
        assert d['ai_choice'] == 0  # First with equal probs
        assert d['ai_confidence'] == 0.25
        assert d['value'] == 1.5
        assert d['state'] == 'pending'
        assert 'action_probs' in d
        assert len(d['action_probs']) == 4


class TestHoldResolution:
    """Test HoldResolution lifecycle."""
    
    def test_resolution_accept(self):
        """Test resolution when accepting AI choice."""
        probs = np.array([0.1, 0.7, 0.2])  # AI chooses 1
        point = HoldPoint(
            action_probs=probs,
            value=0.0,
            observation={},
            brain_id="test",
        )
        
        resolution = HoldResolution(
            hold_point=point,
            action=1,  # Same as AI
            was_override=False,
        )
        
        assert resolution.action == 1
        assert resolution.was_override is False
        assert resolution.override_source is None
        assert resolution.merkle_root is not None
    
    def test_resolution_override(self):
        """Test resolution when overriding AI choice."""
        probs = np.array([0.1, 0.7, 0.2])  # AI chooses 1
        point = HoldPoint(
            action_probs=probs,
            value=0.0,
            observation={},
            brain_id="test",
        )
        
        resolution = HoldResolution(
            hold_point=point,
            action=2,  # Override to action 2
            was_override=True,
            override_source="human",
        )
        
        assert resolution.action == 2
        assert resolution.was_override is True
        assert resolution.override_source == "human"
    
    def test_resolution_merkle_links(self):
        """Test resolution merkle links to hold point."""
        probs = np.array([0.5, 0.5])
        point = HoldPoint(
            action_probs=probs,
            value=0.0,
            observation={},
            brain_id="test",
        )
        
        resolution = HoldResolution(
            hold_point=point,
            action=0,
            was_override=False,
        )
        
        # Resolution merkle should include hold point merkle
        d = resolution.to_dict()
        assert d['hold_merkle'] == point.merkle_root
    
    def test_resolution_to_dict(self):
        """Test HoldResolution serializes correctly."""
        probs = np.array([0.3, 0.7])
        point = HoldPoint(
            action_probs=probs,
            value=0.5,
            observation={},
            brain_id="test",
        )
        
        resolution = HoldResolution(
            hold_point=point,
            action=0,
            was_override=True,
            override_source="policy",
            hold_duration=2.5,
        )
        
        d = resolution.to_dict()
        
        assert d['action'] == 0
        assert d['ai_choice'] == 1
        assert d['was_override'] is True
        assert d['override_source'] == "policy"
        assert d['hold_duration'] == 2.5


class TestHoldSingleton:
    """Test Hold singleton behavior."""
    
    def test_hold_is_singleton(self):
        """Test Hold returns same instance."""
        h1 = Hold.get()
        h2 = Hold.get()
        h3 = Hold()
        
        assert h1 is h2
        assert h2 is h3
    
    def test_hold_auto_accept_mode(self):
        """Test non-blocking auto-accept mode."""
        hold = Hold.get()
        hold.auto_accept = True  # Don't block
        
        probs = np.array([0.1, 0.1, 0.8])
        resolution = hold.yield_point(
            action_probs=probs,
            value=0.0,
            observation={},
            brain_id="test",
        )
        
        # Should return immediately with AI choice
        assert resolution.action == 2
        assert resolution.was_override is False
        
        # Reset
        hold.auto_accept = False
    
    def test_hold_non_blocking_param(self):
        """Test blocking=False returns immediately."""
        hold = Hold.get()
        hold.auto_accept = False  # Even with this False
        
        probs = np.array([0.9, 0.1])
        resolution = hold.yield_point(
            action_probs=probs,
            value=0.0,
            observation={},
            brain_id="test",
            blocking=False,  # This overrides
        )
        
        assert resolution.action == 0
        assert resolution.was_override is False
    
    def test_hold_listener_registration(self):
        """Test listeners are called on hold."""
        hold = Hold.get()
        hold.auto_accept = True
        
        received = []
        
        def listener(point):
            received.append(point)
        
        hold.register_listener(listener)
        
        probs = np.array([0.5, 0.5])
        hold.yield_point(
            action_probs=probs,
            value=0.0,
            observation={"test": "data"},
            brain_id="test",
            blocking=False,
        )
        
        # In auto_accept mode, listeners might not be called
        # but we can test registration works
        hold.unregister_listener(listener)
        assert listener not in hold._listeners
    
    def test_hold_counts_increment(self):
        """Test hold/override counts track correctly."""
        hold = Hold.get()
        initial_count = hold._hold_count
        hold.auto_accept = True
        
        probs = np.array([0.5, 0.5])
        hold.yield_point(
            action_probs=probs,
            value=0.0,
            observation={},
            brain_id="test",
        )
        
        assert hold._hold_count == initial_count + 1


class TestHoldAwareMixin:
    """Test HoldAwareMixin integration."""
    
    def test_mixin_provides_hold_methods(self):
        """Test mixin adds hold capability."""
        class TestBrain(HoldAwareMixin):
            def __init__(self):
                super().__init__()
                self.id = "test_brain"
            
            def forward(self, inputs):
                return {"action_probs": [0.5, 0.5], "value": 0.0}
        
        brain = TestBrain()
        
        # Should have hold-related methods
        assert hasattr(brain, 'forward_with_hold')
        assert hasattr(brain, 'enable_hold')
        assert hasattr(brain, 'disable_hold')
    
    def test_mixin_forward_with_hold(self):
        """Test mixin forward_with_hold calls through to Hold."""
        class TestBrain(HoldAwareMixin):
            def __init__(self):
                super().__init__()
                self.id = "test_mixin_brain"
            
            def forward(self, inputs):
                return {"action_probs": [0.1, 0.9], "value": 1.0}
        
        brain = TestBrain()
        hold = Hold.get()
        hold.auto_accept = True
        
        output = brain.forward_with_hold({"step": 1}, blocking=False)
        
        assert output['action'] == 1
        assert 'hold_resolution' in output


class TestHoldIntegration:
    """Integration tests for HOLD workflow."""
    
    def test_full_workflow_accept(self):
        """Test complete hold -> accept workflow."""
        hold = Hold.get()
        hold.auto_accept = True
        
        # Simulate a sequence of decisions
        actions = []
        probs_sequence = [
            np.array([0.8, 0.1, 0.1]),
            np.array([0.1, 0.7, 0.2]),
            np.array([0.2, 0.2, 0.6]),
        ]
        
        for i, probs in enumerate(probs_sequence):
            resolution = hold.yield_point(
                action_probs=probs,
                value=float(i),
                observation={"step": i},
                brain_id="sequence_test",
            )
            actions.append(resolution.action)
        
        assert actions == [0, 1, 2]  # AI choices
    
    def test_action_names_in_hold(self):
        """Test hold points can include action names."""
        probs = np.array([0.25, 0.25, 0.25, 0.25])
        action_names = ["up", "down", "left", "right"]
        
        point = HoldPoint(
            action_probs=probs,
            value=0.0,
            observation={"action_names": action_names},
            brain_id="test",
        )
        
        # Action names stored in observation
        assert point.observation["action_names"] == action_names
        assert point.ai_choice in range(4)
    
    def test_imagined_futures(self):
        """Test hold points can include trajectory predictions."""
        probs = np.array([0.5, 0.5])
        futures = [
            {"action": 0, "predicted_reward": 1.0, "states": [[0,0], [0,1]]},
            {"action": 1, "predicted_reward": 0.5, "states": [[0,0], [1,0]]},
        ]
        
        point = HoldPoint(
            action_probs=probs,
            value=0.0,
            observation={},
            brain_id="test",
            imagined_futures=futures,
        )
        
        assert point.imagined_futures is not None
        assert len(point.imagined_futures) == 2
        assert point.imagined_futures[0]["predicted_reward"] == 1.0


class TestHoldEdgeCases:
    """Test edge cases and error handling."""
    
    def test_empty_probs(self):
        """Test with single action."""
        probs = np.array([1.0])
        point = HoldPoint(
            action_probs=probs,
            value=0.0,
            observation={},
            brain_id="test",
        )
        
        assert point.ai_choice == 0
        assert point.ai_confidence == 1.0
    
    def test_uniform_probs(self):
        """Test with uniform distribution."""
        probs = np.array([0.2, 0.2, 0.2, 0.2, 0.2])
        point = HoldPoint(
            action_probs=probs,
            value=0.0,
            observation={},
            brain_id="test",
        )
        
        assert point.ai_choice == 0  # argmax returns first
        assert abs(point.ai_confidence - 0.2) < 0.01
    
    def test_large_action_space(self):
        """Test with many actions."""
        n_actions = 100
        probs = np.zeros(n_actions)
        probs[73] = 1.0  # Set one to max
        
        point = HoldPoint(
            action_probs=probs,
            value=0.0,
            observation={},
            brain_id="test",
        )
        
        assert point.ai_choice == 73
        assert point.ai_confidence == 1.0


class TestHoldStatistics:
    """Test hold statistics tracking."""
    
    def test_stats_access(self):
        """Test accessing hold statistics."""
        hold = Hold.get()
        
        # These should be accessible
        assert hasattr(hold, '_hold_count')
        assert hasattr(hold, '_override_count')
        assert isinstance(hold._hold_count, int)
        assert isinstance(hold._override_count, int)


# Run with: pytest tests/test_hold.py -v
if __name__ == "__main__":
    pytest.main([__file__, "-v"])
