#!/usr/bin/env python3
"""
helix_cdc/runtime/policy_trainer.py

WO-ECHO-POLICY-HARNESS-01 - Policy Training Infrastructure

Provides REINFORCE training for the policy harness:
- ExperienceBuffer: Ring buffer for SARS (state, action, reward, next_state) tuples
- compute_returns: Discounted return calculation
- train_reinforce: REINFORCE with baseline training loop
- collect_experiences_from_jsonl: Parse telemetry JSONL into experiences

Usage:
    from helix_cdc.runtime.policy_trainer import (
        ExperienceBuffer,
        train_reinforce,
        collect_experiences_from_jsonl,
    )

    # Collect experiences from telemetry
    experiences = collect_experiences_from_jsonl("receipts/runtime/*.jsonl")

    # Train
    buffer = ExperienceBuffer(max_size=10000)
    buffer.extend(experiences)

    harness = PolicyHarness.default()
    train_reinforce(harness, buffer, epochs=100)
"""

from __future__ import annotations

import glob
import json
import statistics
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import numpy as np

from .policy_harness import (
    PolicyState,
    PolicyAction,
    PolicyMLP,
    PolicyHarness,
    StateNormalizer,
    compute_reward,
    _import_torch,
)


# =============================================================================
# Experience Buffer
# =============================================================================

@dataclass
class Experience:
    """Single experience tuple (s, a, r, s', done)."""
    state: PolicyState
    action: PolicyAction
    reward: float
    next_state: PolicyState
    done: bool = False
    log_prob: Optional[float] = None
    correctness_ok: bool = True


class ExperienceBuffer:
    """
    Ring buffer for experience storage.

    Stores (state, action, reward, next_state, done) tuples
    for training the policy.
    """

    def __init__(self, max_size: int = 10000):
        self.max_size = max_size
        self.buffer: List[Experience] = []
        self._idx = 0

    def add(self, experience: Experience) -> None:
        """Add single experience."""
        if len(self.buffer) < self.max_size:
            self.buffer.append(experience)
        else:
            self.buffer[self._idx] = experience
        self._idx = (self._idx + 1) % self.max_size

    def extend(self, experiences: List[Experience]) -> None:
        """Add multiple experiences."""
        for exp in experiences:
            self.add(exp)

    def sample(self, batch_size: int) -> List[Experience]:
        """Sample random batch."""
        indices = np.random.randint(0, len(self.buffer), size=batch_size)
        return [self.buffer[i] for i in indices]

    def get_all(self) -> List[Experience]:
        """Get all experiences."""
        return list(self.buffer)

    def clear(self) -> None:
        """Clear buffer."""
        self.buffer = []
        self._idx = 0

    def __len__(self) -> int:
        return len(self.buffer)


# =============================================================================
# Return Computation
# =============================================================================

def compute_returns(
    rewards: List[float],
    gamma: float = 0.99,
    normalize: bool = True,
) -> List[float]:
    """
    Compute discounted returns from rewards.

    Args:
        rewards: List of rewards [r_0, r_1, ..., r_T]
        gamma: Discount factor
        normalize: If True, normalize returns to zero mean, unit variance

    Returns:
        List of discounted returns [G_0, G_1, ..., G_T]
    """
    T = len(rewards)
    returns = [0.0] * T

    # Backwards cumulative sum with discounting
    G = 0.0
    for t in reversed(range(T)):
        G = rewards[t] + gamma * G
        returns[t] = G

    if normalize and len(returns) > 1:
        mean = statistics.mean(returns)
        std = statistics.stdev(returns)
        if std > 1e-8:
            returns = [(r - mean) / std for r in returns]

    return returns


# =============================================================================
# Training Loop
# =============================================================================

@dataclass
class TrainingStats:
    """Statistics from training."""
    epoch: int = 0
    loss: float = 0.0
    mean_reward: float = 0.0
    mean_return: float = 0.0
    n_experiences: int = 0
    grad_norm: float = 0.0


