#!/usr/bin/env python3
"""
cdna_chat.py — CDNA Chat Backend for Basin API

Pure backend: does not know HTTP, headers, or FastAPI.
Basin stays a thin HTTP wrapper; all CDNA specifics live here.

Usage:
    from api.backends.cdna_chat import run_cdna_chat, CDNAChatRequest

    request = CDNAChatRequest(
        prompt="Hello",
        max_tokens=100,
        seed=42,
        gguf_path="tmp/mistral_fp8combined_canonical.gguf",
        manifest_path="seeds/hybrid_manifest_v2_pilot.json",
    )
    response = run_cdna_chat(request)

Environment:
    CDNA_GGUF           - Path to GGUF file
    CDNA_MANIFEST       - Path to CDNA manifest
    CDNA_RECEIPTS_DIR   - Where to write receipts
    CDNA_DEPTH          - AUTO|CORE|DEEP
    CDNA_RAM_MB         - RAM budget for inference
"""

import os
import sys
import json
import hashlib
import time
from pathlib import Path
from datetime import datetime, timezone
from dataclasses import dataclass, field, asdict
from typing import Optional, Dict, Any, List

import numpy as np

# Add project root to path
sys.path.insert(0, str(Path(__file__).parent.parent.parent))
os.environ["HELIX_ALLOW_FOREIGN_INFERENCE"] = "1"

# Architecture constants (Mistral 7B)
D_MODEL = 4096
D_HEAD = 128
N_HEADS = 32
N_KV_HEADS = 8
N_LAYERS = 32
INTERMEDIATE_SIZE = 14336
VOCAB_SIZE = 32000
NORM_EPS = 1e-5
ROPE_THETA = 1000000.0


# =============================================================================
# Data Classes
# =============================================================================

@dataclass
class CDNAChatRequest:
    """Request for CDNA chat inference."""
    prompt: str
    max_tokens: int = 100
    seed: int = 42
    temperature: float = 0.0  # 0 = deterministic
    top_p: float = 1.0
    gguf_path: str = ""
    manifest_path: str = ""
    depth_mode: str = "AUTO"  # AUTO, CORE, DEEP
    ram_mb: Optional[int] = None
    receipt_dir: str = "receipts/basin"


@dataclass
class CDNAChatResponse:
    """Response from CDNA chat inference."""
    text: str
    tokens_generated: int
    usage: Dict[str, int] = field(default_factory=dict)
    timings: Dict[str, float] = field(default_factory=dict)
    receipt_path: Optional[str] = None
    receipt_sha256: Optional[str] = None
    debug: Optional[Dict[str, Any]] = None
    error: Optional[str] = None


# =============================================================================
# Core Functions
# =============================================================================

def sha256_string(s: str) -> str:
    """SHA256 hash of string."""
    return hashlib.sha256(s.encode()).hexdigest()[:16]


def sha256_file(path: Path) -> str:
    """SHA256 hash of file."""
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(65536), b""):
            h.update(chunk)
    return h.hexdigest()[:16]


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


def silu(x: np.ndarray) -> np.ndarray:
    """SiLU activation: x * sigmoid(x)"""
    return x * (1 / (1 + np.exp(-x)))


def apply_rope_single_head(x: np.ndarray, theta: float = ROPE_THETA) -> np.ndarray:
    """Apply RoPE to a single head [seq, d_head]."""
    seq_len, d_head = x.shape
    dim = d_head // 2

    freqs = 1.0 / (theta ** (np.arange(0, dim, dtype=np.float32) / dim))
    positions = np.arange(seq_len, dtype=np.float32)
    angles = np.outer(positions, freqs)
    freqs_cis = np.exp(1j * angles).astype(np.complex64)

    x_reshaped = x.reshape(seq_len, dim, 2)
    x_complex = x_reshaped[..., 0] + 1j * x_reshaped[..., 1]
    x_rotated = x_complex * freqs_cis
    x_out = np.stack([x_rotated.real, x_rotated.imag], axis=-1)
    return x_out.reshape(seq_len, d_head).astype(x.dtype)


