import torch
import torch.nn as nn
import triton
import triton.language as tl
from collections import OrderedDict
from typing import Dict, List, Optional, Any, Tuple

# --- Paged KV Cache Management ---

class PagedKVCacheManager:
    """Manages KV cache using a paged/block-based approach."""
    def __init__(self, num_blocks: int, block_size: int, num_heads: int, head_dim: int, device: str = "cuda", dtype: torch.dtype = torch.float16):
        self.num_blocks = num_blocks
        self.block_size = block_size
        self.num_heads = num_heads
        self.head_dim = head_dim
        self.device = torch.device(device)
        self.dtype = dtype
        self.k_pool = torch.zeros((num_blocks, block_size, num_heads, head_dim), device=self.device, dtype=self.dtype)
        self.v_pool = torch.zeros((num_blocks, block_size, num_heads, head_dim), device=self.device, dtype=self.dtype)
        self.free_blocks = list(range(num_blocks))
        self.allocated_blocks: Dict[str, List[int]] = {}

    def allocate(self, session_id: str, num_required_tokens: int):
        num_required_blocks = (num_required_tokens + self.block_size - 1) // self.block_size
        if len(self.free_blocks) < num_required_blocks: raise MemoryError("Out of KV Cache Blocks!")
        blocks = [self.free_blocks.pop() for _ in range(num_required_blocks)]
        self.allocated_blocks[session_id] = blocks
        return blocks

    def free(self, session_id: str):
        if session_id in self.allocated_blocks:
            self.free_blocks.extend(self.allocated_blocks[session_id])
            del self.allocated_blocks[session_id]

# --- Automated Expert Management with LRU & Dynamic Quantization ---

class LRUExpertCache:
    """Manages expert placement with on-the-fly INT8 quantization."""
    def __init__(self, capacity: int, device: str = "cuda", quantize: bool = True):
        self.capacity = capacity
        self.device = torch.device(device)
        self.quantize = quantize
        self.cache = OrderedDict()
        self.cpu_store = {}

    def _quantize_module(self, module: nn.Module) -> Dict[str, torch.Tensor]:
        quantized_dict = {}
        with torch.no_grad():
            for name, param in module.state_dict().items():
                if self.quantize and param.dim() > 1:
                    scale = param.abs().max() / 127.0
                    q_param = (param / (scale + 1e-8)).round().to(torch.int8)
                    quantized_dict[name] = (q_param, scale.to(torch.float16))
                else: quantized_dict[name] = param.cpu().half()
        return quantized_dict

    def _dequantize_module(self, module: nn.Module, quantized_dict: Dict[str, Any]):
        state_dict = {}
        with torch.no_grad():
            for name, data in quantized_dict.items():
                if isinstance(data, tuple):
                    q_data, scale = data
                    state_dict[name] = (q_data.float() * scale.float()).half().to(self.device)
                else: state_dict[name] = data.to(self.device)
        module.load_state_dict(state_dict)

    def get_expert(self, expert_id: str, expert_module: nn.Module) -> nn.Module:
        if expert_id in self.cache:
            self.cache.move_to_end(expert_id)
            return self.cache[expert_id]
        if len(self.cache) >= self.capacity:
            evict_id, evict_module = self.cache.popitem(last=False)
            self.cpu_store[evict_id] = self._quantize_module(evict_module)
            evict_module.to("cpu")
        if expert_id in self.cpu_store:
            self._dequantize_module(expert_module, self.cpu_store[expert_id])
            del self.cpu_store[expert_id]
        expert_module.to(self.device)
        self.cache[expert_id] = expert_module
        return expert_module

