#!/usr/bin/env python3
"""
generate.py - Minimal Streaming Generation for TinyLlama

WO-GENERATE-V1: Autoregressive text generation using CDNA streaming path.
No llama.cpp fallback. Pure Python. KV cache managed explicitly.

Usage:
    python -m helix_cdc.generate \
        --model ~/helix-cdc/models/gguf/tinyllama-1.1b-chat-v1.0.Q8_0.gguf \
        --cdna-dir /tmp/tinyllama_cdna \
        --prompt "Paris is the capital of" \
        --max-new-tokens 8

Success criterion: Produces coherent text using streaming CDNA weights.
"""

from __future__ import annotations

import argparse
import json
import os
import resource
import time
from dataclasses import dataclass, field, asdict
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import numpy as np

# Tripwire: ensure no llama.cpp imports
os.environ["HELIX_GENERATE_NATIVE"] = "1"


# =============================================================================
# KV Cache
# =============================================================================

@dataclass
class KVCache:
    """Per-layer KV cache for autoregressive generation."""
    k: np.ndarray  # [seq_len, n_kv_heads, d_head]
    v: np.ndarray  # [seq_len, n_kv_heads, d_head]

    @classmethod
    def empty(cls, n_kv_heads: int, d_head: int) -> "KVCache":
        """Create empty cache."""
        return cls(
            k=np.zeros((0, n_kv_heads, d_head), dtype=np.float32),
            v=np.zeros((0, n_kv_heads, d_head), dtype=np.float32),
        )

    def append(self, k_new: np.ndarray, v_new: np.ndarray) -> None:
        """Append new K,V positions. Validates shapes."""
        if self.k.shape[0] == 0:
            # First append - initialize from new
            self.k = k_new
            self.v = v_new
        else:
            assert k_new.shape[1:] == self.k.shape[1:], \
                f"K shape mismatch: new {k_new.shape[1:]} vs cache {self.k.shape[1:]}"
            assert v_new.shape[1:] == self.v.shape[1:], \
                f"V shape mismatch: new {v_new.shape[1:]} vs cache {self.v.shape[1:]}"
            self.k = np.concatenate([self.k, k_new], axis=0)
            self.v = np.concatenate([self.v, v_new], axis=0)

    @property
    def seq_len(self) -> int:
        return self.k.shape[0]

    def size_bytes(self) -> int:
        return self.k.nbytes + self.v.nbytes


# =============================================================================
# Generation Receipt
# =============================================================================

@dataclass
class GenerateReceipt:
    """Receipt for generation run."""
    model: str
    cdna_dir: str
    prompt: str
    output_text: str
    output_tokens: List[int]
    tokens_generated: int
    peak_rss_mb: float
    prefill_ms: float
    decode_total_ms: float
    decode_ms_per_token: float
    backend: str = "streaming_cdna"
    kv_cache_mb: float = 0.0
    num_blocks: int = 22
    timestamp: str = field(default_factory=lambda: datetime.now(timezone.utc).isoformat())
    status: str = "PASS"


# =============================================================================
# RMS Norm
# =============================================================================

def rms_norm(x: np.ndarray, weight: np.ndarray, eps: float = 1e-5) -> np.ndarray:
    """RMSNorm: x * weight / RMS(x)"""
    rms = np.sqrt(np.mean(x ** 2, axis=-1, keepdims=True) + eps)
    return (x / rms) * weight


# =============================================================================
# RoPE
# =============================================================================

