#!/usr/bin/env python3
"""MiniMax M2 block-FP8 -> Kunlun compressed-tensors dynamic W8A8.

CPU only; no model instantiation, calibration, or accelerator allocation.
Checkpoint weight_scale = absmax / 127 (NOT the Kunlun runtime's absmax).
"""
import argparse
from contextlib import ExitStack
import hashlib
import json
import math
from pathlib import Path
import re
import shutil
import struct
import time

VERSION = 1
PROJECTION = re.compile(
    r"model\.layers\.\d+\.(?:self_attn\.(?:q|k|v|o)_proj|"
    r"block_sparse_moe\.experts\.\d+\.w[123])\.weight$"
)
DTYPE_BYTES = {"F8_E4M3": 1, "BF16": 2, "F16": 2, "F32": 4, "I8": 1}


def read_json(path):
    return json.loads(path.read_text())


def write_json(path, value):
    tmp = path.with_name(path.name + ".tmp")
    tmp.write_text(json.dumps(value, indent=2, ensure_ascii=False) + "\n")
    tmp.replace(path)


def read_header(path):
    with path.open("rb") as stream:
        raw = stream.read(8)
        if len(raw) != 8:
            raise ValueError(f"Truncated safetensors: {path}")
        size = struct.unpack("<Q", raw)[0]
        if size > 100 * 1024 * 1024:
            raise ValueError(f"Invalid safetensors header: {path}")
        header = json.loads(stream.read(size))
    header.pop("__metadata__", None)
    return header


def quantization_config():
    return {
        "quant_method": "compressed-tensors",
        "format": "int-quantized",
        "quantization_status": "compressed",
        "config_groups": {"group_0": {
            "targets": ["Linear"],
            "weights": {"num_bits": 8, "type": "int", "symmetric": True,
                        "strategy": "channel", "dynamic": False},
            "input_activations": {"num_bits": 8, "type": "int", "symmetric": True,
                                  "strategy": "token", "dynamic": True},
            "output_activations": None,
        }},
        "ignore": ["lm_head", "re:.*\\.gate$"],
        "kv_cache_scheme": None,
    }


def plan_source(source):
    config = read_json(source / "config.json")
    quant = config.get("quantization_config", {})
    if (config.get("model_type") != "minimax_m2"
            or quant.get("quant_method") != "fp8"
            or quant.get("fmt") != "float8_e4m3fn"
            or quant.get("activation_scheme") != "dynamic"):
        raise ValueError("Expected MiniMax M2 dynamic block FP8 E4M3 checkpoint")
    block = quant.get("weight_block_size")
    if block != [128, 128]:
        raise ValueError(f"Expected FP8 block size [128, 128], got {block}")
    index_path = source / "model.safetensors.index.json"
    index = read_json(index_path)
    mapping = index["weight_map"]
    specs = {}
    source_files = sorted(set(mapping.values()))
    digest = hashlib.sha256()
    for path in (source / "config.json", index_path):
        digest.update(path.read_bytes())
    for filename in source_files:
        path = (source / filename).resolve()
        if not path.is_relative_to(source):
            raise ValueError(f"Shard outside source: {filename}")
        header = read_header(path)
        stat = path.stat()
        digest.update(f"{filename}:{stat.st_size}:{stat.st_mtime_ns}".encode())
        for key, spec in header.items():
            if mapping.get(key) == filename:
                specs[key] = spec
    if set(specs) != set(mapping):
        raise ValueError("Source index contains missing tensors")
    scales = set()
    fp8_count = 0
    output_bytes = 0
    for key, spec in specs.items():
        dtype, shape = spec["dtype"], spec["shape"]
        if dtype == "F8_E4M3":
            if not PROJECTION.fullmatch(key) or len(shape) != 2:
                raise ValueError(f"Unexpected FP8 tensor: {key} {shape}")
            scale_key = key.removesuffix("weight") + "weight_scale_inv"
            expected = [math.ceil(shape[0] / block[0]), math.ceil(shape[1] / block[1])]
            scale_spec = specs.get(scale_key, {})
            if scale_spec.get("shape") != expected or scale_spec.get("dtype") != "F32":
                raise ValueError(f"Missing/invalid block scale: {scale_key}")
            scales.add(scale_key)
            fp8_count += 1
            output_bytes += math.prod(shape) + shape[0] * 4
        elif key.endswith("weight_scale_inv"):
            continue
        else:
            if PROJECTION.fullmatch(key):
                raise ValueError(f"Expected FP8 projection: {key}, got {dtype}")
            if dtype not in DTYPE_BYTES:
                raise ValueError(f"Unsupported source dtype: {key}: {dtype}")
            output_bytes += math.prod(shape) * DTYPE_BYTES[dtype]
    if scales != {key for key in specs if key.endswith("weight_scale_inv")}:
        raise ValueError("Orphan block scales in checkpoint")
    expected_count = config["num_hidden_layers"] * (4 + config["num_local_experts"] * 3)
    if fp8_count != expected_count:
        raise ValueError(f"Expected {expected_count} FP8 projections, found {fp8_count}")
    groups = {}
    for key in sorted(specs):
        if key not in scales:
            groups.setdefault(mapping[key], []).append(key)
    groups = sorted(groups.items())
    return config, mapping, specs, groups, digest.hexdigest(), output_bytes, fp8_count