class ManagedMoELayer(nn.Module):
    def __init__(self, moe_layer: nn.Module, layer_id: str, expert_cache: LRUExpertCache):
        super().__init__()
        self.moe_layer = moe_layer
        self.layer_id = layer_id
        self.expert_cache = expert_cache
        self._is_graph_capture = False
        
        # Real-time Telemetry for UI
        self.usage_stats = {
            'deep': [0] * getattr(moe_layer, 'num_deep_experts', 4),
            'hot': [0] * getattr(moe_layer, 'num_experts', 8),
            'cold': [0] * getattr(moe_layer, 'num_experts', 8)
        }

    def __getattr__(self, name: str) -> Any:
        try:
            return super().__getattr__(name)
        except AttributeError:
            return getattr(self.moe_layer, name)

    def forward(self, hidden_states: torch.Tensor, **kwargs) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
        top_k_probs, top_k_indices, router_logits = self.moe_layer.router(hidden_states)
        batch_size, seq_len, hidden_size = hidden_states.shape
        hidden_flat = hidden_states.view(-1, hidden_size)
        original_dtype = hidden_states.dtype
        final_output = torch.zeros_like(hidden_flat)
        
        # In graph capture mode, we MUST NOT trigger expert cache movement (CPU-GPU sync/Python overhead)
        if self._is_graph_capture:
            # Replay-safe path: Use experts already in memory, or mask out those that aren't
            # Production Grade: We should have pre-cached the necessary experts for the graph
            for expert_idx in range(self.moe_layer.num_experts):
                expert = self.moe_layer.experts[expert_idx]
                # Only run if expert is on the correct device
                if expert.weight.device == hidden_flat.device:
                    for k in range(self.moe_layer.num_experts_per_tok):
                        mask = (top_k_indices[:, k] == expert_idx)
                        if mask.any():
                            expert_output = expert(hidden_flat[mask])
                            final_output[mask] += (top_k_probs[mask, k:k+1] * expert_output).to(original_dtype)
            shared_output = self.moe_layer.shared_expert(hidden_flat)
            final_output += shared_output.to(original_dtype)
            return final_output.view(batch_size, seq_len, hidden_size), None

        # Dynamic path (Standard Inference)
        total_experts = getattr(self.moe_layer, 'num_experts', 0)
        num_deep_experts = getattr(self.moe_layer, 'num_deep_experts', 0)
        
        unique_experts = torch.unique(top_k_indices).tolist()
        for expert_idx in unique_experts:
            # Deep Expert Path
            if expert_idx >= total_experts:
                deep_idx = expert_idx - total_experts
                if deep_idx < num_deep_experts:
                    self.usage_stats['deep'][deep_idx] += 1
                    expert = self.moe_layer.deep_experts[deep_idx]
                    for k in range(self.moe_layer.num_experts_per_tok):
                        mask = (top_k_indices[:, k] == expert_idx)
                        if mask.any():
                            expert_output = expert(hidden_flat[mask])
                            final_output[mask] += (top_k_probs[mask, k:k+1] * expert_output).to(original_dtype)
            # Standard Routed Expert Path
            else:
                expert_id = f"{self.layer_id}.expert.{expert_idx}"
                # Telemetry tracking for LRU Cache
                if expert_id in self.expert_cache.cache:
                    self.usage_stats['hot'][expert_idx] += 1
                else:
                    self.usage_stats['cold'][expert_idx] += 1
                
                managed_expert = self.expert_cache.get_expert(expert_id, self.moe_layer.experts[expert_idx])
                for k in range(self.moe_layer.num_experts_per_tok):
                    mask = (top_k_indices[:, k] == expert_idx)
                    if mask.any():
                        expert_output = managed_expert(hidden_flat[mask])
                        final_output[mask] += (top_k_probs[mask, k:k+1] * expert_output).to(original_dtype)
        
        shared_output = self.moe_layer.shared_expert(hidden_flat)
        final_output += shared_output.to(original_dtype)
        return final_output.view(batch_size, seq_len, hidden_size), None

# --- Production CUDA Graph Management ---