def run_block(X: np.ndarray, block_idx: int, accessor) -> np.ndarray:
    """Run a single transformer block."""
    batch, seq, _ = X.shape

    # Load weights
    attn_norm = accessor.load_tensor(f"blk.{block_idx}.attn_norm.weight")
    ffn_norm = accessor.load_tensor(f"blk.{block_idx}.ffn_norm.weight")
    Wq = accessor.load_tensor(f"blk.{block_idx}.attn_q.weight")
    Wk = accessor.load_tensor(f"blk.{block_idx}.attn_k.weight")
    Wv = accessor.load_tensor(f"blk.{block_idx}.attn_v.weight")
    Wo = accessor.load_tensor(f"blk.{block_idx}.attn_output.weight")
    Wgate = accessor.load_tensor(f"blk.{block_idx}.ffn_gate.weight")
    Wup = accessor.load_tensor(f"blk.{block_idx}.ffn_up.weight")
    Wdown = accessor.load_tensor(f"blk.{block_idx}.ffn_down.weight")

    # Attention
    X_norm = rms_norm(X, attn_norm)
    Q = X_norm @ Wq
    K = X_norm @ Wk
    V = X_norm @ Wv

    # Apply RoPE per-head
    Q_heads = Q.reshape(batch, seq, N_HEADS, D_HEAD)
    K_heads = K.reshape(batch, seq, N_KV_HEADS, D_HEAD)

    Q_rope_heads = np.zeros_like(Q_heads)
    K_rope_heads = np.zeros_like(K_heads)
    for h in range(N_HEADS):
        Q_rope_heads[0, :, h, :] = apply_rope_single_head(Q_heads[0, :, h, :])
    for h in range(N_KV_HEADS):
        K_rope_heads[0, :, h, :] = apply_rope_single_head(K_heads[0, :, h, :])

    # GQA attention
    Qh = Q_rope_heads.transpose(0, 2, 1, 3)
    Kh = K_rope_heads.transpose(0, 2, 1, 3)
    Vh = V.reshape(batch, seq, N_KV_HEADS, D_HEAD).transpose(0, 2, 1, 3)

    scale = 1.0 / np.sqrt(D_HEAD)
    group_size = N_HEADS // N_KV_HEADS
    context_heads = np.zeros((batch, N_HEADS, seq, D_HEAD), dtype=np.float32)

    for hkv in range(N_KV_HEADS):
        h0 = hkv * group_size
        h1 = (hkv + 1) * group_size
        Qg = Qh[:, h0:h1, :, :]
        Kg = Kh[:, hkv:hkv+1, :, :]
        Vg = Vh[:, hkv:hkv+1, :, :]

        scores = np.einsum('bgsd,bkdt->bgst', Qg, Kg.transpose(0, 1, 3, 2)) * scale
        causal_mask = np.triu(np.ones((seq, seq), dtype=np.float32), k=1) * -1e9
        scores = scores + causal_mask
        scores_max = scores.max(axis=-1, keepdims=True)
        exp_scores = np.exp(scores - scores_max)
        weights = exp_scores / exp_scores.sum(axis=-1, keepdims=True)
        ctx = np.einsum('bgst,bktd->bgsd', weights, Vg)
        context_heads[:, h0:h1, :, :] = ctx

    context = context_heads.transpose(0, 2, 1, 3).reshape(batch, seq, N_HEADS * D_HEAD)
    attn_out = context @ Wo
    X_mid = X + attn_out

    # FFN
    X_mid_norm = rms_norm(X_mid, ffn_norm)
    gate = X_mid_norm @ Wgate
    up = X_mid_norm @ Wup
    hidden = silu(gate) * up
    ffn_out = hidden @ Wdown

    return X_mid + ffn_out


def route_depth(prompt: str, depth_mode: str) -> List[int]:
    """
    Route depth based on mode.
    Returns list of block indices to run.
    """
    if depth_mode == "DEEP":
        return list(range(N_LAYERS))  # All 32 blocks
    elif depth_mode == "CORE":
        return list(range(16))  # First 16 blocks
    elif depth_mode == "AUTO":
        # Simple heuristic: longer prompts get more depth
        # TODO: Replace with Se-based routing
        if len(prompt) > 500:
            return list(range(N_LAYERS))
        elif len(prompt) > 100:
            return list(range(24))
        else:
            return list(range(N_LAYERS))  # Default to full for accuracy
    else:
        return list(range(N_LAYERS))