def quantize_weight(weight, block_scale, block=(128, 128), chunk_rows=1024):
    """Recover W=FP8*block_scale in FP32, then round(W/(row_absmax/127))."""
    import torch
    if weight.ndim != 2 or weight.dtype != torch.float8_e4m3fn:
        raise ValueError("Expected 2-D FP8 E4M3 weight")
    rows, cols = weight.shape
    if tuple(block_scale.shape) != (math.ceil(rows / block[0]), math.ceil(cols / block[1])):
        raise ValueError("Block scale shape mismatch")
    if chunk_rows < 1 or not torch.isfinite(block_scale).all() or (block_scale < 0).any():
        raise ValueError("Invalid chunk size or block scales")
    result = torch.empty((rows, cols), dtype=torch.int8, device="cpu")
    scales = torch.empty((rows, 1), dtype=torch.float32, device="cpu")
    col_blocks = torch.arange(cols, device="cpu") // block[1]
    for start in range(0, rows, chunk_rows):
        end = min(rows, start + chunk_rows)
        row_blocks = torch.arange(start, end, device="cpu") // block[0]
        value = weight[start:end].to(torch.float32)
        value.mul_(block_scale[row_blocks[:, None], col_blocks[None, :]])
        if not torch.isfinite(value).all():
            raise ValueError("Non-finite dequantized FP8 weight")
        scale = value.abs().amax(dim=1, keepdim=True).div_(127.0)
        # A zero row uses scale=1, q=0. Avoid zero scale and underflow.
        scale.clamp_(min=torch.finfo(torch.float32).tiny)
        scale = torch.where(value.abs().amax(dim=1, keepdim=True) == 0, 1.0, scale)
        result[start:end] = value.div_(scale).round_().clamp_(-127, 127).to(torch.int8)
        scales[start:end] = scale
    return result, scales


def output_specs(keys, specs):
    expected = {}
    for key in keys:
        spec = specs[key]
        expected[key] = ("I8" if spec["dtype"] == "F8_E4M3" else spec["dtype"], spec["shape"])
        if spec["dtype"] == "F8_E4M3":
            expected[key.removesuffix("weight") + "weight_scale"] = ("F32", [spec["shape"][0], 1])
    return expected


def validate_shard(path, expected):
    header = read_header(path)
    if set(header) != set(expected):
        raise ValueError(f"Tensor names mismatch: {path}")
    for key, (dtype, shape) in expected.items():
        if header[key]["dtype"] != dtype or header[key]["shape"] != shape:
            raise ValueError(f"Tensor dtype/shape mismatch: {path}: {key}")
    # safe_open validates offsets and file length without reading tensor data.
    from safetensors import safe_open
    with safe_open(str(path), framework="pt", device="cpu"):
        pass


