#!/usr/bin/env python3
"""
tests/test_policy_harness.py

WO-ECHO-POLICY-GOLDEN-01 - Unit tests for policy harness

Updated for WO-ECHO-POLICY-GOLDEN-01:
- Reduced action space (4 knobs, 24 combinations)
- Fixed env var names (HELIX_USE_CPP_KERNEL, HELIX_USE_FUSED_MATMUL)
- Removed tests for deprecated knobs (sample_rate, cache_size, block_rows)

Tests:
- PolicyState: creation, serialization, tensor conversion
- PolicyAction: creation, value mapping, random generation, index conversion
- PolicyMLP: forward pass, parameter count, save/load (4 heads)
- PolicyHarness: state extraction, action selection, apply_action
- compute_reward: correctness gate, throughput improvement, penalties
- ExperienceBuffer: add, sample, ring buffer behavior
- Training: REINFORCE convergence on synthetic data
"""

import json
import os
import sys
import tempfile
from pathlib import Path

import numpy as np
import pytest

# Add project root to path
PROJECT_ROOT = Path(__file__).parent.parent
sys.path.insert(0, str(PROJECT_ROOT))


# =============================================================================
# PolicyState Tests
# =============================================================================

class TestPolicyState:
    def test_creation_defaults(self):
        from helix_cdc.runtime.policy_harness import PolicyState

        state = PolicyState()
        assert state.tok_per_sec == 0.0
        assert state.cache_hit_rate == 0.0
        assert state.regen_hit_rate == 1.0  # Default to 1.0 (no regen)

    def test_creation_with_values(self):
        from helix_cdc.runtime.policy_harness import PolicyState

        state = PolicyState(
            tok_per_sec=42.5,
            mean_step_ms=23.4,
            p95_step_ms=45.6,
            steps_in_window=100,
            cache_hit_rate=0.85,
            cache_misses=15,
            cache_hits=85,
            regen_hit_rate=0.95,
            regen_time_pct=5.0,
            regen_calls=5,
            rss_delta_mb=10.0,
            rss_mb=500.0,
            se_mean=0.3,
            se_std=0.05,
        )
        assert state.tok_per_sec == 42.5
        assert state.steps_in_window == 100

    def test_to_array(self):
        from helix_cdc.runtime.policy_harness import PolicyState

        state = PolicyState(tok_per_sec=50.0, cache_hits=100)
        arr = state.to_array()
        assert arr.shape == (14,)
        assert arr.dtype == np.float32
        assert arr[0] == 50.0  # tok_per_sec is first

    def test_to_tensor(self):
        from helix_cdc.runtime.policy_harness import PolicyState

        state = PolicyState(tok_per_sec=50.0)
        tensor = state.to_tensor()
        assert tensor.shape == (14,)
        assert tensor[0].item() == 50.0

    def test_from_dict(self):
        from helix_cdc.runtime.policy_harness import PolicyState

        d = {
            "tok_per_sec": 30.0,
            "mean_step_ms": 20.0,
            "cache_hit_rate": 0.9,
        }
        state = PolicyState.from_dict(d)
        assert state.tok_per_sec == 30.0
        assert state.cache_hit_rate == 0.9

    def test_to_dict_roundtrip(self):
        from helix_cdc.runtime.policy_harness import PolicyState

        state = PolicyState(
            tok_per_sec=42.5,
            se_mean=0.3,
        )
        d = state.to_dict()
        state2 = PolicyState.from_dict(d)
        assert state2.tok_per_sec == state.tok_per_sec
        assert state2.se_mean == state.se_mean


# =============================================================================
# PolicyAction Tests
# =============================================================================