def run_forward(
    X: np.ndarray,
    accessor,
    blocks: List[int],
) -> np.ndarray:
    """Run forward pass through specified blocks."""
    current_X = X.copy()

    for block_idx in blocks:
        current_X = run_block(current_X, block_idx, accessor)

    # Final norm + logits
    output_norm = accessor.load_tensor("output_norm.weight")
    lm_head = accessor.load_tensor("output.weight")

    hidden_normed = rms_norm(current_X, output_norm)
    logits = hidden_normed @ lm_head  # [batch, seq, vocab]

    return logits


def greedy_decode_one(
    logits: np.ndarray,
    temperature: float = 0.0,
    seed: int = 42,
) -> int:
    """Decode one token from logits."""
    last_logits = logits[0, -1, :]  # [vocab]

    if temperature == 0.0:
        # Greedy
        return int(np.argmax(last_logits))
    else:
        # Sample with temperature
        rng = np.random.default_rng(seed)
        probs = np.exp(last_logits / temperature)
        probs = probs / probs.sum()
        return int(rng.choice(len(probs), p=probs))


# =============================================================================
# Main Entry Point
# =============================================================================

def run_cdna_chat(req: CDNAChatRequest) -> CDNAChatResponse:
    """
    Pure backend: does not know HTTP, headers, or FastAPI.
    - chooses blocks via depth routing
    - runs inference through proven pipeline
    - returns text + receipt pointer
    """
    from helix_cdc.regrow.tokenizer_bridge import MistralTokenizerBridge
    from helix_cdc.regrow.tensor_accessor import TensorAccessor

    start_time = time.time()
    timings = {}

    # Resolve paths
    gguf_path = Path(req.gguf_path) if req.gguf_path else Path.home() / "helix-cdc/tmp/mistral_fp8combined_canonical.gguf"

    # Validate
    if not gguf_path.exists():
        return CDNAChatResponse(
            text="",
            tokens_generated=0,
            error=f"GGUF not found: {gguf_path}",
        )

    try:
        # Initialize
        t0 = time.time()
        bridge = MistralTokenizerBridge.get_bridge(gguf_path)
        accessor = TensorAccessor.from_gguf(gguf_path)
        timings["init_ms"] = (time.time() - t0) * 1000

        # Tokenize
        t0 = time.time()
        tokens, _ = bridge.encode(req.prompt, add_bos=True)
        prompt_tokens = len(tokens)
        timings["tokenize_ms"] = (time.time() - t0) * 1000

        # Route depth
        blocks = route_depth(req.prompt, req.depth_mode)

        # Get embeddings
        t0 = time.time()
        embeddings, _ = bridge.lookup_embeddings(tokens)
        X = embeddings.reshape(1, len(tokens), -1).astype(np.float32)
        timings["embed_ms"] = (time.time() - t0) * 1000

        # Generate tokens
        t0 = time.time()
        generated_tokens = []

        for i in range(req.max_tokens):
            # Forward pass
            logits = run_forward(X, accessor, blocks)

            # Decode
            next_token = greedy_decode_one(
                logits,
                temperature=req.temperature,
                seed=req.seed + i,
            )

            # Check for EOS (token 2 for Mistral)
            if next_token == 2:
                break

            generated_tokens.append(next_token)

            # Extend input for next iteration
            next_emb, _ = bridge.lookup_embeddings([next_token])
            next_emb = next_emb.reshape(1, 1, -1).astype(np.float32)
            X = np.concatenate([X, next_emb], axis=1)

        timings["generate_ms"] = (time.time() - t0) * 1000

        # Decode to text
        t0 = time.time()
        output_text, _ = bridge.decode(generated_tokens)
        timings["decode_ms"] = (time.time() - t0) * 1000

        # Usage stats
        usage = {
            "prompt_tokens": prompt_tokens,
            "completion_tokens": len(generated_tokens),
            "total_tokens": prompt_tokens + len(generated_tokens),
        }

        # Write receipt
        receipt_path = None
        receipt_sha256 = None

        if req.receipt_dir:
            receipt_dir = Path(req.receipt_dir)
            receipt_dir.mkdir(parents=True, exist_ok=True)
            timestamp = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S")
            receipt_path = receipt_dir / f"cdna_chat_{timestamp}.json"

            receipt = {
                "schema": "cdna_chat_receipt_v1",
                "timestamp": datetime.now(timezone.utc).isoformat(),
                "request": {
                    "prompt_hash": sha256_string(req.prompt),
                    "prompt_len": len(req.prompt),
                    "max_tokens": req.max_tokens,
                    "seed": req.seed,
                    "temperature": req.temperature,
                    "depth_mode": req.depth_mode,
                },
                "response": {
                    "tokens_generated": len(generated_tokens),
                    "output_hash": sha256_string(output_text),
                },
                "usage": usage,
                "timings": timings,
                "blocks_used": blocks,
                "gguf_path": str(gguf_path),
            }

            with open(receipt_path, "w") as f:
                json.dump(receipt, f, indent=2)

            receipt_sha256 = sha256_file(receipt_path)
            receipt_path = str(receipt_path)

        timings["total_ms"] = (time.time() - start_time) * 1000

        return CDNAChatResponse(
            text=output_text,
            tokens_generated=len(generated_tokens),
            usage=usage,
            timings=timings,
            receipt_path=receipt_path,
            receipt_sha256=receipt_sha256,
            debug={
                "blocks_used": blocks,
                "depth_mode": req.depth_mode,
            },
        )

    except Exception as e:
        return CDNAChatResponse(
            text="",
            tokens_generated=0,
            error=str(e),
            timings={"total_ms": (time.time() - start_time) * 1000},
        )


