#!/usr/bin/env python3
"""Project Pyrite Phase 2: AntenneFrance Mixtral Streaming Engine Runner.

Runs Mixtral-8x22B (141B) and Mixtral-8x7B (47B) on 16GB Apple Silicon Macs
using NVMe On-Demand Expert Streaming.
RAM is strictly capped at ~2.3 GiB (8x22B) or ~1.2 GiB (8x7B).
"""

import gc
import json
import os
import sys
import time
from pathlib import Path
from typing import Optional

# Add current dir to path
sys.path.insert(0, str(Path(__file__).parent))

import mlx.core as mx
import mlx.nn as nn
from mlx_lm import stream_generate
from mlx_lm.models.mixtral import Model, ModelArgs
from mlx_lm.tokenizer_utils import TokenizerWrapper
from transformers import AutoTokenizer

from expert_registry import ModelSafetensorsIndex
from moe_stream_layer import StreamingMoELayer


def create_streaming_mixtral_model(model_path: Path, max_cache_experts_per_layer: int = 4):
    """Instantiate and wire the Mixtral MoE-Stream model architecture."""
    model_path = Path(model_path)
    print("=" * 70)
    print(f" Initializing AntenneFrance Mixtral Streaming Engine: {model_path.name}")
    print("=" * 70)

    # 1. Load Config
    with open(model_path / "config.json") as f:
        config = json.load(f)

    model_args = ModelArgs.from_dict(config)
    print(
        f"Architecture: {model_args.model_type} ({model_args.num_hidden_layers} layers, "
        f"{model_args.num_local_experts} experts/layer, Top-{model_args.num_experts_per_tok} active)"
    )

    # 2. Build Model Safetensors Registry
    print("Building zero-copy NVMe safetensors registry...")
    registry = ModelSafetensorsIndex(model_path)
    registry.open()

    # 3. Load Resident Weights (skip switch_mlp)
    print("Loading resident weights (Filtering out all switch_mlp routed experts)...")
    resident_weights = {}
    skipped_tensors = 0

    for weight_name in registry.weight_map.keys():
        if "switch_mlp" in weight_name:
            skipped_tensors += 1
            continue
        resident_weights[weight_name] = registry.get_tensor(weight_name)

    print(
        f"Resident tensors: {len(resident_weights):,} | "
        f"Offloaded expert tensors: {skipped_tensors:,} (100% on NVMe SSD, Zero RAM bloat!)"
    )

    # 4. Instantiate Base Model
    print("Instantiating base model without full expert RAM allocation...")
    saved_experts = model_args.num_local_experts
    model_args.num_local_experts = 1  # prevent allocating 8 full experts per layer in memory
    model = Model(model_args)
    model_args.num_local_experts = saved_experts

    # Auto-quantize resident attention layers and gate
    print("Quantizing resident attention, gates, and embeddings...")
    def auto_class_predicate(p, m):
        if "switch_mlp" in p or "block_sparse_moe.switch_mlp" in p:
            return False
        if not hasattr(m, "to_quantized"):
            return False
        w_key = f"{p}.weight"
        s_key = f"{p}.scales"
        if w_key in resident_weights and s_key in resident_weights:
            w = resident_weights[w_key]
            s = resident_weights[s_key]
            in_dim = m.weight.shape[1] if hasattr(m, "weight") else (m.input_dims if hasattr(m, "input_dims") else None)
            if in_dim is None:
                return False
            bits = (w.shape[-1] * 32) // in_dim
            group_size = in_dim // s.shape[-1]
            return {"bits": bits, "group_size": group_size, "mode": "affine"}
        return False

    nn.quantize(model, class_predicate=auto_class_predicate)

    # Load resident weights into base model
    print("Loading resident weights into base model...")
    model.load_weights(list(resident_weights.items()), strict=False)
    del resident_weights
    gc.collect()

    # 5. Inject StreamingMoELayer into each layer
    print(f"Injecting StreamingMoELayers ({max_cache_experts_per_layer} experts LRU cache/layer)...")
    for layer_idx in range(model_args.num_hidden_layers):
        l_prefix = f"model.layers.{layer_idx}.block_sparse_moe"
        gate_w = registry.get_tensor(f"{l_prefix}.gate.weight")
        gate_s = registry.get_tensor(f"{l_prefix}.gate.scales") if f"{l_prefix}.gate.scales" in registry.weight_map else None
        gate_b = registry.get_tensor(f"{l_prefix}.gate.biases") if f"{l_prefix}.gate.biases" in registry.weight_map else None

        stream_moe = StreamingMoELayer(
            layer_idx=layer_idx,
            gate_weight=gate_w,
            gate_scales=gate_s,
            gate_biases=gate_b,
            shared_experts=None,  # Mixtral has no shared experts
            index_reader=registry,
            top_k=model_args.num_experts_per_tok,
            moe_prefix="block_sparse_moe",
            softmax_before_topk=False,  # Mixtral: top-k logits then softmax
            max_cache_experts=max_cache_experts_per_layer,
        )

        model.layers[layer_idx].block_sparse_moe = stream_moe

    print("AntenneFrance Mixtral Streaming Engine successfully mounted on Apple Metal!")
    return model, registry


def main():
    if len(sys.argv) < 2:
        print("Usage: python3 mixtral_runner.py <model_dir_or_hf_repo> [prompt]")
        print("Example: python3 mixtral_runner.py mlx-community/Mixtral-8x22B-Instruct-v0.1-4bit")
        sys.exit(1)

    model_path = Path(sys.argv[1])
    prompt = sys.argv[2] if len(sys.argv) > 2 else "Bonjour ! Pouvez-vous vous présenter et expliquer le Projet Pyrite en français ?"

    # If model_path does not exist locally, prompt user or guide download
    if not model_path.exists():
        print(f"Model path {model_path} not found locally.")
        print(f"To stream from local NVMe, download via huggingface-cli:")
        print(f"  huggingface-cli download {sys.argv[1]} --local-dir /Volumes/NVMe2TB/AI-LLM/models/{Path(sys.argv[1]).name}")
        sys.exit(1)

    # Load Tokenizer
    print("Loading tokenizer...")
    raw_tok = AutoTokenizer.from_pretrained(str(model_path), trust_remote_code=True)
    tokenizer = TokenizerWrapper(raw_tok)

    # Build Streaming Model
    model, registry = create_streaming_mixtral_model(model_path, max_cache_experts_per_layer=4)

    print("\n" + "=" * 70)
    print(f" Prompt: {prompt.strip()}")
    print("=" * 70)
    print("Streaming generation started:\n")

    response_gen = stream_generate(
        model,
        tokenizer,
        prompt=prompt,
        max_tokens=64,
    )

    last_resp = None
    for resp in response_gen:
        print(resp.text, end="", flush=True)
        last_resp = resp

    print("\n\n" + "=" * 70)
    if last_resp is not None:
        print(f" Prompt TPS: {last_resp.prompt_tps:.2f} tok/s ({last_resp.prompt_tokens} tokens)")
        print(f" Generation Speed: {last_resp.generation_tps:.2f} tok/s ({last_resp.generation_tokens} tokens)")
        print(f" Peak Memory: {last_resp.peak_memory:.2f} GB")
    print("=" * 70)

    registry.close()


if __name__ == "__main__":
    main()