def train_reinforce(
    harness: PolicyHarness,
    buffer: ExperienceBuffer,
    epochs: int = 100,
    lr: float = 0.001,
    gamma: float = 0.99,
    baseline: str = "mean",  # "mean" or "none"
    grad_clip: float = 1.0,
    verbose: bool = True,
) -> List[TrainingStats]:
    """
    Train policy using REINFORCE with baseline.

    Args:
        harness: PolicyHarness with model to train
        buffer: Experience buffer with collected experiences
        epochs: Number of training epochs
        lr: Learning rate
        gamma: Discount factor for returns
        baseline: Baseline type ("mean" or "none")
        grad_clip: Gradient clipping threshold
        verbose: If True, print progress

    Returns:
        List of TrainingStats per epoch
    """
    torch, _ = _import_torch()

    model = harness.model
    model.train()

    optimizer = torch.optim.Adam(model.parameters(), lr=lr)

    experiences = buffer.get_all()
    if not experiences:
        raise ValueError("Empty experience buffer")

    # Extract rewards and compute returns
    rewards = [exp.reward for exp in experiences]
    returns = compute_returns(rewards, gamma=gamma, normalize=True)

    # Convert to tensors
    states = torch.stack([
        torch.tensor(harness.normalizer.normalize(exp.state), dtype=torch.float32)
        for exp in experiences
    ])

    # Build action tensors (for log prob calculation)
    # WO-ECHO-POLICY-GOLDEN-01: Only 4 knobs now
    actions = {
        "verify_policy": torch.tensor([exp.action.verify_policy for exp in experiences]),
        "cpp_kernel": torch.tensor([exp.action.cpp_kernel for exp in experiences]),
        "fused_matmul": torch.tensor([exp.action.fused_matmul for exp in experiences]),
        "fibpi_mode": torch.tensor([exp.action.fibpi_mode for exp in experiences]),
    }

    returns_tensor = torch.tensor(returns, dtype=torch.float32)

    # Apply baseline
    if baseline == "mean":
        returns_tensor = returns_tensor - returns_tensor.mean()

    stats_list = []

    for epoch in range(epochs):
        optimizer.zero_grad()

        # Forward pass
        probs = model(states)

        # Compute log probabilities for taken actions
        log_probs = torch.zeros(len(experiences))
        for name, action_indices in actions.items():
            head_probs = probs[name]  # [batch, n_actions]
            # Gather log probs for taken actions
            log_probs += torch.log(
                head_probs.gather(1, action_indices.unsqueeze(1)).squeeze(1) + 1e-8
            )

        # REINFORCE loss: -E[log_prob * return]
        loss = -(log_probs * returns_tensor).mean()

        # Backward
        loss.backward()

        # Gradient clipping
        grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)

        # Update
        optimizer.step()

        # Stats
        stats = TrainingStats(
            epoch=epoch,
            loss=loss.item(),
            mean_reward=statistics.mean(rewards),
            mean_return=statistics.mean(returns),
            n_experiences=len(experiences),
            grad_norm=grad_norm.item() if isinstance(grad_norm, torch.Tensor) else grad_norm,
        )
        stats_list.append(stats)

        if verbose and (epoch % 10 == 0 or epoch == epochs - 1):
            print(f"Epoch {epoch:4d} | Loss: {stats.loss:.4f} | "
                  f"Mean Return: {stats.mean_return:.4f} | "
                  f"Grad Norm: {stats.grad_norm:.4f}")

    model.eval()
    return stats_list


# =============================================================================
# Experience Collection from JSONL
# =============================================================================

def collect_experiences_from_jsonl(
    pattern: str,
    min_snapshots: int = 2,
) -> List[Experience]:
    """
    Parse telemetry JSONL files into experiences.

    Looks for consecutive snapshots and policy_action events to build
    (state, action, reward, next_state) tuples.

    Args:
        pattern: Glob pattern for JSONL files (e.g., "receipts/runtime/*.jsonl")
        min_snapshots: Minimum snapshots needed to form experiences

    Returns:
        List of Experience objects
    """
    experiences = []

    # Find all matching files
    files = glob.glob(pattern)
    if not files:
        print(f"Warning: No files match pattern '{pattern}'")
        return experiences

    for filepath in files:
        file_experiences = _parse_single_jsonl(filepath, min_snapshots)
        experiences.extend(file_experiences)

    return experiences