# =============================================================================
# Startup Validation
# =============================================================================

def validate_cdna_backend(gguf_path: str = None) -> Dict[str, Any]:
    """
    Startup validation for CDNA backend.
    Run on server boot to fail fast if config is broken.
    """
    from helix_cdc.regrow.tensor_accessor import TensorAccessor

    gguf_path = Path(gguf_path) if gguf_path else Path.home() / "helix-cdc/tmp/mistral_fp8combined_canonical.gguf"

    result = {
        "valid": True,
        "gguf_exists": False,
        "gguf_readable": False,
        "accessor_ok": False,
        "errors": [],
    }

    # Check GGUF exists
    if not gguf_path.exists():
        result["valid"] = False
        result["errors"].append(f"GGUF not found: {gguf_path}")
        return result

    result["gguf_exists"] = True

    # Check GGUF readable
    try:
        with open(gguf_path, "rb") as f:
            magic = f.read(4)
        if magic != b"GGUF":
            result["valid"] = False
            result["errors"].append(f"Invalid GGUF magic: {magic.hex()}")
            return result
        result["gguf_readable"] = True
    except Exception as e:
        result["valid"] = False
        result["errors"].append(f"Failed to read GGUF: {e}")
        return result

    # Check accessor works
    try:
        accessor = TensorAccessor.from_gguf(gguf_path)
        # Quick smoke test: load one tensor
        _ = accessor.load_tensor("blk.0.attn_norm.weight")
        result["accessor_ok"] = True
    except Exception as e:
        result["valid"] = False
        result["errors"].append(f"Accessor failed: {e}")

    return result


# =============================================================================
# CLI Test
# =============================================================================

if __name__ == "__main__":
    print("=" * 70)
    print("CDNA Chat Backend Test")
    print("=" * 70)

    # Validate
    print("\n[1] Validating backend...")
    validation = validate_cdna_backend()
    print(f"    Valid: {validation['valid']}")
    if validation["errors"]:
        print(f"    Errors: {validation['errors']}")
        sys.exit(1)

    # Test inference
    print("\n[2] Running test inference...")
    req = CDNAChatRequest(
        prompt="Hello",
        max_tokens=5,
        seed=42,
        depth_mode="DEEP",
    )

    response = run_cdna_chat(req)

    print(f"\n[3] Response:")
    print(f"    Text: {repr(response.text)}")
    print(f"    Tokens: {response.tokens_generated}")
    print(f"    Usage: {response.usage}")
    print(f"    Timings: {response.timings}")
    print(f"    Receipt: {response.receipt_path}")

    if response.error:
        print(f"    Error: {response.error}")
        sys.exit(1)

    print("\n✅ CDNA Chat Backend OK")
