from __future__ import annotations

import argparse
import json
import shutil
import sys
from pathlib import Path
from typing import Dict

import mlx.core as mx
import mlx.nn as nn
from mlx.utils import tree_flatten

REPO_ROOT = Path(__file__).resolve().parents[1]
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))

from modeling_eurobert_mlx import EuroBertConfig, build_model  # noqa: E402


TOKENIZER_FILES = [
    "tokenizer.json",
    "tokenizer_config.json",
    "special_tokens_map.json",
]


def _save_weights(path: Path, model, metadata: Dict[str, str]) -> None:
    params = tree_flatten(model.parameters(), destination={})
    mx.save_safetensors(str(path), params, metadata=metadata)


def _eval_model(model) -> None:
    params = tree_flatten(model.parameters(), destination={})
    if params:
        mx.eval(*params.values())


def _load_source_config(source_dir: Path) -> dict:
    with open(source_dir / "config.json", "r", encoding="utf-8") as f:
        config = json.load(f)
    weights = mx.load(str(source_dir / "model.safetensors"))
    config["num_labels"] = int(weights["classifier.weight"].shape[0])
    config.setdefault("id2label", {str(i): f"LABEL_{i}" for i in range(config["num_labels"])})
    config.setdefault("label2id", {v: int(k) for k, v in config["id2label"].items()})
    return config


def convert(source_dir: Path, output_dir: Path, include_4bit: bool) -> None:
    source_dir = source_dir.resolve()
    output_dir = output_dir.resolve()
    output_dir.mkdir(parents=True, exist_ok=True)

    config = _load_source_config(source_dir)
    with open(output_dir / "config.json", "w", encoding="utf-8") as f:
        json.dump(config, f, indent=2, sort_keys=True)
        f.write("\n")

    for filename in TOKENIZER_FILES:
        shutil.copy2(source_dir / filename, output_dir / filename)
    if (source_dir / "pulpie-orange.png").exists():
        shutil.copy2(source_dir / "pulpie-orange.png", output_dir / "pulpie-orange.png")

    mlx_config = {
        "format": "mlx",
        "source_model": "feyninc/pulpie-orange-small",
        "architecture": "EuroBertForTokenClassification",
        "implementation": "modeling_eurobert_mlx.py",
        "variants": {
            "model-bf16.safetensors": {
                "dtype": "bfloat16",
                "quantization": None,
                "description": "Native 16-bit BF16 MLX weights converted from the source safetensors.",
            }
        },
    }

    model = build_model(EuroBertConfig(**config))
    model.load_weights(str(source_dir / "model.safetensors"))
    _eval_model(model)
    _save_weights(
        output_dir / "model-bf16.safetensors",
        model,
        {
            "format": "mlx",
            "variant": "bf16",
            "source_model": "feyninc/pulpie-orange-small",
        },
    )

    quantized_specs = [("8bit", 8)]
    if include_4bit:
        quantized_specs.append(("4bit", 4))

    for name, bits in quantized_specs:
        q_model = build_model(EuroBertConfig(**config))
        q_model.load_weights(str(source_dir / "model.safetensors"))
        nn.quantize(q_model, group_size=64, bits=bits, mode="affine")
        _eval_model(q_model)
        weight_name = f"model-{name}.safetensors"
        _save_weights(
            output_dir / weight_name,
            q_model,
            {
                "format": "mlx",
                "variant": name,
                "source_model": "feyninc/pulpie-orange-small",
                "quantization": json.dumps(
                    {"bits": bits, "group_size": 64, "mode": "affine"},
                    sort_keys=True,
                ),
            },
        )
        mlx_config["variants"][weight_name] = {
            "dtype": "quantized",
            "quantization": {"bits": bits, "group_size": 64, "mode": "affine"},
            "description": f"MLX affine weight quantized {bits}-bit variant.",
        }

    with open(output_dir / "mlx_config.json", "w", encoding="utf-8") as f:
        json.dump(mlx_config, f, indent=2, sort_keys=True)
        f.write("\n")


def main() -> None:
    parser = argparse.ArgumentParser(description="Convert Pulpie Orange Small to MLX.")
    parser.add_argument("--source-dir", type=Path, default=Path("source_model"))
    parser.add_argument("--output-dir", type=Path, default=Path("hf_out"))
    parser.add_argument("--include-4bit", action="store_true")
    args = parser.parse_args()
    convert(args.source_dir, args.output_dir, include_4bit=args.include_4bit)


if __name__ == "__main__":
    main()