#!/usr/bin/env python3
"""
Model Runner for helix64 v0.3.0

Unified inference interface for:
  - GGUF models (.gguf) - Direct GGUF loading via gguf library
  - Helix capsules (.hxz) - CDC-compressed regenerative seeds

This module is called by helix64 run --model and produces receipts.

Usage:
    from helix_cdc.drivers.model_runner import ModelRunner

    runner = ModelRunner.from_path("model.gguf", backend="cpu")
    result = runner.generate(prompt="Hello", max_tokens=64)
    receipt = runner.get_receipt()
"""

from __future__ import annotations

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

# Lazy imports to avoid loading torch on import
_torch = None
_np = None


def _lazy_torch():
    global _torch
    if _torch is None:
        import torch
        _torch = torch
    return _torch


def _lazy_np():
    global _np
    if _np is None:
        import numpy as np
        _np = np
    return _np


@dataclass
class ModelConfig:
    """Configuration for model loading."""
    path: Path
    kind: str  # "gguf" or "hxz"
    backend: str = "cpu"  # "cpu" or "cuda"
    device: str = "cpu"  # torch device string
    sha256: str = ""  # Model hash (first 1MB)
    size_bytes: int = 0


@dataclass
class InferenceResult:
    """Result of a generate() call."""
    prompt: str
    output: str
    tokens_generated: int
    elapsed_ms: float
    model_config: ModelConfig
    success: bool = True
    error: Optional[str] = None


