"""Convert attention/MLP matrices one block at a time, retaining native precision islands."""

from __future__ import annotations

import argparse
import hashlib
import json
import logging
from pathlib import Path

import mlx.core as mx

from .config import DiTConfig
from .runtime import configure_memory
from .weights import Linear, WeightStore

log = logging.getLogger(__name__)


def convert(model_path, destination, *, bits=8, group_size=64):
    configure_memory()
    model_path, destination = Path(model_path), Path(destination)
    cfg = DiTConfig.load(model_path / "transformer/config.json")
    source = WeightStore(model_path / "transformer")
    if destination.exists() and any(destination.iterdir()):
        raise FileExistsError(f"Conversion destination must be empty: {destination}")
    destination.mkdir(parents=True, exist_ok=True)
    prefixes = [f"blocks.{i}" for i in range(cfg.num_layers)]
    prefixes += [
        f"token_refiner.blocks.{i}" for i in range(cfg.token_refiner_num_layers)
    ]
    metadata = {
        "format": "h3-mlx-core-v1",
        "bits": bits,
        "group_size": group_size,
        "qkv_layout": "qkv-head-dim",
        "source_fingerprint": source.fingerprint(),
        "source_config_sha256": hashlib.sha256(
            (model_path / "transformer/config.json").read_bytes()
        ).hexdigest(),
        "precision_policy": "Only BF16 attention/MLP matrices quantized; AdaLN and FP32 islands remain in the original checkpoint",
    }
    files = []
    for i, prefix in enumerate(prefixes):
        tensors = {}
        for name in (
            "norm1.weight",
            "norm2.weight",
            "attn.q_norm.weight",
            "attn.k_norm.weight",
        ):
            tensors[f"{prefix}.{name}"] = source.get(f"{prefix}.{name}")
        for name in ("attn.qkv_proj", "attn.out_proj", "mlp.fc1", "mlp.fc2"):
            grouped = (
                (cfg.num_attention_heads, cfg.attention_head_dim)
                if name == "attn.qkv_proj"
                else None
            )
            layer = Linear.load(
                source,
                f"{prefix}.{name}",
                bits=bits,
                group_size=group_size,
                grouped_heads=grouped,
            )
            for suffix, value in (
                ("weight", layer.weight),
                ("scales", layer.scales),
                ("biases", layer.qbiases),
                ("bias", layer.bias),
            ):
                if value is not None:
                    tensors[f"{prefix}.{name}.{suffix}"] = value
            del layer
        mx.eval(tensors)
        file = destination / f"core-{i:03d}.safetensors"
        tmp = file.with_suffix(".tmp.safetensors")
        mx.save_safetensors(str(tmp), tensors)
        tmp.replace(file)
        digest = hashlib.sha256()
        with file.open("rb") as handle:
            for data in iter(lambda: handle.read(8 * 1024**2), b""):
                digest.update(data)
        files.append(
            {
                "file": file.name,
                "bytes": file.stat().st_size,
                "sha256": digest.hexdigest(),
            }
        )
        del tensors
        mx.clear_cache()
        log.info("Converted block %d/%d", i + 1, len(prefixes))
    metadata["files"] = files
    (destination / "quantization.json").write_text(
        json.dumps(metadata, indent=2) + "\n"
    )
    return metadata


def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--model", required=True)
    p.add_argument("--output", required=True)
    p.add_argument("--bits", type=int, choices=(4, 8), default=8)
    p.add_argument("--group-size", type=int, choices=(32, 64, 128), default=64)
    a = p.parse_args()
    logging.basicConfig(level=logging.INFO)
    convert(a.model, a.output, bits=a.bits, group_size=a.group_size)


if __name__ == "__main__":
    main()