class TestPolicyAction:
    def test_default(self):
        from helix_cdc.runtime.policy_harness import PolicyAction

        action = PolicyAction.default()
        assert action.verify_policy == 2  # trust_cached
        assert action.cpp_kernel == 0
        assert action.fused_matmul == 0
        assert action.fibpi_mode == 0

    def test_random(self):
        from helix_cdc.runtime.policy_harness import PolicyAction

        rng = np.random.default_rng(42)
        action = PolicyAction.random(rng)

        # WO-ECHO-POLICY-GOLDEN-01: Check bounds for 4 knobs
        assert 0 <= action.verify_policy <= 2
        assert 0 <= action.cpp_kernel <= 1
        assert 0 <= action.fused_matmul <= 1
        assert 0 <= action.fibpi_mode <= 1

    def test_to_dict_actual_values(self):
        from helix_cdc.runtime.policy_harness import PolicyAction

        # WO-ECHO-POLICY-GOLDEN-01: Only 4 knobs
        action = PolicyAction(
            verify_policy=0,  # always
            cpp_kernel=1,     # on
            fused_matmul=1,   # on
            fibpi_mode=1,     # se_bucket
        )
        d = action.to_dict()

        assert d["verify_policy"] == "always"
        assert d["cpp_kernel"] is True
        assert d["fused_matmul"] is True
        assert d["fibpi_mode"] == "se_bucket"
        # Removed: sample_rate, cache_size, block_rows

    def test_from_dict_actual_values(self):
        from helix_cdc.runtime.policy_harness import PolicyAction

        # WO-ECHO-POLICY-GOLDEN-01: Only 4 knobs
        d = {
            "verify_policy": "sampled",
            "cpp_kernel": True,
            "fused_matmul": False,
            "fibpi_mode": "se_bucket",
        }
        action = PolicyAction.from_dict(d)

        assert action.verify_policy == 1  # sampled
        assert action.cpp_kernel == 1
        assert action.fused_matmul == 0
        assert action.fibpi_mode == 1

    def test_action_index_roundtrip(self):
        """WO-ECHO-POLICY-GOLDEN-01: Test index conversion."""
        from helix_cdc.runtime.policy_harness import PolicyAction, TOTAL_ACTION_COMBINATIONS

        assert TOTAL_ACTION_COMBINATIONS == 24

        # Test all indices
        for i in range(TOTAL_ACTION_COMBINATIONS):
            action = PolicyAction.from_index(i)
            assert action.to_index() == i

    def test_all_actions(self):
        """WO-ECHO-POLICY-GOLDEN-01: Test all_actions enumeration."""
        from helix_cdc.runtime.policy_harness import PolicyAction, TOTAL_ACTION_COMBINATIONS

        all_actions = PolicyAction.all_actions()
        assert len(all_actions) == TOTAL_ACTION_COMBINATIONS
        assert len(all_actions) == 24

        # Check uniqueness
        indices = [a.to_index() for a in all_actions]
        assert len(set(indices)) == 24

    def test_verify_policy_str(self):
        """WO-ECHO-POLICY-GOLDEN-01: Test verify_policy_str property."""
        from helix_cdc.runtime.policy_harness import PolicyAction

        action = PolicyAction(verify_policy=0)
        assert action.verify_policy_str == "always"

        action = PolicyAction(verify_policy=1)
        assert action.verify_policy_str == "sampled"

        action = PolicyAction(verify_policy=2)
        assert action.verify_policy_str == "trust_cached"


# =============================================================================
# PolicyMLP Tests
# =============================================================================

class TestPolicyMLP:
    def test_creation(self):
        from helix_cdc.runtime.policy_harness import PolicyMLP

        model = PolicyMLP(state_dim=14, hidden_dim=32)
        assert model.state_dim == 14
        assert model.hidden_dim == 32

    def test_parameter_count(self):
        from helix_cdc.runtime.policy_harness import PolicyMLP

        model = PolicyMLP(state_dim=14, hidden_dim=32)
        n_params = model.count_parameters()

        # WO-ECHO-POLICY-GOLDEN-01: Reduced from 7 to 4 heads
        # Expected:
        # shared: 14*32 + 32 + 32*32 + 32 = 448 + 32 + 1024 + 32 = 1536
        # heads: (32*3+3) + (32*2+2) + (32*2+2) + (32*2+2) = 99 + 66 + 66 + 66 = 297
        # Total ~1833
        assert n_params > 1500
        assert n_params < 2500

    def test_forward(self):
        from helix_cdc.runtime.policy_harness import PolicyMLP, PolicyState
        import torch

        model = PolicyMLP(state_dim=14, hidden_dim=32)
        state = PolicyState(tok_per_sec=50.0, cache_hit_rate=0.8)

        probs = model(state.to_tensor())

        # WO-ECHO-POLICY-GOLDEN-01: Only 4 heads
        assert "verify_policy" in probs
        assert "cpp_kernel" in probs
        assert "fused_matmul" in probs
        assert "fibpi_mode" in probs

        # Removed heads should not exist
        assert "sample_rate" not in probs
        assert "cache_size" not in probs
        assert "block_rows" not in probs

        # Check probabilities sum to 1
        assert abs(probs["verify_policy"].sum().item() - 1.0) < 1e-5
        assert abs(probs["cpp_kernel"].sum().item() - 1.0) < 1e-5

    def test_forward_batch(self):
        from helix_cdc.runtime.policy_harness import PolicyMLP
        import torch

        model = PolicyMLP(state_dim=14, hidden_dim=32)
        batch = torch.randn(8, 14)  # Batch of 8 states

        probs = model(batch)

        # WO-ECHO-POLICY-GOLDEN-01: Only 4 heads
        assert probs["verify_policy"].shape == (8, 3)
        assert probs["cpp_kernel"].shape == (8, 2)
        assert probs["fused_matmul"].shape == (8, 2)
        assert probs["fibpi_mode"].shape == (8, 2)

    def test_state_dict_roundtrip(self):
        from helix_cdc.runtime.policy_harness import PolicyMLP
        import torch

        model1 = PolicyMLP(state_dim=14, hidden_dim=32)
        state_dict = model1.state_dict()

        model2 = PolicyMLP.from_state_dict(state_dict)

        # Check outputs match
        x = torch.randn(14)
        probs1 = model1(x)
        probs2 = model2(x)

        for key in probs1:
            assert torch.allclose(probs1[key], probs2[key])