class ModelRunner:
    """
    Unified model runner for helix64.

    Supports:
      - GGUF models (via gguf library + transformers)
      - HXZ capsules (via Helix regrow)
    """

    def __init__(self, config: ModelConfig):
        self.config = config
        self.model = None
        self.tokenizer = None
        self.metadata: Dict[str, Any] = {}
        self._load_start_time = 0.0
        self._load_elapsed_ms = 0.0

    @classmethod
    def from_path(
        cls,
        model_path: str,
        backend: str = "cpu",
    ) -> "ModelRunner":
        """
        Create ModelRunner from file path.

        Args:
            model_path: Path to model file (.gguf or .hxz)
            backend: "cpu" or "cuda"

        Returns:
            ModelRunner ready for generate()
        """
        path = Path(model_path)
        if not path.exists():
            raise FileNotFoundError(f"Model not found: {path}")

        # Determine kind
        suffix = path.suffix.lower()
        if suffix == ".gguf":
            kind = "gguf"
        elif suffix == ".hxz":
            kind = "hxz"
        else:
            raise ValueError(f"Unknown model format: {suffix} (expected .gguf or .hxz)")

        # Compute hash (first 1MB)
        with open(path, "rb") as f:
            sha256 = hashlib.sha256(f.read(1024 * 1024)).hexdigest()

        # Determine device
        torch = _lazy_torch()
        if backend == "cuda" and torch.cuda.is_available():
            device = "cuda"
        else:
            device = "cpu"

        config = ModelConfig(
            path=path,
            kind=kind,
            backend=backend,
            device=device,
            sha256=sha256,
            size_bytes=path.stat().st_size,
        )

        runner = cls(config)
        runner._load_model()
        return runner

    def _load_model(self) -> None:
        """Load model based on kind."""
        self._load_start_time = time.time()

        if self.config.kind == "gguf":
            self._load_gguf()
        elif self.config.kind == "hxz":
            self._load_hxz()
        else:
            raise ValueError(f"Unknown model kind: {self.config.kind}")

        self._load_elapsed_ms = (time.time() - self._load_start_time) * 1000

    def _load_gguf(self) -> None:
        """Load GGUF model using transformers + gguf library."""
        torch = _lazy_torch()

        # Check if gguf library is available
        try:
            from gguf import GGUFReader
        except ImportError:
            raise ImportError(
                "GGUF support requires the 'gguf' package.\n"
                "Install with: pip install gguf"
            )

        # Use our GGUF loader from helix_cdc.loaders
        try:
            from helix_cdc.loaders.tinyllama_genome0 import gguf_to_pytorch_state_dict
        except ImportError:
            # Fallback if loader not available
            raise ImportError(
                "GGUF loader not found. Expected helix_cdc.loaders.tinyllama_genome0"
            )

        # Load state dict
        print(f"[ModelRunner] Loading GGUF: {self.config.path.name}")
        state_dict, mapping_info = gguf_to_pytorch_state_dict(str(self.config.path))

        # Infer model type from tensor names
        model_type = self._infer_model_type(state_dict)
        print(f"[ModelRunner] Detected model type: {model_type}")

        # Create model architecture
        if model_type == "llama":
            self._create_llama_model(state_dict)
        else:
            raise ValueError(f"Unsupported GGUF model type: {model_type}")

        self.metadata.update({
            "source": "GGUF",
            "tensors_loaded": len(state_dict),
            "mapping_info": {
                "mapped": len(mapping_info.get("mapped", {})),
                "unmapped": len(mapping_info.get("unmapped", [])),
            },
        })

    def _infer_model_type(self, state_dict: Dict[str, Any]) -> str:
        """Infer model type from state dict keys."""
        keys = list(state_dict.keys())

        # Check for LLaMA-style keys
        if any("model.layers" in k for k in keys):
            return "llama"
        if any("token_embd" in k or "blk." in k for k in keys):
            return "llama"  # GGUF format

        return "unknown"

    def _create_llama_model(self, state_dict: Dict[str, Any]) -> None:
        """Create LlamaForCausalLM and load state dict."""
        torch = _lazy_torch()

        # Try to import transformers
        try:
            from transformers import AutoConfig, AutoTokenizer, LlamaForCausalLM
        except ImportError:
            raise ImportError(
                "Transformers library required for LLaMA models.\n"
                "Install with: pip install transformers"
            )

        # Infer config from state dict shapes
        # Look for embed_tokens or equivalent
        embed_key = None
        for k in state_dict:
            if "embed_tokens" in k or "token_embd" in k:
                embed_key = k
                break

        if embed_key is None:
            raise ValueError("Cannot find embedding tensor to infer vocab size")

        embed_shape = state_dict[embed_key].shape
        vocab_size = embed_shape[0]
        hidden_size = embed_shape[1]

        # Count layers
        layer_indices = set()
        for k in state_dict:
            import re
            m = re.search(r"layers\.(\d+)\.", k)
            if m:
                layer_indices.add(int(m.group(1)))
        num_layers = len(layer_indices) if layer_indices else 22  # TinyLlama default

        print(f"[ModelRunner] Inferred config: vocab={vocab_size}, hidden={hidden_size}, layers={num_layers}")

        # Try TinyLlama config first (most common for small models)
        try:
            config = AutoConfig.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0")
            self.tokenizer = AutoTokenizer.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0")
        except Exception as e:
            print(f"[ModelRunner] Warning: Could not load TinyLlama config: {e}")
            # Create minimal config
            from transformers import LlamaConfig
            config = LlamaConfig(
                vocab_size=vocab_size,
                hidden_size=hidden_size,
                num_hidden_layers=num_layers,
            )
            self.tokenizer = None

        # Create model
        self.model = LlamaForCausalLM(config)

        # Load weights
        missing, unexpected = self.model.load_state_dict(state_dict, strict=False)

        if missing:
            print(f"[ModelRunner] Missing keys: {len(missing)}")
        if unexpected:
            print(f"[ModelRunner] Unexpected keys: {len(unexpected)}")

        # Move to device
        self.model = self.model.to(device=self.config.device)
        self.model.eval()

        # Freeze
        for p in self.model.parameters():
            p.requires_grad = False

        self.metadata["model_type"] = "llama"
        self.metadata["vocab_size"] = vocab_size
        self.metadata["hidden_size"] = hidden_size
        self.metadata["num_layers"] = num_layers

    def _load_hxz(self) -> None:
        """Load Helix capsule (.hxz) via regrow pipeline."""
        torch = _lazy_torch()

        print(f"[ModelRunner] Loading HXZ capsule: {self.config.path.name}")

        # Use the HXZ model loader
        try:
            from helix_cdc.regrow.hxz_model_loader import (
                hxz_to_pytorch_state_dict,
                inspect_hxz_header,
            )
        except ImportError as e:
            print(f"[ModelRunner] HXZ loader not available: {e}")
            self.model = None
            self.tokenizer = None
            self.metadata.update({
                "source": "HXZ_CAPSULE",
                "error": f"HXZ loader not available: {e}",
            })
            return

        # Inspect header first
        header = inspect_hxz_header(str(self.config.path))
        print(f"[ModelRunner] HXZ format: {header.format}, schema: {header.schema}")

        # Load state dict
        try:
            state_dict, mapping_info = hxz_to_pytorch_state_dict(str(self.config.path))
        except ValueError as e:
            print(f"[ModelRunner] HXZ load error: {e}")
            self.model = None
            self.tokenizer = None
            self.metadata.update({
                "source": "HXZ_CAPSULE",
                "format": header.format,
                "schema": header.schema,
                "error": str(e),
            })
            return

        # Convert numpy arrays to torch tensors
        torch_state_dict = {}
        for name, arr in state_dict.items():
            torch_state_dict[name] = torch.from_numpy(arr)

        # Infer model type from state dict keys
        model_type = self._infer_hxz_model_type(torch_state_dict)
        print(f"[ModelRunner] Detected HXZ model type: {model_type}")

        # Create appropriate model
        if model_type == "lstm":
            self._create_lstm_model(torch_state_dict)
        elif model_type == "mistral_7b":
            self._create_mistral_7b_from_vault(torch_state_dict)
        elif model_type == "llama":
            self._create_llama_model({k: v.numpy() for k, v in torch_state_dict.items()})
        else:
            # Store state dict for manual use
            self.model = None
            self.tokenizer = None
            print(f"[ModelRunner] Unknown HXZ model type '{model_type}' - state dict available")

        hxz_meta = mapping_info.get("hxz_metadata", {})
        self.metadata.update({
            "source": "HXZ_CAPSULE",
            "format": header.format,
            "schema": header.schema,
            "tensors_loaded": len(state_dict),
            "model_type": model_type,
            "hxz_created_utc": hxz_meta.get("created_utc"),
        })

    def _infer_hxz_model_type(self, state_dict: Dict[str, Any]) -> str:
        """Infer model type from HXZ state dict keys."""
        keys = list(state_dict.keys())

        # LSTM model
        if any("lstm" in k.lower() for k in keys):
            return "lstm"

        # Check for Mistral-7B specifically (291 tensors with GQA structure)
        if len(keys) >= 280 and any("model.layers" in k for k in keys):
            # Count k_proj tensors - if present with shape [1024, 4096], it's Mistral GQA
            return "mistral_7b"

        # LLaMA-style model
        if any("model.layers" in k for k in keys):
            return "llama"
        if any("embed_tokens" in k or "token_embd" in k for k in keys):
            return "llama"

        # GPT-style model
        if any("wte" in k or "wpe" in k for k in keys):
            return "gpt"

        return "unknown"

    def _create_lstm_model(self, state_dict: Dict[str, Any]) -> None:
        """Create LSTM model from HXZ state dict."""
        torch = _lazy_torch()
        import torch.nn as nn

        # Infer LSTM config from state dict shapes
        # Look for lstm.weight_ih_l0 or similar
        hidden_size = None
        input_size = None
        num_layers = 0
        vocab_size = None

        for name, tensor in state_dict.items():
            if "lstm.weight_ih_l" in name:
                # weight_ih shape: (4*hidden_size, input_size)
                hidden_size = tensor.shape[0] // 4
                if "l0" in name:
                    input_size = tensor.shape[1]
                num_layers = max(num_layers, int(name.split("_l")[-1]) + 1)
            elif "emb.weight" in name:
                vocab_size = tensor.shape[0]
                input_size = tensor.shape[1]

        if hidden_size is None or input_size is None:
            print(f"[ModelRunner] Cannot infer LSTM config from state dict")
            self.model = None
            return

        print(f"[ModelRunner] LSTM config: vocab={vocab_size}, input={input_size}, hidden={hidden_size}, layers={num_layers}")

        # Create model
        class SimpleLSTM(nn.Module):
            def __init__(self, vocab_size, embed_dim, hidden_dim, num_layers):
                super().__init__()
                self.emb = nn.Embedding(vocab_size, embed_dim)
                self.lstm = nn.LSTM(embed_dim, hidden_dim, num_layers, batch_first=True)
                self.head = nn.Linear(hidden_dim, vocab_size)

            def forward(self, x):
                x = self.emb(x)
                x, _ = self.lstm(x)
                return self.head(x)

        model = SimpleLSTM(vocab_size or 256, input_size, hidden_size, num_layers)

        # Load weights
        missing, unexpected = model.load_state_dict(state_dict, strict=False)
        if missing:
            print(f"[ModelRunner] Missing keys: {len(missing)}")
        if unexpected:
            print(f"[ModelRunner] Unexpected keys: {len(unexpected)}")

        model = model.to(self.config.device)
        model.eval()
        for p in model.parameters():
            p.requires_grad = False

        self.model = model
        self.tokenizer = None  # LSTM doesn't have a standard tokenizer

        self.metadata["model_type"] = "lstm"
        self.metadata["vocab_size"] = vocab_size or 256
        self.metadata["hidden_size"] = hidden_size
        self.metadata["num_layers"] = num_layers

    def _create_mistral_7b_from_vault(self, state_dict: Dict[str, Any]) -> None:
        """
        Create Mistral-7B LlamaForCausalLM from DNA-expanded vault state dict.

        Uses exact Mistral-7B architecture:
        - vocab_size: 32000
        - hidden_size: 4096
        - intermediate_size: 14336
        - num_hidden_layers: 32
        - num_attention_heads: 32
        - num_key_value_heads: 8 (GQA)
        """
        torch = _lazy_torch()

        try:
            from transformers import LlamaConfig, LlamaForCausalLM, AutoTokenizer
        except ImportError:
            raise ImportError(
                "Transformers library required for Mistral-7B.\n"
                "Install with: pip install transformers"
            )

        print("[ModelRunner] Creating Mistral-7B from vault...")

        # Exact Mistral-7B-Instruct-v0.2 config
        config = LlamaConfig(
            vocab_size=32000,
            hidden_size=4096,
            intermediate_size=14336,
            num_hidden_layers=32,
            num_attention_heads=32,
            num_key_value_heads=8,  # GQA
            max_position_embeddings=32768,
            rms_norm_eps=1e-5,
            rope_theta=10000.0,
            torch_dtype=torch.float32,
        )

        # Create model
        print(f"[ModelRunner] Building Mistral-7B architecture: {config.num_hidden_layers} layers, {config.hidden_size} hidden")
        self.model = LlamaForCausalLM(config)

        # Convert numpy arrays to torch tensors
        torch_sd = {}
        for name, arr in state_dict.items():
            if isinstance(arr, torch.Tensor):
                torch_sd[name] = arr
            else:
                torch_sd[name] = torch.from_numpy(arr)

        # Load weights
        missing, unexpected = self.model.load_state_dict(torch_sd, strict=False)

        print(f"[ModelRunner] Loaded {len(torch_sd)} tensors")
        if missing:
            print(f"[ModelRunner] Missing keys: {len(missing)}")
            if len(missing) <= 5:
                for k in missing:
                    print(f"  - {k}")
        if unexpected:
            print(f"[ModelRunner] Unexpected keys: {len(unexpected)}")

        # Try to load tokenizer
        try:
            self.tokenizer = AutoTokenizer.from_pretrained(
                "mistralai/Mistral-7B-Instruct-v0.2",
                trust_remote_code=True
            )
            print("[ModelRunner] Loaded Mistral tokenizer")
        except Exception as e:
            print(f"[ModelRunner] Warning: Could not load Mistral tokenizer: {e}")
            # Try fallback to any LlamaTokenizer
            try:
                self.tokenizer = AutoTokenizer.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0")
                print("[ModelRunner] Using TinyLlama tokenizer as fallback")
            except Exception:
                self.tokenizer = None
                print("[ModelRunner] No tokenizer available")

        # Move to device
        self.model = self.model.to(device=self.config.device)
        self.model.eval()

        # Freeze
        for p in self.model.parameters():
            p.requires_grad = False

        self.metadata["model_type"] = "mistral_7b"
        self.metadata["vocab_size"] = 32000
        self.metadata["hidden_size"] = 4096
        self.metadata["intermediate_size"] = 14336
        self.metadata["num_layers"] = 32
        self.metadata["num_attention_heads"] = 32
        self.metadata["num_kv_heads"] = 8
        self.metadata["dna_regenerated"] = True

        # Calculate total params
        total_params = sum(p.numel() for p in self.model.parameters())
        self.metadata["total_parameters"] = total_params
        print(f"[ModelRunner] Mistral-7B ready: {total_params:,} parameters")

    def generate(
        self,
        prompt: str,
        max_tokens: int = 64,
        temperature: float = 0.7,
        do_sample: bool = True,
    ) -> InferenceResult:
        """
        Generate text from prompt.

        Args:
            prompt: Input text
            max_tokens: Maximum new tokens to generate
            temperature: Sampling temperature
            do_sample: Whether to sample (False = greedy)

        Returns:
            InferenceResult with output text and metrics
        """
        if self.model is None:
            # Model not loaded
            model_type = self.metadata.get("model_type", "unknown")
            return InferenceResult(
                prompt=prompt,
                output=f"[Model not loaded (type: {model_type})]",
                tokens_generated=0,
                elapsed_ms=0.0,
                model_config=self.config,
                success=False,
                error=f"Model not loaded (type: {model_type})",
            )

        # Check if LSTM model (uses different generation path)
        model_type = self.metadata.get("model_type", "")
        if model_type == "lstm":
            return self._generate_lstm(prompt, max_tokens, temperature, do_sample)

        if self.tokenizer is None:
            return InferenceResult(
                prompt=prompt,
                output="[No tokenizer available for LLM generation]",
                tokens_generated=0,
                elapsed_ms=0.0,
                model_config=self.config,
                success=False,
                error="Tokenizer not loaded",
            )

        torch = _lazy_torch()

        start_time = time.time()

        try:
            # Tokenize
            inputs = self.tokenizer(prompt, return_tensors="pt")
            input_ids = inputs["input_ids"].to(self.config.device)
            input_len = input_ids.shape[1]

            # Generate
            with torch.no_grad():
                outputs = self.model.generate(
                    input_ids,
                    max_new_tokens=max_tokens,
                    temperature=temperature if do_sample else 1.0,
                    do_sample=do_sample,
                    pad_token_id=self.tokenizer.eos_token_id,
                )

            # Decode
            full_output = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
            # Get only generated part
            tokens_generated = outputs.shape[1] - input_len

            elapsed_ms = (time.time() - start_time) * 1000

            return InferenceResult(
                prompt=prompt,
                output=full_output,
                tokens_generated=tokens_generated,
                elapsed_ms=elapsed_ms,
                model_config=self.config,
                success=True,
            )

        except Exception as e:
            elapsed_ms = (time.time() - start_time) * 1000
            return InferenceResult(
                prompt=prompt,
                output="",
                tokens_generated=0,
                elapsed_ms=elapsed_ms,
                model_config=self.config,
                success=False,
                error=str(e),
            )

    def _generate_lstm(
        self,
        prompt: str,
        max_tokens: int,
        temperature: float,
        do_sample: bool,
    ) -> InferenceResult:
        """Generate text using LSTM model (char-level autoregressive)."""
        torch = _lazy_torch()
        import torch.nn.functional as F

        start_time = time.time()

        try:
            vocab_size = self.metadata.get("vocab_size", 256)

            # Simple char-level encoding (treat each byte as token)
            input_bytes = prompt.encode("utf-8", errors="replace")
            input_ids = torch.tensor(
                [min(b, vocab_size - 1) for b in input_bytes],
                dtype=torch.long
            ).unsqueeze(0).to(self.config.device)

            # Autoregressive generation
            generated_ids = []
            current_input = input_ids

            with torch.no_grad():
                for _ in range(max_tokens):
                    logits = self.model(current_input)
                    # Get last position's logits
                    last_logits = logits[0, -1, :]

                    if do_sample and temperature > 0:
                        # Sample with temperature
                        probs = F.softmax(last_logits / temperature, dim=-1)
                        next_id = torch.multinomial(probs, 1)
                    else:
                        # Greedy
                        next_id = torch.argmax(last_logits, dim=-1, keepdim=True)

                    generated_ids.append(next_id.item())
                    # Append to input
                    current_input = torch.cat([current_input, next_id.unsqueeze(0)], dim=1)

            # Decode generated bytes
            output_bytes = bytes([min(max(0, g), 255) for g in generated_ids])
            output_text = prompt + output_bytes.decode("utf-8", errors="replace")

            elapsed_ms = (time.time() - start_time) * 1000

            return InferenceResult(
                prompt=prompt,
                output=output_text,
                tokens_generated=len(generated_ids),
                elapsed_ms=elapsed_ms,
                model_config=self.config,
                success=True,
            )

        except Exception as e:
            elapsed_ms = (time.time() - start_time) * 1000
            return InferenceResult(
                prompt=prompt,
                output="",
                tokens_generated=0,
                elapsed_ms=elapsed_ms,
                model_config=self.config,
                success=False,
                error=str(e),
            )

    def get_receipt(self, result: Optional[InferenceResult] = None) -> Dict[str, Any]:
        """
        Generate receipt for this runner's operations.

        Args:
            result: Optional InferenceResult to include

        Returns:
            Receipt dict suitable for JSON serialization
        """
        receipt = {
            "protocol_version": "helix_cdc:v0.1.0",
            "schema_version": "helix64_model_runner:v1",
            "model": {
                "path": str(self.config.path),
                "kind": self.config.kind,
                "sha256_1mb": self.config.sha256,
                "size_bytes": self.config.size_bytes,
            },
            "runtime": {
                "backend": self.config.backend,
                "device": self.config.device,
                "load_time_ms": self._load_elapsed_ms,
            },
            "metadata": self.metadata,
            "provenance": {
                "timestamp_utc": datetime.now(timezone.utc).isoformat(),
            },
        }

        if result:
            receipt["inference"] = {
                "prompt": result.prompt,
                "output": result.output,
                "tokens_generated": result.tokens_generated,
                "elapsed_ms": result.elapsed_ms,
                "success": result.success,
            }
            if result.error:
                receipt["inference"]["error"] = result.error

        return receipt


# Entry point for testing
if __name__ == "__main__":
    import sys

    if len(sys.argv) < 2:
        print("Usage: python -m helix_cdc.drivers.model_runner <model.gguf|model.hxz>")
        sys.exit(1)

    model_path = sys.argv[1]
    prompt = sys.argv[2] if len(sys.argv) > 2 else "Hello, my name is"

    print(f"Loading: {model_path}")
    runner = ModelRunner.from_path(model_path)

    print(f"\nGenerating from prompt: '{prompt}'")
    result = runner.generate(prompt=prompt, max_tokens=32)

    print(f"\nResult:")
    print(f"  Success: {result.success}")
    print(f"  Output: {result.output}")
    print(f"  Tokens: {result.tokens_generated}")
    print(f"  Time: {result.elapsed_ms:.1f}ms")

    if result.error:
        print(f"  Error: {result.error}")

    # Print receipt
    receipt = runner.get_receipt(result)
    print(f"\nReceipt:")
    print(json.dumps(receipt, indent=2))