def apply_rope(
    x: np.ndarray,  # [seq, n_heads, d_head]
    start_pos: int,
    theta: float = 10000.0,
) -> np.ndarray:
    """Apply Rotary Position Embedding."""
    seq_len, n_heads, d_head = x.shape

    # Compute frequencies
    freqs = 1.0 / (theta ** (np.arange(0, d_head, 2, dtype=np.float32) / d_head))

    # Position indices
    positions = np.arange(start_pos, start_pos + seq_len, dtype=np.float32)

    # Angles: [seq, d_head/2]
    angles = np.outer(positions, freqs)

    # Complex rotation
    cos = np.cos(angles)  # [seq, d_head/2]
    sin = np.sin(angles)  # [seq, d_head/2]

    # Reshape x for rotation: split into pairs
    x_reshape = x.reshape(seq_len, n_heads, d_head // 2, 2)
    x0 = x_reshape[..., 0]  # [seq, n_heads, d_head/2]
    x1 = x_reshape[..., 1]  # [seq, n_heads, d_head/2]

    # Expand cos/sin for broadcasting: [seq, 1, d_head/2]
    cos = cos[:, np.newaxis, :]
    sin = sin[:, np.newaxis, :]

    # Apply rotation
    out0 = x0 * cos - x1 * sin
    out1 = x0 * sin + x1 * cos

    # Reassemble
    out = np.stack([out0, out1], axis=-1)  # [seq, n_heads, d_head/2, 2]
    out = out.reshape(seq_len, n_heads, d_head)

    return out


# =============================================================================
# Attention with KV Cache
# =============================================================================

def attention_with_cache(
    X: np.ndarray,  # [seq, d_model]
    Wq: np.ndarray,  # [d_model, d_model]
    Wk: np.ndarray,  # [d_kv, d_model] (transposed for GQA)
    Wv: np.ndarray,  # [d_kv, d_model] (transposed for GQA)
    Wo: np.ndarray,  # [d_model, d_model]
    kv_cache: Optional[KVCache],
    start_pos: int,
    n_heads: int = 32,
    n_kv_heads: int = 4,
    d_head: int = 64,
) -> Tuple[np.ndarray, KVCache]:
    """
    Multi-head attention with GQA and KV cache.

    Returns: (output [seq, d_model], updated_kv_cache)
    """
    seq_len = X.shape[0]

    # GGUF stores all projections as [out_features, in_features], so transpose all
    # Q projection: [seq, d_model] @ [d_model, n_heads*d_head].T -> [seq, n_heads*d_head]
    Q = X @ Wq.T  # [seq, n_heads * d_head]
    Q = Q.reshape(seq_len, n_heads, d_head)  # [seq, n_heads, d_head]

    # K, V projection (GQA: fewer KV heads)
    # Wk, Wv are [n_kv_heads*d_head, d_model], so X @ W.T -> [seq, d_kv]
    K_new = X @ Wk.T  # [seq, n_kv_heads * d_head]
    V_new = X @ Wv.T  # [seq, n_kv_heads * d_head]
    K_new = K_new.reshape(seq_len, n_kv_heads, d_head)
    V_new = V_new.reshape(seq_len, n_kv_heads, d_head)

    # Apply RoPE to Q and K_new
    Q = apply_rope(Q, start_pos)
    K_new = apply_rope(K_new, start_pos)

    # Update KV cache
    if kv_cache is None:
        kv_cache = KVCache.empty(n_kv_heads, d_head)
    kv_cache.append(K_new, V_new)

    # Get full K, V from cache
    K = kv_cache.k  # [total_seq, n_kv_heads, d_head]
    V = kv_cache.v  # [total_seq, n_kv_heads, d_head]

    # GQA: repeat K,V heads to match Q heads
    # n_heads = 32, n_kv_heads = 4, so each KV head serves 8 Q heads
    n_rep = n_heads // n_kv_heads
    K = np.repeat(K, n_rep, axis=1)  # [total_seq, n_heads, d_head]
    V = np.repeat(V, n_rep, axis=1)  # [total_seq, n_heads, d_head]

    # Attention scores: Q @ K^T
    # Q: [seq, n_heads, d_head]
    # K: [total_seq, n_heads, d_head]
    # scores: [n_heads, seq, total_seq]
    scale = 1.0 / np.sqrt(d_head)

    # Transpose for matmul: [n_heads, seq, d_head] @ [n_heads, d_head, total_seq]
    Q_t = Q.transpose(1, 0, 2)  # [n_heads, seq, d_head]
    K_t = K.transpose(1, 2, 0)  # [n_heads, d_head, total_seq]
    scores = np.matmul(Q_t, K_t) * scale  # [n_heads, seq, total_seq]

    # Causal mask (for prefill, later tokens can't attend to earlier)
    total_seq = K.shape[0]
    if seq_len > 1:
        # Prefill mode: apply causal mask
        mask = np.triu(np.ones((seq_len, total_seq), dtype=np.float32) * -1e9, k=1)
        # Shift mask based on position
        if start_pos > 0:
            mask = np.zeros((seq_len, total_seq), dtype=np.float32)
            for i in range(seq_len):
                mask[i, start_pos + i + 1:] = -1e9
        scores = scores + mask[np.newaxis, :, :]
    # For single token decode (seq_len=1), no mask needed - can attend to all prior

    # Softmax
    scores_max = scores.max(axis=-1, keepdims=True)
    scores_exp = np.exp(scores - scores_max)
    attn_weights = scores_exp / scores_exp.sum(axis=-1, keepdims=True)  # [n_heads, seq, total_seq]

    # Apply attention to V
    # V: [total_seq, n_heads, d_head] -> V_t: [n_heads, total_seq, d_head]
    V_t = V.transpose(1, 0, 2)  # [n_heads, total_seq, d_head]
    # attn_weights @ V_t: [n_heads, seq, total_seq] @ [n_heads, total_seq, d_head] -> [n_heads, seq, d_head]
    context = np.matmul(attn_weights, V_t)  # [n_heads, seq, d_head]

    # Reshape back: [seq, n_heads * d_head]
    context = context.transpose(1, 0, 2).reshape(seq_len, n_heads * d_head)

    # Output projection
    output = context @ Wo.T  # [seq, d_model]

    return output, kv_cache


# =============================================================================
# FFN (SwiGLU)
# =============================================================================

def ffn_swiglu(
    X: np.ndarray,  # [seq, d_model]
    W_gate: np.ndarray,  # [d_ffn, d_model] (GGUF format)
    W_up: np.ndarray,  # [d_ffn, d_model] (GGUF format)
    W_down: np.ndarray,  # [d_model, d_ffn] (GGUF format)
) -> np.ndarray:
    """SwiGLU FFN: down(silu(gate(x)) * up(x))"""
    # GGUF stores gate/up as [d_ffn, d_model], so transpose for X @ W.T
    gate = X @ W_gate.T  # [seq, d_ffn]
    up = X @ W_up.T  # [seq, d_ffn]

    # SiLU activation on gate
    silu_gate = gate * (1.0 / (1.0 + np.exp(-gate)))

    # Element-wise multiply
    hidden = silu_gate * up  # [seq, d_ffn]

    # Down projection: W_down is [d_model, d_ffn], so W_down.T is [d_ffn, d_model]
    # hidden [seq, d_ffn] @ W_down.T [d_ffn, d_model] -> [seq, d_model]
    output = hidden @ W_down.T  # [seq, d_model]

    return output


# =============================================================================
# Single Block Forward with KV Cache
# =============================================================================

def forward_block_with_cache(
    X: np.ndarray,  # [seq, d_model]
    block_idx: int,
    start_pos: int,
    kv_cache: Optional[KVCache],
    weights: Dict[str, np.ndarray],
    n_heads: int = 32,
    n_kv_heads: int = 4,
    d_head: int = 64,
    norm_eps: float = 1e-5,
) -> Tuple[np.ndarray, KVCache]:
    """
    Single transformer block forward with KV cache.

    Pre-norm architecture:
    1. X_norm = RMSNorm(X, attn_norm)
    2. attn_out, kv = attention(X_norm, kv_cache)
    3. X_mid = X + attn_out
    4. X_mid_norm = RMSNorm(X_mid, ffn_norm)
    5. ffn_out = ffn(X_mid_norm)
    6. output = X_mid + ffn_out
    """
    # Load weights for this block
    attn_norm = weights[f"blk.{block_idx}.attn_norm.weight"]
    ffn_norm = weights[f"blk.{block_idx}.ffn_norm.weight"]
    Wq = weights[f"blk.{block_idx}.attn_q.weight"]
    Wk = weights[f"blk.{block_idx}.attn_k.weight"]
    Wv = weights[f"blk.{block_idx}.attn_v.weight"]
    Wo = weights[f"blk.{block_idx}.attn_output.weight"]
    W_gate = weights[f"blk.{block_idx}.ffn_gate.weight"]
    W_up = weights[f"blk.{block_idx}.ffn_up.weight"]
    W_down = weights[f"blk.{block_idx}.ffn_down.weight"]

    # Step 1: Attention norm
    X_norm = rms_norm(X, attn_norm, norm_eps)

    # Step 2: Attention with KV cache
    attn_out, kv_cache = attention_with_cache(
        X_norm, Wq, Wk, Wv, Wo,
        kv_cache, start_pos,
        n_heads, n_kv_heads, d_head,
    )

    # Step 3: Residual
    X_mid = X + attn_out

    # Step 4: FFN norm
    X_mid_norm = rms_norm(X_mid, ffn_norm, norm_eps)

    # Step 5: FFN
    ffn_out = ffn_swiglu(X_mid_norm, W_gate, W_up, W_down)

    # Step 6: Residual
    output = X_mid + ffn_out

    return output, kv_cache


# =============================================================================
# Weight Loading
# =============================================================================

def load_weights_from_gguf(gguf_path: Path) -> Dict[str, np.ndarray]:
    """Load all weights from GGUF file."""
    from gguf import GGUFReader
    from gguf.quants import dequantize

    reader = GGUFReader(str(gguf_path))
    weights = {}

    for t in reader.tensors:
        # Dequantize
        try:
            data = dequantize(t.data, t.tensor_type)
        except:
            data = t.data.copy()
            if hasattr(t.tensor_type, 'name'):
                if t.tensor_type.name == "F16":
                    data = data.view(np.float16).astype(np.float32)
                elif t.tensor_type.name == "F32":
                    data = data.view(np.float32)

        weights[t.name] = data.astype(np.float32)

    return weights


def load_weights_from_cdna(cdna_dir: Path, manifest: dict) -> Dict[str, np.ndarray]:
    """Load weights from CDNA shards (streaming-capable)."""
    from helix_cdc.regrow.stream_xw_matmul import stream_xw_from_cdna

    # For now, we load full tensors from CDNA
    # TODO: Actually stream block-by-block
    weights = {}

    for shard in manifest.get("shards", []):
        tensor_name = shard["tensor_name"]
        cdna_path = cdna_dir / shard["file"]

        if cdna_path.exists():
            # Read indices + codebook, reconstruct
            from helix_cdc.regrow.cdna_reader import CDNAv2Reader
            reader = CDNAv2Reader(cdna_path)
            codebook = reader.codebook

            # Reconstruct full tensor
            all_indices = []
            for block_indices, start_row in reader.stream_blocks():
                all_indices.append(block_indices)
            indices = np.concatenate(all_indices, axis=0)
            weights[tensor_name] = codebook[indices]

    # Load 1D tensors (norms) from npy files
    for f in cdna_dir.glob("*.npy"):
        tensor_name = f.stem.replace("_", ".")
        weights[tensor_name] = np.load(f)

    return weights


# =============================================================================
# Tokenizer
# =============================================================================

class SimpleTokenizer:
    """Minimal tokenizer using llama.cpp."""

    def __init__(self, gguf_path: Path):
        from llama_cpp import Llama
        self._llm = Llama(
            model_path=str(gguf_path),
            n_ctx=512,
            n_gpu_layers=0,
            verbose=False,
        )
        self.eos_token_id = self._llm.token_eos()
        self.bos_token_id = self._llm.token_bos()

    def encode(self, text: str, add_bos: bool = True) -> List[int]:
        tokens = self._llm.tokenize(text.encode(), add_bos=add_bos)
        return tokens

    def decode(self, tokens: List[int]) -> str:
        return self._llm.detokenize(tokens).decode("utf-8", errors="replace")


# =============================================================================
# Main Generate Function
# =============================================================================

def generate(
    prompt: str,
    gguf_path: Path,
    cdna_dir: Optional[Path] = None,
    max_new_tokens: int = 32,
    temperature: float = 0.0,
    use_cdna: bool = False,
) -> GenerateReceipt:
    """
    Autoregressive generation using native Python path.

    No llama.cpp inference. Pure numpy + KV cache.
    """
    t0_total = time.perf_counter()

    # Get initial RSS
    def get_rss_mb() -> float:
        return resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024

    baseline_rss = get_rss_mb()
    peak_rss = baseline_rss

    # Load weights
    print("Loading weights...")
    t0 = time.perf_counter()
    if use_cdna and cdna_dir and cdna_dir.exists():
        manifest = json.loads((cdna_dir / "manifest.json").read_text())
        weights = load_weights_from_cdna(cdna_dir, manifest)
        backend = "streaming_cdna"
    else:
        weights = load_weights_from_gguf(gguf_path)
        backend = "gguf_direct"
    load_time = (time.perf_counter() - t0) * 1000
    print(f"  Loaded {len(weights)} tensors in {load_time:.0f}ms")
    peak_rss = max(peak_rss, get_rss_mb())

    # Model config (TinyLlama)
    num_blocks = 22
    d_model = 2048
    n_heads = 32
    n_kv_heads = 4
    d_head = 64
    vocab_size = 32000

    # Tokenizer
    print("Loading tokenizer...")
    tokenizer = SimpleTokenizer(gguf_path)

    # Tokenize prompt
    token_ids = tokenizer.encode(prompt)
    print(f"  Prompt: {len(token_ids)} tokens")

    # Embedding lookup
    embed_weight = weights["token_embd.weight"]  # [d_model, vocab_size] in GGUF
    # GGUF stores as [d_model, vocab_size], we need [vocab_size, d_model]
    if embed_weight.shape[0] == d_model:
        embed_weight = embed_weight.T

    X = embed_weight[token_ids]  # [seq, d_model]
    print(f"  Embedded: {X.shape}")

    # Prefill: run through all blocks
    print(f"Prefill ({num_blocks} blocks)...")
    t0_prefill = time.perf_counter()

    kv_caches: Dict[int, KVCache] = {}
    for blk in range(num_blocks):
        X, kv = forward_block_with_cache(
            X, blk, start_pos=0, kv_cache=None, weights=weights,
            n_heads=n_heads, n_kv_heads=n_kv_heads, d_head=d_head,
        )
        kv_caches[blk] = kv
        if blk % 5 == 0:
            print(f"  Block {blk}: X={X.shape}, KV seq={kv.seq_len}")

    prefill_ms = (time.perf_counter() - t0_prefill) * 1000
    print(f"  Prefill done in {prefill_ms:.0f}ms")
    peak_rss = max(peak_rss, get_rss_mb())

    # Final norm + output projection
    output_norm = weights["output_norm.weight"]
    output_weight = weights["output.weight"]  # [d_model, vocab_size] in GGUF
    if output_weight.shape[0] == d_model:
        output_weight = output_weight.T  # [vocab_size, d_model]

    # Get logits for last position
    X_last = X[-1:]  # [1, d_model]
    X_normed = rms_norm(X_last, output_norm)
    logits = X_normed @ output_weight.T  # [1, vocab_size]

    # Sample first token
    if temperature == 0:
        next_token = int(np.argmax(logits[0]))
    else:
        probs = np.exp(logits[0] / temperature)
        probs = probs / probs.sum()
        next_token = int(np.random.choice(len(probs), p=probs))

    output_tokens = [next_token]
    print(f"  First token: {next_token} = '{tokenizer.decode([next_token])}'")

    # Decode loop
    print(f"Decoding (max {max_new_tokens} tokens)...")
    t0_decode = time.perf_counter()

    for step in range(max_new_tokens - 1):
        if next_token == tokenizer.eos_token_id:
            print(f"  EOS at step {step}")
            break

        # Embed single token
        X = embed_weight[next_token:next_token+1]  # [1, d_model]
        pos = len(token_ids) + step

        # Run through all blocks
        for blk in range(num_blocks):
            X, kv = forward_block_with_cache(
                X, blk, start_pos=pos, kv_cache=kv_caches[blk], weights=weights,
                n_heads=n_heads, n_kv_heads=n_kv_heads, d_head=d_head,
            )
            kv_caches[blk] = kv

        # Get logits
        X_normed = rms_norm(X, output_norm)
        logits = X_normed @ output_weight.T  # [1, vocab_size]

        # Sample
        if temperature == 0:
            next_token = int(np.argmax(logits[0]))
        else:
            probs = np.exp(logits[0] / temperature)
            probs = probs / probs.sum()
            next_token = int(np.random.choice(len(probs), p=probs))

        output_tokens.append(next_token)

        if (step + 1) % 4 == 0:
            print(f"  Step {step+1}: '{tokenizer.decode(output_tokens)}'")

    decode_total_ms = (time.perf_counter() - t0_decode) * 1000
    decode_ms_per_token = decode_total_ms / len(output_tokens) if output_tokens else 0

    # Decode output text
    output_text = tokenizer.decode(output_tokens)
    print(f"\nGenerated: {output_text}")

    # Calculate KV cache size
    kv_cache_bytes = sum(kv.size_bytes() for kv in kv_caches.values())
    kv_cache_mb = kv_cache_bytes / (1024 * 1024)

    peak_rss = max(peak_rss, get_rss_mb())

    total_time = (time.perf_counter() - t0_total) * 1000
    print(f"\nTiming:")
    print(f"  Prefill: {prefill_ms:.0f}ms")
    print(f"  Decode: {decode_total_ms:.0f}ms ({decode_ms_per_token:.0f}ms/token)")
    print(f"  Total: {total_time:.0f}ms")
    print(f"  Peak RSS: {peak_rss:.0f}MB")
    print(f"  KV cache: {kv_cache_mb:.1f}MB")

    return GenerateReceipt(
        model=str(gguf_path.name),
        cdna_dir=str(cdna_dir) if cdna_dir else "",
        prompt=prompt,
        output_text=output_text,
        output_tokens=output_tokens,
        tokens_generated=len(output_tokens),
        peak_rss_mb=peak_rss,
        prefill_ms=prefill_ms,
        decode_total_ms=decode_total_ms,
        decode_ms_per_token=decode_ms_per_token,
        backend=backend,
        kv_cache_mb=kv_cache_mb,
        num_blocks=num_blocks,
    )


# =============================================================================
# CLI
# =============================================================================

def main():
    parser = argparse.ArgumentParser(
        description="Minimal streaming generation for TinyLlama"
    )
    parser.add_argument(
        "--model", required=True,
        help="Path to GGUF model file"
    )
    parser.add_argument(
        "--cdna-dir",
        help="Path to CDNA directory (optional, uses GGUF if not provided)"
    )
    parser.add_argument(
        "--prompt", required=True,
        help="Input prompt"
    )
    parser.add_argument(
        "--max-new-tokens", type=int, default=32,
        help="Maximum new tokens to generate"
    )
    parser.add_argument(
        "--temperature", type=float, default=0.0,
        help="Sampling temperature (0 = greedy)"
    )
    parser.add_argument(
        "--receipt-out", default="generate_receipt.json",
        help="Path to save receipt JSON"
    )
    parser.add_argument(
        "--use-cdna", action="store_true",
        help="Use CDNA weights instead of GGUF"
    )

    args = parser.parse_args()

    gguf_path = Path(args.model).expanduser()
    cdna_dir = Path(args.cdna_dir).expanduser() if args.cdna_dir else None

    if not gguf_path.exists():
        print(f"ERROR: Model not found: {gguf_path}")
        return 1

    receipt = generate(
        prompt=args.prompt,
        gguf_path=gguf_path,
        cdna_dir=cdna_dir,
        max_new_tokens=args.max_new_tokens,
        temperature=args.temperature,
        use_cdna=args.use_cdna,
    )

    # Save receipt
    receipt_path = Path(args.receipt_out)
    receipt_path.write_text(json.dumps(asdict(receipt), indent=2))
    print(f"\nReceipt saved: {receipt_path}")

    return 0


if __name__ == "__main__":
    exit(main())