# =============================================================================
# compute_reward Tests
# =============================================================================

class TestComputeReward:
    def test_correctness_failure(self):
        from helix_cdc.runtime.policy_harness import (
            PolicyState, PolicyAction, compute_reward
        )

        prev = PolicyState()
        next_ = PolicyState(tok_per_sec=100.0)  # Great improvement
        action = PolicyAction.default()

        reward = compute_reward(prev, action, next_, correctness_ok=False)
        assert reward == -10.0  # Huge penalty

    def test_throughput_improvement(self):
        from helix_cdc.runtime.policy_harness import (
            PolicyState, PolicyAction, compute_reward
        )

        prev = PolicyState(tok_per_sec=30.0)
        next_ = PolicyState(tok_per_sec=50.0)  # +20 tok/s
        action = PolicyAction.default()

        reward = compute_reward(prev, action, next_, correctness_ok=True)
        assert reward > 0  # Positive reward for improvement

    def test_latency_penalty(self):
        from helix_cdc.runtime.policy_harness import (
            PolicyState, PolicyAction, compute_reward
        )

        prev = PolicyState(p95_step_ms=20.0)
        next_ = PolicyState(p95_step_ms=40.0)  # 2x increase (>1.5x threshold)
        action = PolicyAction.default()

        reward = compute_reward(prev, action, next_, correctness_ok=True)
        # Should have latency penalty
        assert reward <= 0.3  # Base reward minus penalty

    def test_cache_bonus(self):
        from helix_cdc.runtime.policy_harness import (
            PolicyState, PolicyAction, compute_reward
        )

        prev = PolicyState(tok_per_sec=30.0)
        next_ = PolicyState(tok_per_sec=30.0, cache_hit_rate=0.95)  # >0.9
        action = PolicyAction.default()

        reward = compute_reward(prev, action, next_, correctness_ok=True)
        assert reward > 0  # Cache bonus


# =============================================================================
# PolicyHarness Tests
# =============================================================================