def _parse_single_jsonl(filepath: str, min_snapshots: int) -> List[Experience]:
    """Parse single JSONL file into experiences."""
    experiences = []

    # Read all events
    events = []
    with open(filepath, "r") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                events.append(json.loads(line))
            except json.JSONDecodeError:
                continue

    # Filter to snapshots and policy actions
    snapshots = [
        e for e in events
        if e.get("schema") == "echo.runtime.telemetry.v1"
    ]
    policy_actions = {
        e.get("token_index"): e.get("action")
        for e in events
        if e.get("event") == "policy_action"
    }

    if len(snapshots) < min_snapshots:
        return experiences

    # Build experiences from consecutive snapshots
    for i in range(len(snapshots) - 1):
        snap = snapshots[i]
        next_snap = snapshots[i + 1]

        # Extract state
        state = PolicyState.from_telemetry_snapshot(snap)
        next_state = PolicyState.from_telemetry_snapshot(next_snap)

        # Check for policy action at this token index
        token_idx = snap.get("token_index", 0)
        action_dict = policy_actions.get(token_idx)

        if action_dict:
            action = PolicyAction.from_dict(action_dict)
        else:
            # Default action if none recorded
            action = PolicyAction.default()

        # Compute reward (assume correctness OK if no verification failure recorded)
        correctness_ok = True  # TODO: Check verification events
        reward = compute_reward(state, action, next_state, correctness_ok)

        # Check if this is terminal
        done = (i == len(snapshots) - 2)

        exp = Experience(
            state=state,
            action=action,
            reward=reward,
            next_state=next_state,
            done=done,
            correctness_ok=correctness_ok,
        )
        experiences.append(exp)

    return experiences


# =============================================================================
# Training Receipt
# =============================================================================

def create_training_receipt(
    stats_list: List[TrainingStats],
    harness: PolicyHarness,
    input_pattern: str,
    output_path: str,
    extra: Optional[Dict[str, Any]] = None,
) -> Dict[str, Any]:
    """
    Create receipt documenting training run.

    Args:
        stats_list: List of training stats
        harness: Trained harness
        input_pattern: Input file pattern used
        output_path: Where checkpoint was saved
        extra: Additional fields

    Returns:
        Receipt dict
    """
    final_stats = stats_list[-1] if stats_list else TrainingStats()

    receipt = {
        "schema": "echo.policy.training.v0",
        "work_order": "WO-ECHO-POLICY-HARNESS-01",
        "timestamp": datetime.utcnow().isoformat() + "Z",
        "input": {
            "pattern": input_pattern,
        },
        "output": {
            "checkpoint_path": output_path,
        },
        "model_config": {
            "state_dim": harness.model.state_dim,
            "hidden_dim": harness.model.hidden_dim,
            "n_params": harness.model.count_parameters(),
        },
        "training": {
            "n_experiences": final_stats.n_experiences,
            "n_epochs": len(stats_list),
            "final_loss": final_stats.loss,
            "mean_reward": final_stats.mean_reward,
        },
        "normalizer": {
            "count": harness.normalizer.count,
        },
    }

    if extra:
        receipt.update(extra)

    return receipt


# =============================================================================
# Synthetic Experience Generation (for testing/bootstrap)
# =============================================================================

def generate_synthetic_experiences(
    n_experiences: int = 1000,
    seed: int = 42,
) -> List[Experience]:
    """
    Generate synthetic experiences for testing/bootstrap.

    Creates experiences with realistic state distributions
    and rewards that correlate with throughput improvements.

    Args:
        n_experiences: Number of experiences to generate
        seed: Random seed

    Returns:
        List of synthetic experiences
    """
    rng = np.random.default_rng(seed)
    experiences = []

    # Base state parameters
    base_tok_per_sec = 30.0
    base_step_ms = 30.0

    for i in range(n_experiences):
        # Generate random state
        state = PolicyState(
            tok_per_sec=base_tok_per_sec + rng.normal(0, 10),
            mean_step_ms=base_step_ms + rng.normal(0, 5),
            p95_step_ms=base_step_ms * 1.5 + rng.normal(0, 10),
            steps_in_window=int(rng.integers(50, 150)),
            cache_hit_rate=rng.uniform(0.6, 0.95),
            cache_misses=int(rng.integers(5, 50)),
            cache_hits=int(rng.integers(50, 150)),
            regen_hit_rate=rng.uniform(0.8, 1.0),
            regen_time_pct=rng.uniform(0, 10),
            regen_calls=int(rng.integers(0, 20)),
            rss_delta_mb=rng.normal(0, 20),
            rss_mb=rng.uniform(400, 600),
            se_mean=rng.uniform(0.2, 0.4),
            se_std=rng.uniform(0.01, 0.1),
        )

        # Generate random action
        action = PolicyAction.random(rng)

        # Simulate effect of action on next state
        # WO-ECHO-POLICY-GOLDEN-01: Only wirable knobs affect throughput
        # cpp_kernel and fused_matmul improve throughput
        tps_bonus = 0.0
        if action.cpp_kernel:
            tps_bonus += rng.uniform(5, 15)
        if action.fused_matmul:
            tps_bonus += rng.uniform(3, 10)

        # se_bucket mode may improve routing efficiency
        if action.fibpi_mode == 1:  # se_bucket
            tps_bonus += rng.uniform(0, 5)

        next_state = PolicyState(
            tok_per_sec=state.tok_per_sec + tps_bonus + rng.normal(0, 5),
            mean_step_ms=state.mean_step_ms - tps_bonus * 0.3 + rng.normal(0, 2),
            p95_step_ms=state.p95_step_ms - tps_bonus * 0.2 + rng.normal(0, 5),
            steps_in_window=state.steps_in_window + int(rng.integers(-10, 20)),
            cache_hit_rate=min(0.99, state.cache_hit_rate + rng.normal(0, 0.02)),
            cache_misses=max(0, state.cache_misses + int(rng.integers(-5, 5))),
            cache_hits=state.cache_hits + int(rng.integers(0, 20)),
            regen_hit_rate=state.regen_hit_rate + rng.normal(0, 0.02),
            regen_time_pct=max(0, state.regen_time_pct + rng.normal(0, 1)),
            regen_calls=max(0, state.regen_calls + int(rng.integers(-2, 5))),
            rss_delta_mb=rng.normal(0, 15),
            rss_mb=state.rss_mb + rng.normal(0, 10),
            se_mean=state.se_mean + rng.normal(0, 0.02),
            se_std=state.se_std + rng.normal(0, 0.005),
        )

        # Compute reward
        correctness_ok = rng.random() > 0.05  # 5% failure rate
        reward = compute_reward(state, action, next_state, correctness_ok)

        exp = Experience(
            state=state,
            action=action,
            reward=reward,
            next_state=next_state,
            done=(i == n_experiences - 1),
            correctness_ok=correctness_ok,
        )
        experiences.append(exp)

    return experiences