class GraphRunner:
    """
    SOTA CUDA Graph Runner for elite low-latency inference.
    Captures a model's forward pass within a graph to eliminate Python/CPU overhead.
    """
    def __init__(self, model: nn.Module, batch_size: int = 1, max_seq_len: int = 1):
        self.model = model
        self.batch_size = batch_size
        self.max_seq_len = max_seq_len
        self.graph = None
        self.static_inputs = {}
        self.static_outputs = None
        self.device = next(model.parameters()).device
        self.dtype = next(model.parameters()).dtype

    def capture(self, **sample_inputs):
        """Captures the model execution into a CUDA Graph."""
        print(f"[XoronEngine] ⚡ Capturing CUDA Graph (Batch: {self.batch_size}, Seq: {self.max_seq_len})...")
        
        # 1. Allocate Static Buffers
        self.static_inputs = {
            name: tensor.clone().detach().to(self.device) 
            for name, tensor in sample_inputs.items()
        }
        
        # 2. Warmup Runs
        # Crucial for stable graph capture: ensures memory allocators and kernels are ready
        torch.cuda.synchronize()
        with torch.no_grad():
            for _ in range(3):
                self.model(**self.static_inputs)
        torch.cuda.synchronize()

        # 3. Capture Graph
        self.graph = torch.cuda.CUDAGraph()
        with torch.cuda.graph(self.graph):
            self.static_outputs = self.model(**self.static_inputs)
        
        torch.cuda.synchronize()
        print("[XoronEngine] ✅ Graph Capture Complete.")

    def run(self, **inputs):
        """Replays the captured graph with new inputs."""
        # 1. Copy inputs into static buffers
        for name, tensor in inputs.items():
            if name in self.static_inputs:
                self.static_inputs[name].copy_(tensor)
        
        # 2. Replay Graph
        self.graph.replay()
        
        return self.static_outputs

# --- SOTA Multimodal Triton Kernels ---

@triton.jit
def fused_cross_attention_kernel(
    Q, K, V, Out,
    stride_qb, stride_qh, stride_qm, stride_qk,
    stride_kb, stride_kh, stride_kn, stride_kk,
    stride_vb, stride_vh, stride_vn, stride_vk,
    stride_ob, stride_oh, stride_om, stride_ok,
    n_heads, d_model, seq_len_q, seq_len_kv,
    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_DMODEL: tl.constexpr,
):
    """Fused Cross-Attention for Text-to-Image/Video/Audio Guidance."""
    pid = tl.program_id(0)
    num_pid_m = tl.cdiv(seq_len_q, BLOCK_M)
    pid_m = pid % num_pid_m
    pid_h = pid // num_pid_m
    rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    rn = tl.arange(0, BLOCK_N)
    rk = tl.arange(0, BLOCK_DMODEL)
    Q_ptr = Q + pid_h * stride_qh + rm[:, None] * stride_qm + rk[None, :] * stride_qk
    K_ptr = K + pid_h * stride_kh + rn[None, :] * stride_kn + rk[:, None] * stride_kk
    m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float('inf')
    l_i = tl.zeros([BLOCK_M], dtype=tl.float32)
    acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32)
    for start_n in range(0, seq_len_kv, BLOCK_N):
        k = tl.load(K + pid_h * stride_kh + (start_n + tl.arange(0, BLOCK_N))[None, :] * stride_kn + rk[:, None] * stride_kk)
        qk = tl.dot(tl.load(Q_ptr).to(tl.float16), k.to(tl.float16)) * 0.125
        m_ij = tl.max(qk, 1)
        p = tl.exp(qk - m_ij[:, None])
        l_ij = tl.sum(p, 1)
        m_i_new = tl.maximum(m_i, m_ij)
        alpha = tl.exp(m_i - m_i_new)
        beta = tl.exp(m_ij - m_i_new)
        l_i = l_i * alpha + l_ij * beta
        v = tl.load(V + pid_h * stride_vh + (start_n + tl.arange(0, BLOCK_N))[:, None] * stride_vn + rk[None, :] * stride_vk)
        acc = acc * alpha[:, None] + tl.dot(p.to(tl.float16), v.to(tl.float16))
        m_i = m_i_new
    acc /= l_i[:, None]
    tl.store(Out + pid_h * stride_oh + rm[:, None] * stride_om + rk[None, :] * stride_ok, acc)

# Native SOTA wrappers removed (optimized internally instead)