class TestPolicyHarness:
    def test_default(self):
        from helix_cdc.runtime.policy_harness import PolicyHarness

        harness = PolicyHarness.default()
        assert harness.model is not None
        assert harness.normalizer is not None

    def test_select_action_greedy(self):
        from helix_cdc.runtime.policy_harness import PolicyHarness, PolicyState

        harness = PolicyHarness.default()
        state = PolicyState(tok_per_sec=50.0)

        action = harness.select_action(state, explore=False)

        assert action is not None
        assert 0 <= action.verify_policy <= 2

    def test_select_action_explore(self):
        from helix_cdc.runtime.policy_harness import PolicyHarness, PolicyState

        harness = PolicyHarness.default()
        state = PolicyState(tok_per_sec=50.0)

        # With epsilon=1.0, should always explore
        actions = [harness.select_action(state, explore=True, epsilon=1.0)
                   for _ in range(10)]

        # Check we get some variety
        verify_policies = set(a.verify_policy for a in actions)
        assert len(verify_policies) > 1 or len(actions) < 3

    def test_apply_action_context(self):
        from helix_cdc.runtime.policy_harness import PolicyHarness, PolicyAction

        harness = PolicyHarness.default()
        action = PolicyAction(verify_policy=0, cpp_kernel=1, fused_matmul=1, fibpi_mode=1)

        # WO-ECHO-POLICY-GOLDEN-01: Fixed env var names
        # Save original env
        orig_cpp = os.environ.get("HELIX_USE_CPP_KERNEL")
        orig_fused = os.environ.get("HELIX_USE_FUSED_MATMUL")
        orig_fibpi = os.environ.get("HELIX_FIBPI_MODE")

        with harness.apply_action(action) as applied:
            # Correct env var names
            assert os.environ.get("HELIX_USE_CPP_KERNEL") == "1"
            assert os.environ.get("HELIX_USE_FUSED_MATMUL") == "1"
            assert os.environ.get("HELIX_FIBPI_MODE") == "se_bucket"

            # verify_policy NOT set via env var (passed as parameter)
            assert applied["verify_policy"] == "always"

        # Should be restored
        assert os.environ.get("HELIX_USE_CPP_KERNEL") == orig_cpp
        assert os.environ.get("HELIX_USE_FUSED_MATMUL") == orig_fused
        assert os.environ.get("HELIX_FIBPI_MODE") == orig_fibpi

    def test_checkpoint_roundtrip(self):
        from helix_cdc.runtime.policy_harness import PolicyHarness, PolicyState

        harness = PolicyHarness.default()

        # Update normalizer
        for _ in range(100):
            state = PolicyState(tok_per_sec=np.random.uniform(20, 60))
            harness.normalizer.update(state)

        with tempfile.NamedTemporaryFile(suffix=".pt", delete=False) as f:
            path = f.name

        try:
            harness.save_checkpoint(path)
            harness2 = PolicyHarness.from_checkpoint(path)

            # Check normalizer preserved
            assert harness2.normalizer.count == harness.normalizer.count

            # Check model produces same output
            state = PolicyState(tok_per_sec=40.0)
            action1 = harness.select_action(state)
            action2 = harness2.select_action(state)
            assert action1.verify_policy == action2.verify_policy
        finally:
            os.unlink(path)


# =============================================================================
# ExperienceBuffer Tests
# =============================================================================

class TestExperienceBuffer:
    def test_add_and_get(self):
        from helix_cdc.runtime.policy_trainer import (
            ExperienceBuffer, Experience
        )
        from helix_cdc.runtime.policy_harness import PolicyState, PolicyAction

        buffer = ExperienceBuffer(max_size=100)

        exp = Experience(
            state=PolicyState(tok_per_sec=30.0),
            action=PolicyAction.default(),
            reward=0.5,
            next_state=PolicyState(tok_per_sec=35.0),
        )
        buffer.add(exp)

        assert len(buffer) == 1
        all_exp = buffer.get_all()
        assert len(all_exp) == 1
        assert all_exp[0].reward == 0.5

    def test_ring_buffer_overflow(self):
        from helix_cdc.runtime.policy_trainer import (
            ExperienceBuffer, Experience
        )
        from helix_cdc.runtime.policy_harness import PolicyState, PolicyAction

        buffer = ExperienceBuffer(max_size=5)

        for i in range(10):
            exp = Experience(
                state=PolicyState(tok_per_sec=float(i)),
                action=PolicyAction.default(),
                reward=float(i),
                next_state=PolicyState(),
            )
            buffer.add(exp)

        assert len(buffer) == 5  # Max size

        # Should have most recent experiences
        rewards = [e.reward for e in buffer.get_all()]
        assert max(rewards) == 9.0

    def test_sample(self):
        from helix_cdc.runtime.policy_trainer import (
            ExperienceBuffer, Experience
        )
        from helix_cdc.runtime.policy_harness import PolicyState, PolicyAction

        buffer = ExperienceBuffer(max_size=100)

        for i in range(50):
            exp = Experience(
                state=PolicyState(),
                action=PolicyAction.default(),
                reward=float(i),
                next_state=PolicyState(),
            )
            buffer.add(exp)

        batch = buffer.sample(batch_size=10)
        assert len(batch) == 10


# =============================================================================
# Training Tests
# =============================================================================