def copy_assets(source, destination):
    excluded = {"config.json", "model.safetensors.index.json", "quantize_config.json",
                "quantization_config.json", "conversion_state.json", "conversion_report.json"}
    for path in source.iterdir():
        if (path.is_file() and path.name not in excluded and not path.name.endswith(".index.json")
                and path.suffix in {".json", ".py", ".jinja", ".txt", ".model", ".tiktoken"}):
            shutil.copy2(path, destination / path.name)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--source", required=True, type=Path)
    parser.add_argument("--output", required=True, type=Path)
    parser.add_argument("--threads", type=int, default=8)
    parser.add_argument("--chunk-rows", type=int, default=1024)
    parser.add_argument("--dry-run", action="store_true", help="Check headers, estimate space; write nothing")
    parser.add_argument("--resume", action="store_true", help="Resume matching conversion_state.json")
    args = parser.parse_args()
    if args.threads < 1 or args.chunk_rows < 1:
        parser.error("threads and chunk-rows must be positive")
    source, output = args.source.resolve(), args.output.resolve()
    if source == output or source in output.parents or output in source.parents:
        parser.error("source and output must be separate, non-nested directories")
    config, mapping, specs, groups, fingerprint, size, count = plan_source(source)
    ancestor = output
    while not ancestor.exists():
        ancestor = ancestor.parent
    free = shutil.disk_usage(ancestor).free
    print(json.dumps({"fp8_weights": count, "output_shards": len(groups),
                      "estimated_output_GiB": round(size / 2**30, 2),
                      "free_disk_GiB": round(free / 2**30, 2)}, indent=2), flush=True)
    if args.dry_run:
        return
    import torch
    from safetensors import safe_open
    from safetensors.torch import save_file
    torch.set_num_threads(args.threads)
    state_path = output / "conversion_state.json"
    identity = {"version": VERSION, "source": str(source), "fingerprint": fingerprint,
                "chunk_rows": args.chunk_rows, "quantization_config": quantization_config()}
    if args.resume:
        state = read_json(state_path)
        if state["identity"] != identity:
            raise ValueError("Source/settings changed; use a new output directory")
    else:
        if output.exists() and any(output.iterdir()):
            raise ValueError("Output must be empty; use --resume for an interrupted conversion")
        state = {"identity": identity, "completed": {}}
    completed_bytes = sum((output / name).stat().st_size for name in state["completed"])
    if free < max(0, size - completed_bytes) + 2 * 2**30:
        raise ValueError("Insufficient output disk space (including 2 GiB reserve)")
    output.mkdir(parents=True, exist_ok=True)
    write_json(state_path, state)
    output_mapping = {}
    total_size = 0
    started = time.monotonic()
    for number, (filename, keys) in enumerate(groups, 1):
        name = f"model-{number:05d}-of-{len(groups):05d}.safetensors"
        path = output / name
        expected = output_specs(keys, specs)
        if name in state["completed"]:
            if path.stat().st_size != state["completed"][name]:
                raise ValueError(f"Completed shard changed: {path}")
            validate_shard(path, expected)
            print(f"[{number}/{len(groups)}] resume: {name}", flush=True)
        else:
            tensors = {}
            with ExitStack() as stack:
                readers = {}

                def get(key):
                    shard = mapping[key]
                    if shard not in readers:
                        readers[shard] = stack.enter_context(safe_open(
                            str(source / shard), framework="pt", device="cpu"))
                    return readers[shard].get_tensor(key)

                with torch.inference_mode():
                    for key in keys:
                        if specs[key]["dtype"] == "F8_E4M3":
                            scale_key = key.removesuffix("weight") + "weight_scale_inv"
                            tensors[key], tensors[key.removesuffix("weight") + "weight_scale"] = (
                                quantize_weight(get(key), get(scale_key), chunk_rows=args.chunk_rows))
                        else:
                            tensors[key] = get(key).contiguous()
                tmp = path.with_name(path.name + ".tmp")
                save_file(tensors, str(tmp), metadata={"format": "pt"})
                validate_shard(tmp, expected)
                tmp.replace(path)
                del tensors
            state["completed"][name] = path.stat().st_size
            write_json(state_path, state)
            print(f"[{number}/{len(groups)}] wrote {name}; elapsed={time.monotonic()-started:.1f}s", flush=True)
        output_mapping.update(dict.fromkeys(expected, name))
        total_size += sum(math.prod(shape) * DTYPE_BYTES[dtype] for dtype, shape in expected.values())
    copy_assets(source, output)
    config["quantization_config"] = quantization_config()
    write_json(output / "config.json", config)
    write_json(output / "model.safetensors.index.json", {
        "metadata": {"total_size": total_size}, "weight_map": output_mapping})
    write_json(output / "conversion_report.json", {
        "source": str(source), "source_fingerprint": fingerprint,
        "fp8_weights_converted": count, "output_tensors": len(output_mapping),
        "tensor_bytes": total_size, "seconds_this_run": time.monotonic() - started,
        "method": "FP8 * block_scale -> FP32 -> symmetric per-output-channel INT8",
        "calibration": "none; dynamic per-token activation quantization at inference",
    })
    print(f"Complete: {output}", flush=True)


if __name__ == "__main__":
    main()