#!/usr/bin/env python3
"""Export the Apochat-tuned Gemma 4 E2B MLX model to a LiteRT .litertlm artifact.

This script is intentionally **not** run on the 16 GB local Mac. It is designed for a
machine with at least 32 GB of CPU RAM or a GPU with 24 GB+ VRAM (e.g. a Hugging Face
Space/Notebook with GPU upgrade).

Pipeline:
1. Load the public MLX-q4 fused snapshot from Hugging Face.
2. Dequantize weights to bfloat16 and save as PyTorch-format safetensors shards.
3. Patch the config so transformers sees a normal bf16 checkpoint.
4. Run `litert convert` with weight-only int4 quantization to produce .litertlm.
5. Upload the resulting artifact to a Hugging Face model repo.

Usage (on a high-memory machine / HF Space):
    pip install -r scripts/requirements_litert_export.txt
    python scripts/export_apochat_litert.py \
        --mlx-repo apoapps/apochat-gemma4-e2b-apochat-tuned-v1 \
        --output-dir ./apochat-litert-build \
        --upload-repo apoapps/apochat-gemma4-e2b-apochat-tuned-v1-litert
"""

from __future__ import annotations

import argparse
import json
import os
import shutil
import subprocess
import sys
import tempfile
from pathlib import Path
from typing import Any

import mlx.core as mx
import numpy as np
from huggingface_hub import HfApi, create_repo, hf_hub_download, upload_file, upload_folder
from safetensors.torch import save_file


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Export Apochat-tuned Gemma 4 E2B to LiteRT")
    parser.add_argument(
        "--mlx-repo",
        default="apoapps/apochat-gemma4-e2b-apochat-tuned-v1",
        help="Hugging Face repo id containing the fused MLX-q4 model",
    )
    parser.add_argument(
        "--revision",
        default=None,
        help="Optional git revision for the MLX repo",
    )
    parser.add_argument(
        "--output-dir",
        default="./apochat-litert-build",
        help="Local directory for intermediate PyTorch checkpoint and final .litertlm",
    )
    parser.add_argument(
        "--upload-repo",
        default="apoapps/apochat-gemma4-e2b-apochat-tuned-v1-litert",
        help="HF repo id where the final .litertlm will be uploaded",
    )
    parser.add_argument(
        "--upload-private",
        action="store_true",
        help="Make the upload repo private",
    )
    parser.add_argument(
        "--skip-upload",
        action="store_true",
        help="Keep the output local; do not upload to HF",
    )
    parser.add_argument(
        "--prefill-lengths",
        default="256",
        help="LiteRT prefill signature lengths (comma separated)",
    )
    parser.add_argument(
        "--cache-length",
        type=int,
        default=1024,
        help="LiteRT KV-cache length",
    )
    parser.add_argument(
        "--quantize-recipe",
        default="weight_only_wi4_afp32",
        help="LiteRT quantization recipe",
    )
    parser.add_argument(
        "--shard-size",
        type=int,
        default=10_000_000_000,
        help="Target size in bytes per PyTorch safetensors shard",
    )
    return parser.parse_args()


def _write_sharded_index(checkpoint_dir: Path) -> None:
    """Write model.safetensors.index.json from shards in a directory."""
    from safetensors import safe_open

    weight_map: dict[str, str] = {}
    shards = sorted(checkpoint_dir.glob("model-?????-of-?????.safetensors"))
    for shard in shards:
        with safe_open(str(shard), framework="np") as f:
            for key in f.keys():
                weight_map[key] = shard.name
    index = {"metadata": {"total_size": sum(s.stat().st_size for s in shards)}, "weight_map": weight_map}
    (checkpoint_dir / "model.safetensors.index.json").write_text(
        json.dumps(index, indent=2, sort_keys=True), encoding="utf-8"
    )


def download_repo_files(repo_id: str, revision: str | None, local_dir: Path) -> None:
    """Download all non-weight files from the MLX repo into local_dir."""
    print(f"Downloading aux files from {repo_id} ...")
    local_dir.mkdir(parents=True, exist_ok=True)
    api = HfApi()
    files = api.list_repo_files(repo_id, repo_type="model", revision=revision)
    for fname in files:
        if fname.endswith(".safetensors"):
            continue
        if fname.endswith(".safetensors.index.json"):
            # We will regenerate the index if sharding, or use a single file.
            continue
        print(f"  {fname}")
        hf_hub_download(
            repo_id=repo_id,
            filename=fname,
            repo_type="model",
            revision=revision,
            local_dir=str(local_dir),
            local_dir_use_symlinks=False,
        )