class TestTraining:
    def test_compute_returns(self):
        from helix_cdc.runtime.policy_trainer import compute_returns

        rewards = [1.0, 1.0, 1.0, 1.0, 1.0]
        returns = compute_returns(rewards, gamma=0.99, normalize=False)

        # First return should be highest (sum of discounted rewards)
        assert returns[0] > returns[-1]

        # With gamma=1.0 and no normalization, first return = 5
        returns_undiscounted = compute_returns([1.0] * 5, gamma=1.0, normalize=False)
        assert abs(returns_undiscounted[0] - 5.0) < 1e-5

    def test_train_reinforce_runs(self):
        """Test that training runs without error."""
        from helix_cdc.runtime.policy_harness import PolicyHarness
        from helix_cdc.runtime.policy_trainer import (
            ExperienceBuffer,
            generate_synthetic_experiences,
            train_reinforce,
        )

        experiences = generate_synthetic_experiences(n_experiences=50, seed=42)

        buffer = ExperienceBuffer(max_size=100)
        buffer.extend(experiences)

        harness = PolicyHarness.default()
        for exp in experiences:
            harness.normalizer.update(exp.state)

        stats = train_reinforce(
            harness=harness,
            buffer=buffer,
            epochs=5,
            lr=0.01,
            verbose=False,
        )

        assert len(stats) == 5
        assert stats[-1].n_experiences == 50


# =============================================================================
# StateNormalizer Tests
# =============================================================================

class TestStateNormalizer:
    def test_update_and_normalize(self):
        from helix_cdc.runtime.policy_harness import StateNormalizer, PolicyState

        normalizer = StateNormalizer()

        # Add some states
        for i in range(100):
            state = PolicyState(
                tok_per_sec=30.0 + np.random.normal(0, 5),
                cache_hit_rate=0.8 + np.random.normal(0, 0.1),
            )
            normalizer.update(state)

        assert normalizer.count == 100

        # Normalized state should be approximately zero mean
        state = PolicyState(tok_per_sec=30.0, cache_hit_rate=0.8)
        normalized = normalizer.normalize(state)
        # tok_per_sec (index 0) should be close to 0
        assert abs(normalized[0]) < 2.0  # Within 2 std

    def test_serialization(self):
        from helix_cdc.runtime.policy_harness import StateNormalizer, PolicyState

        normalizer = StateNormalizer()
        for i in range(50):
            state = PolicyState(tok_per_sec=float(i))
            normalizer.update(state)

        d = normalizer.to_dict()
        normalizer2 = StateNormalizer.from_dict(d)

        assert normalizer2.count == normalizer.count
        assert np.allclose(normalizer2.mean, normalizer.mean)


# =============================================================================
# Integration Tests
# =============================================================================

class TestIntegration:
    def test_full_pipeline(self):
        """Test full train -> save -> load -> verify pipeline."""
        from helix_cdc.runtime.policy_harness import PolicyHarness, PolicyState
        from helix_cdc.runtime.policy_trainer import (
            ExperienceBuffer,
            generate_synthetic_experiences,
            train_reinforce,
            create_training_receipt,
        )

        # Generate experiences
        experiences = generate_synthetic_experiences(n_experiences=100, seed=42)

        # Create buffer
        buffer = ExperienceBuffer(max_size=200)
        buffer.extend(experiences)

        # Create and train harness
        harness = PolicyHarness.default()
        for exp in experiences:
            harness.normalizer.update(exp.state)

        stats = train_reinforce(
            harness=harness,
            buffer=buffer,
            epochs=10,
            lr=0.01,
            verbose=False,
        )

        # Create receipt
        receipt = create_training_receipt(
            stats_list=stats,
            harness=harness,
            input_pattern="synthetic",
            output_path="/tmp/test_policy.pt",
        )

        assert receipt["schema"] == "echo.policy.training.v0"
        assert receipt["training"]["n_epochs"] == 10

        # Save and load
        with tempfile.NamedTemporaryFile(suffix=".pt", delete=False) as f:
            path = f.name

        try:
            harness.save_checkpoint(path, extra={"receipt": receipt})
            harness2 = PolicyHarness.from_checkpoint(path)

            # Test that loaded model works
            state = PolicyState(tok_per_sec=40.0, cache_hit_rate=0.85)
            action = harness2.select_action(state)
            assert action is not None
        finally:
            os.unlink(path)


# =============================================================================
# Main
# =============================================================================

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