"""Quantize a Hugging Face model directory to NVFP4 Hugging Face format.

Reads a Hugging Face model directory (config.json + safetensors) and writes a
quantized Hugging Face model directory. By default (--format hf), writes standard
Hugging Face compressed-tensors format with .weight_packed, .weight_scale, and
.weight_global_scale natively supported by Hugging Face transformers, vLLM, SGLang,
and llama.cpp.

Optionally (--format modelopt / comfy), outputs raw weight tensors (.weight),
weight_scale, weight_scale_2, and .comfy_quant for ModelOpt and ComfyUI.

Weight matrices are quantized using convert_to_quant's learned rounding (or simple
round-to-grid with --simple). Biases are adjusted using calibration data unless
--no-bias-correction is specified. Embeddings, layer norms, and excluded layers
stay in their float format.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import logging
import os
import re
import shutil
import sys
from pathlib import Path

import numpy as np
import torch
from safetensors import safe_open
from safetensors.torch import save_file

from nvfp4gguf import (
    CTQ_NUM_ITER,
    DEFAULT_CONVERT_TO_QUANT,
    DEFAULT_LLAMA_CPP,
    DEFAULT_EARLY_STOP_FLOOR_STALL,
    DEFAULT_TARGET_LOSS,
    DEFAULT_EARLY_STOP_STALL,
    DEFAULT_LR_FLOOR,
    DEFAULT_LR_PATIENCE,
    DEFAULT_LR_THRESHOLD,
    DEFAULT_SCALE_OPTIMIZATION,
    DEFAULT_SCALE_REFINEMENT_ROUNDS,
    QUANT_VERSION,
    TensorQuantizer,
    add_quant_arguments,
    bind_exclude,
    coordinate_descent_refine,
    ctq_skips_key,
    ensure_cuda_interpreter,
    load_learned_converter,
    quant_profile,
    report_reconstruction,
    resolve_optimizer,
    sweep_per_tensor_scale,
    ue4m3_bits,
)

INDEX_NAME = "model.safetensors.index.json"
NON_WEIGHT_PATTERNS = {
    ".safetensors",
    ".bin",
    ".pt",
    ".pth",
    ".h5",
    ".msgpack",
    ".ot",
    ".tmp",
}

log = logging.getLogger("hf_to_nvfp4_hf")


def _checkpoint_file(directory: Path, name: str) -> Path:
    digest = hashlib.sha1(name.encode("utf-8")).hexdigest()[:16]
    safe = re.sub(r"[^A-Za-z0-9._-]+", "_", name)[:80]
    return directory / f"{safe}.{digest}.npz"


def _load_checkpoint(directory: Path, name: str, shape: tuple[int, ...], need_bias: bool):
    path = _checkpoint_file(directory, name)
    if not path.is_file():
        return None
    with np.load(path) as blob:
        if "version" in blob.files and int(blob["version"]) != QUANT_VERSION:
            return None
        if tuple(int(v) for v in blob["shape"]) != tuple(shape):
            return None
        if need_bias and "bias" not in blob.files:
            return None
        codes = np.array(blob["codes"], copy=True)
        block_scale = np.array(blob["block_scale"], copy=True)
        per_tensor = float(blob["per_tensor"])
        bias = np.array(blob["bias"], copy=True) if "bias" in blob.files else None
    return codes, block_scale, per_tensor, bias


def _save_checkpoint(directory: Path, name: str, shape: tuple[int, ...], codes, block_scale, per_tensor, bias) -> None:
    path = _checkpoint_file(directory, name)
    temporary = path.with_name(path.name + ".tmp")
    with temporary.open("wb") as handle:
        save_dict = {
            "version": np.int32(QUANT_VERSION),
            "shape": np.asarray(shape, dtype=np.int64),
            "codes": np.asarray(codes, dtype=np.uint8),
            "block_scale": np.asarray(block_scale, dtype=np.float32),
            "per_tensor": np.float32(per_tensor),
        }
        if bias is not None:
            save_dict["bias"] = np.asarray(bias, dtype=np.float32)
        np.savez_compressed(handle, **save_dict)
    temporary.replace(path)


def discover_shards(model_dir: Path) -> list[str]:
    """Return safetensors shard filenames in the checkpoint directory."""
    index_path = model_dir / INDEX_NAME
    if index_path.is_file():
        try:
            payload = json.loads(index_path.read_text(encoding="utf-8"))
            weight_map = payload.get("weight_map")
            if isinstance(weight_map, dict) and weight_map:
                names = sorted(set(weight_map.values()))
                missing = [name for name in names if not (model_dir / name).is_file()]
                if not missing:
                    return names
        except Exception:
            pass

    single = model_dir / "model.safetensors"
    if single.is_file():
        return ["model.safetensors"]

    shards = sorted(path.name for path in model_dir.glob("*.safetensors") if path.is_file())
    if not shards:
        raise SystemExit(f"No safetensors shards found in {model_dir}")
    return shards


def _assign_output(args) -> None:
    source = Path(args.model).resolve()
    stem = source.name
    passthrough = args.no_quant or args.outtype == "nvfp4"
    if not args.output:
        if passthrough and args.outtype == "auto":
            directory = Path.cwd() / stem
        elif passthrough:
            label = "NVFP4" if args.outtype == "nvfp4" else args.outtype.upper()
            directory = Path.cwd() / f"{stem}-{label}"
        else:
            directory = Path.cwd() / f"{stem}-NVFP4"
    else:
        directory = Path(args.output).resolve()
    args.output = str(directory)


def _apply_directory_profile(args, source: Path) -> None:
    config_path = source / "config.json"
    if not config_path.is_file():
        return
    try:
        config = json.loads(config_path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError):
        return
    names: list[str] = []
    architectures = config.get("architectures")
    if isinstance(architectures, list):
        names.extend(str(item) for item in architectures)
    elif isinstance(architectures, str):
        names.append(architectures)
    model_type = config.get("model_type")
    if isinstance(model_type, str) and model_type:
        names.append(model_type)
    for name in names:
        profile = quant_profile(name)
        if profile is None:
            continue
        if not args.exclude:
            args.exclude = profile.skip
            log.info("%s profile: leaving %s unquantized.", profile.arch, profile.label)
        else:
            log.info("%s profile: --exclude replaces default skip of %s.", profile.arch, profile.label)
        return


def _copy_non_weight_files(source_dir: Path, output_dir: Path) -> None:
    """Copy configs, tokenizers, chat templates, and documentation into output."""
    output_dir.mkdir(parents=True, exist_ok=True)
    for path in source_dir.iterdir():
        if path.is_dir() or path.name.startswith("."):
            continue
        if path.name == "config.json" or path.name == INDEX_NAME:
            continue
        if any(path.name.endswith(ext) for ext in NON_WEIGHT_PATTERNS):
            continue
        dest = output_dir / path.name
        if not dest.exists() or dest.stat().st_mtime < path.stat().st_mtime:
            shutil.copy2(path, dest)


def _write_config_json(source_dir: Path, output_dir: Path, args, is_quantized: bool) -> None:
    config_path = source_dir / "config.json"
    if not config_path.is_file():
        return
    try:
        config = json.loads(config_path.read_text(encoding="utf-8"))
    except Exception as exc:
        log.warning("Could not parse config.json: %s", exc)
        return

    fmt = getattr(args, "format", "hf")

    def _build_ct_quant_config():
        ignore = []
        if not getattr(args, "quantize_output", False):
            ignore.append("lm_head")
        if getattr(args, "exclude", None):
            ignore.append(args.exclude)
        return {
            "quant_method": "compressed-tensors",
            "format": "nvfp4-pack-quantized",
            "quantization_status": "compressed",
            "ignore": ignore,
            "config_groups": {
                "group_0": {
                    "weights": {
                        "num_bits": 4,
                        "type": "float",
                        "strategy": "tensor_group",
                        "group_size": 16,
                        "symmetric": True,
                        "dynamic": False,
                        "actorder": None,
                    },
                    "input_activations": None,
                    "targets": ["Linear"],
                    "format": "nvfp4-pack-quantized",
                }
            },
        }

    if is_quantized:
        if fmt == "hf":
            config["quantization_config"] = _build_ct_quant_config()
        else:
            config["quantization_config"] = {
                "quant_method": "modelopt",
                "quant_algo": "NVFP4",
                "format": "nvfp4-pack-quantized",
                "group_size": 16,
                "bits": 4,
            }
        config["torch_dtype"] = getattr(args, "output_dtype", "bfloat16")
    elif args.no_quant:
        if args.outtype in ("f32", "f16", "bf16"):
            config.pop("quantization_config", None)
            config["torch_dtype"] = "bfloat16" if args.outtype == "bf16" else ("float16" if args.outtype == "f16" else "float32")
        elif args.outtype == "nvfp4":
            if fmt == "hf":
                config["quantization_config"] = _build_ct_quant_config()
            else:
                config.setdefault("quantization_config", {
                    "quant_method": "modelopt",
                    "quant_algo": "NVFP4",
                    "format": "nvfp4-pack-quantized",
                    "group_size": 16,
                    "bits": 4,
                })

    if getattr(args, "context_length", None):
        for key in ("max_position_embeddings", "seq_length", "max_sequence_length"):
            if key in config:
                config[key] = int(args.context_length)

    out_file = output_dir / "config.json"
    out_file.write_text(json.dumps(config, indent=2) + "\n", encoding="utf-8")


def pack_nvfp4_hf(codes: np.ndarray, block_scales: np.ndarray, per_tensor: float):
    """Pack codes and scales into NVFP4 tensors."""
    # Packed weights: low nibble is element 0, high nibble is element 1
    packed = (codes[:, 0::2] & 0x0F) | ((codes[:, 1::2] & 0x0F) << 4)
    packed_tensor = torch.from_numpy(np.ascontiguousarray(packed, dtype=np.uint8))

    scale_bits = ue4m3_bits(block_scales)
    scale_tensor = torch.from_numpy(np.ascontiguousarray(scale_bits, dtype=np.uint8)).view(torch.float8_e4m3fn)
    scale_2_tensor = torch.tensor(float(per_tensor), dtype=torch.float32)

    return packed_tensor, scale_tensor, scale_2_tensor


def _should_quantize(
    key: str,
    shape: tuple[int, ...],
    args,
    convert_to_quant_dir: Path | None,
) -> bool:
    if len(shape) != 2:
        return False
    if not (key.endswith(".weight") or key.endswith(".weight_packed")):
        return False
    if shape[-1] % 64 != 0 or shape[-1] < 64:
        return False
    if key in ("lm_head.weight", "lm_head.weight_packed") and not getattr(args, "quantize_output", False):
        return False
    lowered = key.lower()
    if "embed_tokens" in lowered or "wte" in lowered or "word_embeddings" in lowered:
        return False
    if "norm" in lowered or "ln" in lowered:
        return False
    if getattr(args, "exclude_re", None) and args.exclude_re.search(key):
        return False
    if ctq_skips_key(key, convert_to_quant_dir):
        return False
    return True


def _target_dtype(dtype_str: str) -> torch.dtype:
    mapping = {
        "float32": torch.float32,
        "f32": torch.float32,
        "float16": torch.float16,
        "f16": torch.float16,
        "bfloat16": torch.bfloat16,
        "bf16": torch.bfloat16,
    }
    return mapping.get(dtype_str.lower(), torch.bfloat16)


def export_directory(args) -> None:
    source_dir = Path(args.model).resolve()
    if not source_dir.is_dir():
        raise FileNotFoundError(f"Model directory does not exist: {source_dir}")

    _assign_output(args)
    output_dir = Path(args.output).resolve()
    if output_dir == source_dir:
        raise SystemExit("Output directory cannot be the same as the input directory.")
    if output_dir.exists() and any(output_dir.iterdir()) and not args.overwrite:
        raise SystemExit(f"Output directory {output_dir} already exists and is not empty. Use --overwrite to replace.")

    shards = discover_shards(source_dir)
    _apply_directory_profile(args, source_dir)
    args.exclude_re = re.compile(args.exclude) if args.exclude else None

    logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO, format="%(message)s")
    log.info("Loading %s", source_dir)
    log.info("Found %d safetensors shard(s)", len(shards))
    log.info("Target format: %s", "Hugging Face (compressed-tensors)" if args.format == "hf" else "ModelOpt / ComfyUI")

    output_dir.mkdir(parents=True, exist_ok=True)
    _copy_non_weight_files(source_dir, output_dir)

    is_passthrough = args.no_quant or args.outtype == "nvfp4"
    float_dtype = _target_dtype(args.output_dtype if not is_passthrough else (args.outtype if args.outtype != "auto" else "bfloat16"))

    if is_passthrough:
        log.info("Quantization skipped. Converting tensors to %s", args.outtype)
        weight_map: dict[str, str] = {}
        total_size = 0
        for shard_name in shards:
            log.info("Processing %s", shard_name)
            src_shard = source_dir / shard_name
            out_shard = output_dir / shard_name
            out_tensors: dict[str, torch.Tensor] = {}
            with safe_open(src_shard, framework="pt", device="cpu") as handle:
                for key in handle.keys():
                    t = handle.get_tensor(key)
                    if args.outtype in ("f32", "f16", "bf16") and t.is_floating_point():
                        t = t.to(dtype=float_dtype)
                    out_tensors[key] = t
                    weight_map[key] = shard_name
                    total_size += t.numel() * t.element_size()
            save_file(out_tensors, out_shard)
        _write_config_json(source_dir, output_dir, args, is_quantized=False)
        if len(shards) > 1:
            (output_dir / INDEX_NAME).write_text(
                json.dumps({"metadata": {"total_size": total_size}, "weight_map": weight_map}, indent=2) + "\n",
                encoding="utf-8",
            )
        log.info("Wrote %s", output_dir)
        return

    # Learned NVFP4 Quantization
    ensure_cuda_interpreter(args.device)
    args.optimizer = resolve_optimizer(args.optimizer)

    # Prepare checkpoint directory for per-tensor resume
    checkpoint_dir = Path(args.checkpoint_dir) if args.checkpoint_dir else output_dir / ".nvfp4-work"
    if args.fresh and checkpoint_dir.exists():
        shutil.rmtree(checkpoint_dir)
    checkpoint_dir.mkdir(parents=True, exist_ok=True)

    # Initialize quantizer
    import gguf
    quantizer = TensorQuantizer(args, gguf, load_learned_converter(args.convert_to_quant)[1])

    # First pass: collect all tensor keys and shapes across shards
    all_keys: dict[str, tuple[str, tuple[int, ...]]] = {}
    for shard_name in shards:
        with safe_open(source_dir / shard_name, framework="pt", device="cpu") as handle:
            for key in handle.keys():
                slice_header = handle.get_slice(key)
                all_keys[key] = (shard_name, tuple(int(d) for d in slice_header.get_shape()))

    quant_count = 0
    passthrough_count = 0
    bias_count = 0
    weight_map: dict[str, str] = {}
    total_size = 0
    all_quant_meta: dict[str, dict] = {}

    for s_idx, shard_name in enumerate(shards, start=1):
        log.info("\n[%d/%d] Processing shard: %s", s_idx, len(shards), shard_name)
        src_shard = source_dir / shard_name
        out_shard = output_dir / shard_name
        shard_tensors: dict[str, torch.Tensor] = {}

        with safe_open(src_shard, framework="pt", device="cpu") as handle:
            keys_in_shard = list(handle.keys())

            for key in keys_in_shard:
                tensor = handle.get_tensor(key)
                shape = tuple(int(d) for d in tensor.shape)
                base_key = key.removesuffix(".weight_packed").removesuffix(".weight")

                if _should_quantize(key, shape, args, args.convert_to_quant):
                    log.info("  Quantizing [%s] %s", "x".join(str(d) for d in shape), key)
                    weight_np = tensor.detach().float().cpu().numpy()

                    # Check for sibling bias
                    bias_key = f"{base_key}.bias"
                    has_bias = bias_key in all_keys and not args.no_bias_correction
                    bias_shard = all_keys[bias_key][0] if has_bias else None

                    # Check checkpoint cache
                    cached = _load_checkpoint(checkpoint_dir, key, shape, need_bias=has_bias)
                    if cached is not None:
                        codes, block_scale, per_tensor, bias_corr = cached
                        log.info("    Loaded from checkpoint: %s", key)
                    else:
                        rounds = getattr(args, "scale_refinement_rounds", DEFAULT_SCALE_REFINEMENT_ROUNDS)
                        codes, block_scale, per_tensor, dequant = quantizer.quantize_matrix(weight_np, key)
                        codes, block_scale = coordinate_descent_refine(
                            weight_np, codes, block_scale, per_tensor, bake=False, rounds=rounds, device=quantizer.device,
                        )
                        codes, block_scale, per_tensor = sweep_per_tensor_scale(
                            weight_np, codes, block_scale, per_tensor, bake=False, device=quantizer.device,
                        )

                        bias_corr = None
                        if has_bias:
                            bias_corr = quantizer.bias_correction(weight_np, dequant)

                        _save_checkpoint(checkpoint_dir, key, shape, codes, block_scale, per_tensor, bias_corr)

                    packed_w, scale_1, scale_2 = pack_nvfp4_hf(codes, block_scale, per_tensor)

                    if args.format == "hf":
                        # Standard Hugging Face compressed-tensors layout
                        shard_tensors[f"{base_key}.weight_packed"] = packed_w
                        shard_tensors[f"{base_key}.weight_scale"] = scale_1
                        shard_tensors[f"{base_key}.weight_global_scale"] = torch.tensor(1.0 / float(per_tensor), dtype=torch.float32)
                    else:
                        # ModelOpt / ComfyUI layout
                        shard_tensors[f"{base_key}.weight"] = packed_w
                        shard_tensors[f"{base_key}.weight_scale"] = scale_1
                        shard_tensors[f"{base_key}.weight_scale_2"] = scale_2
                        shard_tensors[f"{base_key}.comfy_quant"] = torch.tensor([1], dtype=torch.uint8)

                    quant_count += 1
                    all_quant_meta[base_key] = {
                        "quant_type": "nvfp4",
                        "block_size": 16,
                        "scale_type": "float8_e4m3fn",
                    }

                    # Bias correction application if bias is in the same shard
                    if has_bias and bias_shard == shard_name and bias_corr is not None:
                        bias_t = handle.get_tensor(bias_key).float()
                        corrected_bias = bias_t - torch.from_numpy(bias_corr).to(bias_t.device)
                        shard_tensors[bias_key] = corrected_bias.to(dtype=float_dtype)
                        bias_count += 1
                else:
                    # Non-quantized tensors (or bias already processed)
                    if key in shard_tensors:
                        continue
                    if tensor.is_floating_point():
                        if tensor.ndim <= 1 and key.endswith("norm.weight"):
                            tensor = tensor.float()
                        else:
                            tensor = tensor.to(dtype=float_dtype)
                    shard_tensors[key] = tensor
                    passthrough_count += 1

        shard_metadata = {}
        if args.format != "hf":
            shard_metadata["_quantization_metadata"] = json.dumps({"format_version": "1.0", "layers": all_quant_meta}, separators=(",", ":"))

        for k, t in shard_tensors.items():
            weight_map[k] = shard_name
            total_size += t.numel() * t.element_size()

        save_file(shard_tensors, out_shard, metadata=shard_metadata if shard_metadata else None)

    _write_config_json(source_dir, output_dir, args, is_quantized=True)

    if len(shards) > 1:
        index_payload = {
            "metadata": {"total_size": total_size},
            "weight_map": weight_map,
        }
        (output_dir / INDEX_NAME).write_text(json.dumps(index_payload, indent=2) + "\n", encoding="utf-8")

    if not args.keep_checkpoints and quant_count:
        shutil.rmtree(checkpoint_dir, ignore_errors=True)

    log.info(
        "\nWrote %s (%d NVFP4 tensors, %d biases adjusted, %d other tensors)",
        output_dir,
        quant_count,
        bias_count,
        passthrough_count,
    )


def self_test(args) -> None:
    """Check that packed blocks roundtrip accurately with dequantizer."""
    logging.basicConfig(level=logging.INFO, format="%(message)s")
    from nvfp4gguf import FP4_TABLE
    import gguf

    log.info("Running NVFP4 roundtrip self-test...")
    _converter_cls, unpack_uint4 = load_learned_converter(args.convert_to_quant)
    quantizer = TensorQuantizer(args, gguf, unpack_uint4)

    generator = np.random.default_rng(0)
    for shape in ((64, 128), (128, 256)):
        weight = generator.normal(size=shape).astype(np.float32) * 0.05
        codes, block_scale, per_tensor, dequant = quantizer.quantize_matrix(weight)
        packed_w, scale_1, scale_2 = pack_nvfp4_hf(codes, block_scale, per_tensor)

        # Dequantize packed tensor
        w_u = packed_w.numpy()
        c0 = w_u & 0x0F
        c1 = (w_u >> 4) & 0x0F
        unpacked_codes = np.empty(shape, dtype=np.uint8)
        unpacked_codes[:, 0::2] = c0
        unpacked_codes[:, 1::2] = c1
        assert np.array_equal(codes, unpacked_codes), "Nibble unpacking mismatch!"

        # Check values
        fp4_vals = FP4_TABLE[unpacked_codes]
        b_scales = block_scale.repeat(16, axis=1)
        recon = fp4_vals * b_scales * per_tensor
        max_diff = np.max(np.abs(recon - dequant))
        assert max_diff < 1e-4, f"Reconstruction difference {max_diff} exceeded tolerance!"
        log.info("Shape %s roundtrip ok (max diff %.3e)", shape, max_diff)

    log.info("self-test passed successfully.")


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Quantize a Hugging Face model directory to NVFP4 Hugging Face format.",
    )
    parser.add_argument("model", nargs="?", help="Hugging Face model directory (config.json + safetensors)")
    parser.add_argument("-o", "--output", help="Output directory. Default: ./<model-name>-NVFP4")
    parser.add_argument("--format", default="hf", choices=("hf", "modelopt", "comfy"),
                        help="Safetensors tensor layout: 'hf' (default) uses Hugging Face compressed-tensors (.weight_packed, .weight_scale, .weight_global_scale). 'modelopt' (or 'comfy') uses raw .weight.")
    parser.add_argument("--overwrite", action="store_true", help="Replace existing files in the output directory")
    parser.add_argument("--no-quant", action="store_true", help="Skip learned rounding. Rewrites or repacks tensors")
    parser.add_argument("--outtype", default="auto", choices=("auto", "f32", "f16", "bf16", "nvfp4"),
                        help="Float type for tensors when --no-quant is used (default: auto)")
    parser.add_argument("--output-dtype", default="bfloat16", choices=("bfloat16", "float16"),
                        help="Dtype for unquantized 2D weights and embeddings (default: bfloat16)")
    parser.add_argument("--context-length", type=int, default=None,
                        help="Override context length stored in config.json")
    parser.add_argument("--model-name", default=None, help="Override model name")
    parser.add_argument("--checkpoint-dir", type=Path, default=None, help="Per-tensor resume directory")
    parser.add_argument("--fresh", action="store_true", help="Delete checkpoint directory before starting")
    parser.add_argument("--keep-checkpoints", action="store_true", help="Keep per-tensor checkpoints upon completion")
    parser.add_argument("--eager", action="store_true", help="Eager loading")
    parser.add_argument("--low-memory", action="store_true", help="Stream each shard instead of loading all at once")
    parser.add_argument("--heur", action="store_true", help="Skip layers heuristic rejects")
    parser.add_argument("--verbose", action="store_true", help="Enable verbose logging")
    parser.add_argument("--self-test", action="store_true", help="Run NVFP4 packing self-test and exit")
    parser.add_argument("--llama-cpp", type=Path, default=DEFAULT_LLAMA_CPP,
                        help="llama.cpp checkout directory (default: ~/llama.cpp)")

    add_quant_arguments(parser)
    args = parser.parse_args(argv)

    if args.self_test:
        if args.device == "auto":
            args.device = "cpu"
        return args

    if not args.model:
        parser.error("model directory is required")

    if args.convert_to_quant is None and DEFAULT_CONVERT_TO_QUANT.is_dir():
        args.convert_to_quant = DEFAULT_CONVERT_TO_QUANT

    return args


def main(argv: list[str] | None = None) -> None:
    args = parse_args(argv)
    if args.self_test:
        self_test(args)
        return
    export_directory(args)


if __name__ == "__main__":
    main()