# WO-ECHO-POLICY-GOLDEN-01: Removed CACHE_SIZE_VALUES import (no longer used)


# =============================================================================
# Self-test
# =============================================================================

if __name__ == "__main__":
    print("=" * 60)
    print("WO-ECHO-POLICY-HARNESS-01 - Policy Trainer Self-Test")
    print("=" * 60)
    print()

    # Test ExperienceBuffer
    print("1. ExperienceBuffer")
    buffer = ExperienceBuffer(max_size=100)
    print(f"   Created buffer with max_size=100")

    # Generate synthetic experiences
    print("\n2. generate_synthetic_experiences")
    experiences = generate_synthetic_experiences(n_experiences=200, seed=42)
    print(f"   Generated {len(experiences)} experiences")

    buffer.extend(experiences)
    print(f"   Buffer size after extend: {len(buffer)}")

    # Test compute_returns
    print("\n3. compute_returns")
    rewards = [exp.reward for exp in experiences[:10]]
    returns = compute_returns(rewards, gamma=0.99)
    print(f"   Rewards[:5]: {[f'{r:.3f}' for r in rewards[:5]]}")
    print(f"   Returns[:5]: {[f'{r:.3f}' for r in returns[:5]]}")

    # Test training
    print("\n4. train_reinforce")
    harness = PolicyHarness.default()

    # Update normalizer with experience states
    for exp in experiences:
        harness.normalizer.update(exp.state)

    print(f"   Normalizer count: {harness.normalizer.count}")

    stats = train_reinforce(
        harness=harness,
        buffer=buffer,
        epochs=20,
        lr=0.01,
        gamma=0.99,
        verbose=True,
    )

    final = stats[-1]
    print(f"\n   Final loss: {final.loss:.4f}")
    print(f"   Mean reward: {final.mean_reward:.4f}")

    # Test training receipt
    print("\n5. create_training_receipt")
    receipt = create_training_receipt(
        stats_list=stats,
        harness=harness,
        input_pattern="synthetic",
        output_path="/tmp/test_policy.pt",
    )
    print(f"   Receipt schema: {receipt['schema']}")
    print(f"   Model params: {receipt['model_config']['n_params']}")
    print(f"   Final loss: {receipt['training']['final_loss']:.4f}")

    # Save and reload
    print("\n6. Save/load trained harness")
    test_path = "/tmp/policy_trainer_selftest.pt"
    harness.save_checkpoint(test_path, extra={"training_receipt": receipt})
    print(f"   Saved to {test_path}")

    harness2 = PolicyHarness.from_checkpoint(test_path)
    print(f"   Loaded successfully")

    # Test action selection
    state = experiences[0].state
    action1 = harness.select_action(state)
    action2 = harness2.select_action(state)
    print(f"   Original harness action: {action1.to_dict()}")
    print(f"   Loaded harness action: {action2.to_dict()}")

    print()
    print("=" * 60)
    print("Self-test PASSED")
    print("=" * 60)