def mlx_key_to_pytorch_key(key: str) -> str:
    """Map MLX Gemma 4 key names to PyTorch / transformers key names."""
    if key.startswith("language_model.model."):
        return "model.language_model." + key[len("language_model.model."):]
    if key.startswith("audio_tower."):
        return "model.audio_tower." + key[len("audio_tower."):]
    if key.startswith("vision_tower."):
        return "model.vision_tower." + key[len("vision_tower."):]
    if key.startswith("embed_audio.embedding_projection"):
        return key.replace("embed_audio.embedding_projection", "model.embed_audio.embedding_projection", 1)
    if key.startswith("embed_vision.embedding_projection"):
        return key.replace("embed_vision.embedding_projection", "model.embed_vision.embedding_projection", 1)
    raise ValueError(f"Unexpected MLX key prefix: {key}")


def patch_config_for_pytorch(config_path: Path) -> None:
    """Remove MLX quantization config and ensure torch_dtype is bfloat16."""
    with open(config_path, "r", encoding="utf-8") as f:
        config: dict[str, Any] = json.load(f)

    config.pop("quantization_config", None)
    text_config = config.get("text_config")
    if isinstance(text_config, dict):
        text_config.pop("quantization_config", None)
    config["torch_dtype"] = "bfloat16"

    with open(config_path, "w", encoding="utf-8") as f:
        json.dump(config, f, indent=2)


def dequantize_mlx_to_pytorch(
    mlx_repo: str,
    revision: str | None,
    output_dir: Path,
    shard_size_bytes: int,
) -> None:
    """Load MLX-q4 weights, dequantize, and write PyTorch safetensors shards."""
    print("Loading MLX-q4 weights ...")
    api = HfApi()
    files = api.list_repo_files(mlx_repo, repo_type="model", revision=revision)
    safetensors_files = [f for f in files if f.endswith(".safetensors")]
    weights: dict[str, mx.array] = {}
    for fname in safetensors_files:
        print(f"  {fname}")
        local_path = hf_hub_download(
            repo_id=mlx_repo,
            filename=fname,
            repo_type="model",
            revision=revision,
        )
        part = mx.load(local_path)
        if isinstance(part, dict):
            weights.update(part)
        else:
            raise RuntimeError(f"Unexpected MLX load result for {fname}: {type(part)}")
    print(f"Total tensors: {len(weights)}")

    # Identify quantized triples: weight + scales + biases.
    quantized: set[str] = set()
    for name in list(weights.keys()):
        if name.endswith(".scales"):
            base = name[: -len(".scales")]
            if f"{base}.biases" in weights:
                quantized.add(base)

    print(f"Quantized groups: {len(quantized)}")

    current_shard: dict[str, Any] = {}
    current_shard_bytes = 0
    shard_index = 0

    def flush_shard() -> None:
        nonlocal current_shard, current_shard_bytes, shard_index
        if not current_shard:
            return
        shard_path = output_dir / f"model-{shard_index:05d}-of-?????.safetensors"
        save_file(current_shard, str(shard_path))
        print(f"  Saved {shard_path.name} ({len(current_shard)} tensors, {current_shard_bytes / 1e9:.2f} GB)")
        current_shard = {}
        current_shard_bytes = 0
        shard_index += 1

    for name, arr in weights.items():
        # Skip scale/bias metadata; we'll consume them with the base weight.
        if name.endswith(".scales") or name.endswith(".biases"):
            continue

        # The quantized group is the key prefix without the final `.weight`.
        group_base = name[: -len(".weight")] if name.endswith(".weight") else name
        is_quantized = group_base in quantized

        if is_quantized:
            scales = weights[f"{group_base}.scales"]
            biases = weights[f"{group_base}.biases"]
            # Dequantize to bfloat16 on the MLX device.
            arr = mx.dequantize(arr, scales, biases, group_size=64, bits=4).astype(mx.bfloat16)
        elif arr.dtype != mx.bfloat16:
            arr = arr.astype(mx.bfloat16)

        torch_tensor = mlx_bfloat16_to_torch(arr)
        pytorch_name = mlx_key_to_pytorch_key(name)
        current_shard[pytorch_name] = torch_tensor
        current_shard_bytes += torch_tensor.nbytes

        if current_shard_bytes >= shard_size_bytes:
            flush_shard()

    flush_shard()

    # Rewrite the final shard names with the actual count.
    shards = sorted(output_dir.glob("model-?????-of-?????.safetensors"))
    total = len(shards)
    if total == 1:
        # Transformers / litert expect a single "model.safetensors" for unsharded checkpoints.
        shards[0].rename(output_dir / "model.safetensors")
    else:
        for i, old in enumerate(shards):
            new = old.with_name(f"model-{i:05d}-of-{total:05d}.safetensors")
            old.rename(new)
        # Generate a fresh index so transformers can load the sharded checkpoint.
        _write_sharded_index(output_dir)

    print(f"Wrote {total} safetensors shard(s) to {output_dir}")