@triton.jit
def video_3d_rope_kernel(
    Q, Cos_T, Sin_T, Cos_H, Sin_H, Cos_W, Sin_W,
    stride_qt, stride_qx, stride_qy, stride_qh, stride_qd,
    n_t, n_x, n_y, n_h, n_d,
    BLOCK_D: tl.constexpr,
):
    """SOTA 3D-RoPE for Video Generation."""
    pid_t = tl.program_id(0)
    pid_x = tl.program_id(1)
    pid_y = tl.program_id(2)
    pid_h = tl.program_id(3) # Head index
    
    # Each thread block handles one head for one (T, X, Y) position
    q_ptr = Q + pid_t * stride_qt + pid_x * stride_qx + pid_y * stride_qy + pid_h * stride_qh
    
    # Offsets for the three chunks of the hidden dimension (D = BLOCK_D)
    # 3D RoPE splits the embedding into 3 parts: T, H, W
    d_third = BLOCK_D // 3
    
    # Part 1: Temporal (T)
    offsets_t_half = tl.arange(0, d_third // 2)
    q1_t = tl.load(q_ptr + offsets_t_half)
    q2_t = tl.load(q_ptr + offsets_t_half + d_third // 2)
    cos_t = tl.load(Cos_T + pid_t * (d_third // 2) + offsets_t_half)
    sin_t = tl.load(Sin_T + pid_t * (d_third // 2) + offsets_t_half)
    out1_t = q1_t * cos_t - q2_t * sin_t
    out2_t = q2_t * cos_t + q1_t * sin_t
    tl.store(q_ptr + offsets_t_half, out1_t)
    tl.store(q_ptr + offsets_t_half + d_third // 2, out2_t)

    # Part 2: Spatial Height (H/X)
    offset_h_base = d_third
    offsets_h_half = tl.arange(0, d_third // 2)
    q1_h = tl.load(q_ptr + offset_h_base + offsets_h_half)
    q2_h = tl.load(q_ptr + offset_h_base + offsets_h_half + d_third // 2)
    cos_h = tl.load(Cos_H + pid_x * (d_third // 2) + offsets_h_half)
    sin_h = tl.load(Sin_H + pid_x * (d_third // 2) + offsets_h_half)
    out1_h = q1_h * cos_h - q2_h * sin_h
    out2_h = q2_h * cos_h + q1_h * sin_h
    tl.store(q_ptr + offset_h_base + offsets_h_half, out1_h)
    tl.store(q_ptr + offset_h_base + offsets_h_half + d_third // 2, out2_h)
    
    # Part 3: Spatial Width (W/Y)
    offset_w_base = d_third * 2
    offsets_w_half = tl.arange(0, d_third // 2)
    q1_w = tl.load(q_ptr + offset_w_base + offsets_w_half)
    q2_w = tl.load(q_ptr + offset_w_base + offsets_w_half + d_third // 2)
    cos_w = tl.load(Cos_W + pid_y * (d_third // 2) + offsets_w_half)
    sin_w = tl.load(Sin_W + pid_y * (d_third // 2) + offsets_w_half)
    out1_w = q1_w * cos_w - q2_w * sin_w
    out2_w = q2_w * cos_w + q1_w * sin_w
    tl.store(q_ptr + offset_w_base + offsets_w_half, out1_w)
    tl.store(q_ptr + offset_w_base + offsets_w_half + d_third // 2, out2_w)


@triton.jit
def waveform_norm_kernel(X, Y, Weight, Bias, N, eps, BLOCK: tl.constexpr):
    """Accelerated raw waveform RMSNorm for Audio ASR/TTS."""
    pid = tl.program_id(0)
    
    # Pointer to the start of this sequence/batch item
    x_ptr = X + pid * N
    y_ptr = Y + pid * N
    
    # Compute variance using block loop
    _sum = tl.zeros([BLOCK], dtype=tl.float32)
    for start in range(0, N, BLOCK):
        offsets = start + tl.arange(0, BLOCK)
        mask = offsets < N
        x = tl.load(x_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
        _sum += x * x
    
    var = tl.sum(_sum, axis=0) / N
    rstd = 1.0 / tl.sqrt(var + eps)
    
    # Apply normalization and learned weights
    for start in range(0, N, BLOCK):
        offsets = start + tl.arange(0, BLOCK)
        mask = offsets < N
        x = tl.load(x_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
        
        # Apply RMS Norm
        x_hat = x * rstd
        
        # Apply affine transform
        w = tl.load(Weight + offsets, mask=mask, other=0.0).to(tl.float32)
        b = tl.load(Bias + offsets, mask=mask, other=0.0).to(tl.float32)
        y = x_hat * w + b
        
        # Store result
        tl.store(y_ptr + offsets, y.to(X.dtype.element_ty), mask=mask)