def mlx_bfloat16_to_torch(arr: mx.array) -> Any:
    """Convert an MLX bfloat16 array to a contiguous torch bfloat16 tensor."""
    import torch

    # MLX bfloat16 cannot be read directly by numpy; bridge via uint16.
    u16 = np.array(arr.astype(mx.uint16))
    if not u16.flags.c_contiguous:
        u16 = np.ascontiguousarray(u16)
    return torch.from_numpy(u16).view(torch.bfloat16)


def run_litert_convert(
    checkpoint_dir: Path,
    output_dir: Path,
    prefill_lengths: str,
    cache_length: int,
    quantize_recipe: str,
) -> Path:
    """Run `litert convert` on the dequantized checkpoint."""
    print("Running litert convert ...")
    cmd = [
        "litert",
        "convert",
        str(checkpoint_dir),
        "--output",
        str(output_dir),
        "--quantize-recipe",
        quantize_recipe,
        "--prefill-lengths",
        prefill_lengths,
        "--cache-length",
        str(cache_length),
        "--bundle-litert-lm",
    ]
    subprocess.run(cmd, check=True)

    litertlm_files = list(output_dir.glob("*.litertlm"))
    if not litertlm_files:
        raise RuntimeError(f"No .litertlm file found in {output_dir}")
    return litertlm_files[0]


def upload_litert_model(repo_id: str, litertlm_path: Path, private: bool) -> str:
    """Upload the .litertlm file to HF and return the git revision."""
    print(f"Uploading {litertlm_path.name} to {repo_id} ...")
    create_repo(repo_id, repo_type="model", private=private, exist_ok=True)
    upload_file(
        repo_id=repo_id,
        repo_type="model",
        path_in_repo=litertlm_path.name,
        path_or_fileobj=str(litertlm_path),
    )
    # Get the new revision.
    api = HfApi()
    info = api.repo_info(repo_id, repo_type="model")
    print(f"Uploaded. Revision: {info.sha}")
    return info.sha


def main() -> int:
    args = parse_args()
    output_dir = Path(args.output_dir).resolve()
    output_dir.mkdir(parents=True, exist_ok=True)

    # Stage 1: prepare a transformers-compatible checkpoint.
    pytorch_dir = output_dir / "pytorch_checkpoint"
    pytorch_dir.mkdir(parents=True, exist_ok=True)

    download_repo_files(args.mlx_repo, args.revision, pytorch_dir)
    patch_config_for_pytorch(pytorch_dir / "config.json")

    dequantize_mlx_to_pytorch(
        args.mlx_repo,
        args.revision,
        pytorch_dir,
        shard_size_bytes=args.shard_size,
    )

    # Stage 2: convert to LiteRT.
    litert_dir = output_dir / "litert_out"
    litert_dir.mkdir(parents=True, exist_ok=True)
    litertlm_path = run_litert_convert(
        pytorch_dir,
        litert_dir,
        args.prefill_lengths,
        args.cache_length,
        args.quantize_recipe,
    )
    print(f"LiteRT artifact: {litertlm_path}")

    # Stage 3: upload.
    if not args.skip_upload:
        upload_litert_model(args.upload_repo, litertlm_path, args.upload_private)

    return 0


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