#!/usr/bin/env python3
"""Create the shared F16 GLM-5.2-Vision mmproj GGUF for llama.cpp."""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import shutil
import struct
import subprocess
import sys
import tempfile
from pathlib import Path
from typing import Any, Iterable, Sequence


DEFAULT_SOURCE = Path("../GLM-5.2-Vision-FP8")
DEFAULT_OUTPUT = Path("mmproj-GLM-5.2-Vision-f16.gguf")
DEFAULT_LLAMA_CPP = Path("../../llama.cpp")
LLAMA_CPP_PR_URL = "https://github.com/ggml-org/llama.cpp/pull/26126"
VISION_FILENAME = "vision_tower.safetensors"
PROJECTOR_FILENAME = "mm_projector.safetensors"
EXPECTED_VISION_TENSORS = 329
EXPECTED_PROJECTOR_TENSORS = 6
EXPECTED_OUTPUT_TENSORS = EXPECTED_VISION_TENSORS + EXPECTED_PROJECTOR_TENSORS
MIN_FREE_BYTES = 2_000_000_000
BUILD_TARGETS = ("llama-server", "llama-mtmd-cli", "llama-mtmd-debug")


class GeneratorError(RuntimeError):
    pass


def info(message: str) -> None:
    print(f"[glm5v-mmproj] {message}", flush=True)


def display_command(command: Sequence[str]) -> str:
    import shlex

    displayed: list[str] = []
    previous = ""
    for part in command:
        if previous == "-c" and len(part) > 120:
            displayed.append("<inline Python>")
        else:
            displayed.append(part)
        previous = part
    return " ".join(shlex.quote(part) for part in displayed)


def run(
    command: Sequence[str],
    *,
    cwd: Path | None = None,
    input_bytes: bytes | None = None,
    capture: bool = False,
    check: bool = True,
) -> subprocess.CompletedProcess[bytes]:
    info(f"Running: {display_command(command)}")
    result = subprocess.run(
        list(command),
        cwd=cwd,
        input=input_bytes,
        stdout=subprocess.PIPE if capture else None,
        stderr=subprocess.PIPE if capture else None,
        check=False,
    )
    if check and result.returncode != 0:
        details = ""
        if capture:
            output = (result.stdout + result.stderr).decode("utf-8", errors="replace").strip()
            if output:
                details = f"\n{output}"
        raise GeneratorError(f"Command failed with exit code {result.returncode}: {display_command(command)}{details}")
    return result


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--source", type=Path, default=DEFAULT_SOURCE, help=f"GLM-5.2-Vision Hugging Face directory (default: {DEFAULT_SOURCE})")
    parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT, help=f"Output GGUF path (default: {DEFAULT_OUTPUT})")
    parser.add_argument("--llama-cpp", type=Path, default=DEFAULT_LLAMA_CPP, help=f"llama.cpp checkout (default: {DEFAULT_LLAMA_CPP})")
    parser.add_argument("--python", type=Path, help="Python interpreter with llama.cpp conversion dependencies")
    parser.add_argument("--jobs", type=int, default=max(1, os.cpu_count() or 1), help="Parallel CMake build jobs")
    parser.add_argument("--force", action="store_true", help="Regenerate and atomically replace an existing output")
    parser.add_argument("--check-only", action="store_true", help="Validate prerequisites without building or converting")
    return parser.parse_args()


def resolve_directory(path: Path, label: str) -> Path:
    resolved = path.expanduser().resolve()
    if not resolved.is_dir():
        raise GeneratorError(f"{label} is not a directory: {resolved}")
    return resolved


def read_json(path: Path) -> dict[str, Any]:
    try:
        with path.open("r", encoding="utf-8") as handle:
            value = json.load(handle)
    except (OSError, json.JSONDecodeError) as exc:
        raise GeneratorError(f"Cannot read JSON file {path}: {exc}") from exc
    if not isinstance(value, dict):
        raise GeneratorError(f"Expected a JSON object in {path}")
    return value


def read_safetensors_header(path: Path) -> dict[str, dict[str, Any]]:
    try:
        file_size = path.stat().st_size
        with path.open("rb") as handle:
            raw_length = handle.read(8)
            if len(raw_length) != 8:
                raise GeneratorError(f"Invalid Safetensors header in {path}")
            header_length = struct.unpack("<Q", raw_length)[0]
            if header_length <= 0 or header_length > file_size - 8:
                raise GeneratorError(f"Invalid Safetensors header length in {path}: {header_length}")
            header = json.loads(handle.read(header_length))
    except (OSError, json.JSONDecodeError) as exc:
        raise GeneratorError(f"Cannot read Safetensors header {path}: {exc}") from exc

    tensors = {name: value for name, value in header.items() if name != "__metadata__"}
    if not tensors or not all(isinstance(value, dict) for value in tensors.values()):
        raise GeneratorError(f"No valid tensor entries found in {path}")

    data_size = file_size - 8 - header_length
    for name, tensor in tensors.items():
        offsets = tensor.get("data_offsets")
        if not isinstance(offsets, list) or len(offsets) != 2:
            raise GeneratorError(f"Invalid data offsets for {name} in {path}")
        start, end = offsets
        if not isinstance(start, int) or not isinstance(end, int) or not 0 <= start <= end <= data_size:
            raise GeneratorError(f"Out-of-range data offsets for {name} in {path}")
    return tensors


def validate_source(source: Path) -> tuple[dict[str, dict[str, Any]], dict[str, dict[str, Any]]]:
    config_path = source / "config.json"
    preprocessor_path = source / "preprocessor_config.json"
    vision_path = source / VISION_FILENAME
    projector_path = source / PROJECTOR_FILENAME
    for path in (config_path, preprocessor_path, vision_path, projector_path):
        if not path.is_file():
            raise GeneratorError(f"Required source file is missing: {path}")

    config = read_json(config_path)
    architectures = config.get("architectures", [])
    if "Glm5vForConditionalGeneration" not in architectures:
        raise GeneratorError(f"Unexpected model architecture in {config_path}: {architectures}")

    text_config = config.get("text_config")
    vision_config = config.get("vision_config")
    if not isinstance(text_config, dict) or text_config.get("hidden_size") != 6144:
        raise GeneratorError("GLM-5.2 text hidden size must be 6144")
    expected_vision = {
        "hidden_size": 1152,
        "intermediate_size": 4304,
        "num_attention_heads": 16,
        "num_hidden_layers": 27,
        "patch_size": 14,
        "merge_kernel_size": [2, 2],
    }
    if not isinstance(vision_config, dict):
        raise GeneratorError("vision_config is missing from config.json")
    mismatches = {key: (vision_config.get(key), expected) for key, expected in expected_vision.items() if vision_config.get(key) != expected}
    if mismatches:
        raise GeneratorError(f"Unexpected GLM-5.2 vision configuration values: {mismatches}")

    preprocessor = read_json(preprocessor_path)
    media_config = preprocessor.get("media_proc_cfg")
    if not isinstance(media_config, dict) or media_config.get("patch_size") != 14:
        raise GeneratorError("Expected MoonViT media_proc_cfg with patch_size 14")

    vision_tensors = read_safetensors_header(vision_path)
    projector_tensors = read_safetensors_header(projector_path)
    if len(vision_tensors) != EXPECTED_VISION_TENSORS:
        raise GeneratorError(f"Expected {EXPECTED_VISION_TENSORS} vision tensors, found {len(vision_tensors)}")
    if len(projector_tensors) != EXPECTED_PROJECTOR_TENSORS:
        raise GeneratorError(f"Expected {EXPECTED_PROJECTOR_TENSORS} projector tensors, found {len(projector_tensors)}")
    if any(not name.startswith("vision_tower.") for name in vision_tensors):
        raise GeneratorError("The vision Safetensors file contains a non-vision tensor")
    if any(not name.startswith("mm_projector.") for name in projector_tensors):
        raise GeneratorError("The projector Safetensors file contains a non-projector tensor")
    dtypes = {tensor.get("dtype") for tensor in [*vision_tensors.values(), *projector_tensors.values()]}
    if dtypes != {"BF16"}:
        raise GeneratorError(f"Expected BF16 vision and projector source tensors, found: {sorted(str(value) for value in dtypes)}")

    info(f"Validated source: {len(vision_tensors)} vision tensors + {len(projector_tensors)} projector tensors, all BF16")
    return vision_tensors, projector_tensors


def validate_glm5v_support(llama_cpp: Path) -> None:
    # GLM-5.2-Vision reuses the Kimi-K2.5 projector; the image begin/end tokens are
    # resolved at runtime from the text model vocab instead of a dedicated projector type
    required_markers = {
        "conversion/__init__.py": '"Glm5vForConditionalGeneration": "kimivl"',
        "conversion/kimivl.py": '@ModelBase.register("Glm5vForConditionalGeneration")',
        "tools/mtmd/clip-impl.h": "PROJECTOR_TYPE_KIMIK25",
        "tools/mtmd/mtmd.cpp": 'lookup_token("<|begin_of_image|>")',
    }
    missing: list[str] = []
    for relative_path, marker in required_markers.items():
        path = llama_cpp / relative_path
        try:
            contents = path.read_text(encoding="utf-8")
        except OSError:
            missing.append(relative_path)
            continue
        if marker not in contents:
            missing.append(relative_path)
    if missing:
        raise GeneratorError(f"llama.cpp does not contain the GLM-5.2-Vision support from {LLAMA_CPP_PR_URL}; missing support in: {', '.join(missing)}")
    info(f"Validated llama.cpp GLM-5.2-Vision support from {LLAMA_CPP_PR_URL}")


def interpreter_has_dependencies(python: Path, llama_cpp: Path) -> bool:
    probe = "import numpy, safetensors, torch, transformers; print('conversion dependencies OK')"
    result = run([str(python), "-c", probe], capture=True, check=False)
    if result.returncode != 0:
        return False
    converter = llama_cpp / "convert_hf_to_gguf.py"
    return run([str(python), str(converter), "--help"], capture=True, check=False).returncode == 0


def select_python(explicit: Path | None, llama_cpp: Path) -> Path:
    if explicit is not None:
        candidates: Iterable[Path] = (explicit.expanduser(),)
    else:
        candidates = (Path(sys.executable), Path("/usr/bin/python3"))

    checked: set[Path] = set()
    for candidate in candidates:
        candidate = candidate.expanduser()
        executable = candidate if candidate.is_absolute() else Path(shutil.which(str(candidate)) or candidate)
        try:
            resolved = executable.resolve(strict=True)
        except OSError:
            continue
        if resolved in checked or not os.access(executable, os.X_OK):
            continue
        checked.add(resolved)
        if interpreter_has_dependencies(executable, llama_cpp):
            info(f"Using converter Python: {executable}")
            return executable

    if explicit is not None:
        raise GeneratorError(f"The requested Python interpreter lacks llama.cpp conversion dependencies: {explicit}")
    raise GeneratorError("No Python interpreter with torch, transformers, numpy, and safetensors was found; use --python PATH")


def validate_build(llama_cpp: Path) -> Path:
    converter = llama_cpp / "convert_hf_to_gguf.py"
    build_dir = llama_cpp / "build"
    if not converter.is_file():
        raise GeneratorError(f"llama.cpp converter is missing: {converter}")
    if not (llama_cpp / ".git").exists():
        raise GeneratorError(f"llama.cpp path is not a Git checkout: {llama_cpp}")
    if not (build_dir / "CMakeCache.txt").is_file():
        raise GeneratorError(f"llama.cpp build directory is not configured: {build_dir}")
    return build_dir


def build_llama_cpp(llama_cpp: Path, build_dir: Path, jobs: int) -> None:
    if jobs < 1:
        raise GeneratorError("--jobs must be at least 1")
    cmake = shutil.which("cmake")
    if cmake is None:
        raise GeneratorError("cmake is not available on PATH")
    info("Refreshing the existing llama.cpp CMake configuration")
    run([cmake, "-S", str(llama_cpp), "-B", str(build_dir)], cwd=llama_cpp)
    info(f"Building GLM5V-capable llama.cpp targets with {jobs} jobs")
    run([cmake, "--build", str(build_dir), "--parallel", str(jobs), "--target", *BUILD_TARGETS], cwd=llama_cpp)
    missing = [name for name in BUILD_TARGETS if not (build_dir / "bin" / name).is_file()]
    if missing:
        raise GeneratorError(f"Build completed without expected executables: {missing}")


def tensor_data_size(tensors: dict[str, dict[str, Any]]) -> int:
    return sum(int(tensor["data_offsets"][1]) - int(tensor["data_offsets"][0]) for tensor in tensors.values())


def stage_mmproj_source(
    source: Path,
    staging: Path,
    vision_tensors: dict[str, dict[str, Any]],
    projector_tensors: dict[str, dict[str, Any]],
) -> None:
    for name in ("config.json", "preprocessor_config.json", "README.md", "generation_config.json"):
        source_file = source / name
        if source_file.is_file():
            shutil.copy2(source_file, staging / name)

    vision_part = "model-00001-of-00002.safetensors"
    projector_part = "model-00002-of-00002.safetensors"
    (staging / vision_part).symlink_to(source / VISION_FILENAME)
    (staging / projector_part).symlink_to(source / PROJECTOR_FILENAME)
    weight_map = {name: vision_part for name in vision_tensors}
    weight_map.update({name: projector_part for name in projector_tensors})
    index = {
        "metadata": {"total_size": tensor_data_size(vision_tensors) + tensor_data_size(projector_tensors)},
        "weight_map": weight_map,
    }
    with (staging / "model.safetensors.index.json").open("w", encoding="utf-8") as handle:
        json.dump(index, handle, indent=2, sort_keys=True)
        handle.write("\n")


GGUF_VALIDATOR = r"""
import json
import sys
from pathlib import Path

sys.path.insert(0, sys.argv[2])
from gguf import GGUFReader

reader = GGUFReader(sys.argv[1], "r")

def value(name):
    field = reader.get_field(name)
    if field is None:
        raise RuntimeError(f"missing GGUF metadata field: {name}")
    return field.contents()

expected = {
    "general.type": "mmproj",
    "general.architecture": "clip",
    "clip.projector_type": "kimik25",
    "clip.has_vision_encoder": True,
    "clip.vision.projection_dim": 6144,
}
for key, wanted in expected.items():
    actual = value(key)
    if actual != wanted:
        raise RuntimeError(f"unexpected {key}: expected {wanted!r}, got {actual!r}")

if len(reader.tensors) != 335:
    raise RuntimeError(f"expected 335 tensors, found {len(reader.tensors)}")

names = {tensor.name for tensor in reader.tensors}
required_names = {"v.patch_embd.weight", "mm.input_norm.weight", "mm.1.weight", "mm.2.weight"}
missing = sorted(required_names - names)
if missing:
    raise RuntimeError(f"missing expected tensors: {missing}")

types = sorted({tensor.tensor_type.name for tensor in reader.tensors})
if not set(types).issubset({"F16", "F32"}):
    raise RuntimeError(f"unexpected tensor types in F16 mmproj: {types}")

print(json.dumps({"tensor_count": len(reader.tensors), "tensor_types": types, "projector_type": value("clip.projector_type"), "projection_dim": value("clip.vision.projection_dim")}))
"""


def validate_output(output: Path, python: Path, llama_cpp: Path) -> dict[str, Any]:
    if not output.is_file():
        raise GeneratorError(f"Output GGUF does not exist: {output}")
    result = run(
        [str(python), "-c", GGUF_VALIDATOR, str(output), str(llama_cpp / "gguf-py")],
        capture=True,
        check=False,
    )
    if result.returncode != 0:
        details = (result.stdout + result.stderr).decode("utf-8", errors="replace").strip()
        raise GeneratorError(f"GGUF validation failed for {output}:\n{details}")
    try:
        summary = json.loads(result.stdout.decode("utf-8"))
    except json.JSONDecodeError as exc:
        raise GeneratorError(f"GGUF validator returned invalid output: {result.stdout!r}") from exc
    info(f"Validated GGUF: {summary['tensor_count']} tensors, types {summary['tensor_types']}, projector {summary['projector_type']}, output dimension {summary['projection_dim']}")
    return summary


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(16 * 1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def convert(
    source: Path,
    output: Path,
    llama_cpp: Path,
    python: Path,
    vision_tensors: dict[str, dict[str, Any]],
    projector_tensors: dict[str, dict[str, Any]],
) -> None:
    output.parent.mkdir(parents=True, exist_ok=True)
    if shutil.disk_usage(output.parent).free < MIN_FREE_BYTES:
        raise GeneratorError(f"At least {MIN_FREE_BYTES:,} free bytes are required in {output.parent}")

    with tempfile.TemporaryDirectory(prefix=".glm5v-mmproj-", dir=output.parent) as temporary:
        staging = Path(temporary)
        stage_mmproj_source(source, staging, vision_tensors, projector_tensors)
        temporary_output = staging / output.name
        command = [
            str(python),
            str(llama_cpp / "convert_hf_to_gguf.py"),
            str(staging),
            "--mmproj",
            "--outtype", "f16",
            "--model-name", "GLM-5.2-Vision",
            "--outfile", str(temporary_output),
        ]
        info("Converting the BF16 MoonViT tower and projector to an F16 GGUF")
        run(command, cwd=llama_cpp)
        validate_output(temporary_output, python, llama_cpp)
        os.replace(temporary_output, output)


def main() -> int:
    args = parse_args()
    if args.jobs < 1:
        raise GeneratorError("--jobs must be at least 1")
    source = resolve_directory(args.source, "Source model")
    llama_cpp = resolve_directory(args.llama_cpp, "llama.cpp checkout")
    output = args.output.expanduser().resolve()

    vision_tensors, projector_tensors = validate_source(source)
    build_dir = validate_build(llama_cpp)
    validate_glm5v_support(llama_cpp)
    python = select_python(args.python, llama_cpp)

    if args.check_only:
        if output.exists():
            validate_output(output, python, llama_cpp)
        info("All prerequisite checks passed; no files were changed")
        return 0

    build_llama_cpp(llama_cpp, build_dir, args.jobs)

    if output.exists() and not args.force:
        try:
            validate_output(output, python, llama_cpp)
        except GeneratorError as exc:
            raise GeneratorError(f"Output already exists but is invalid; use --force to replace it: {exc}") from exc
        info(f"Valid output already exists; use --force to regenerate it: {output}")
    else:
        convert(source, output, llama_cpp, python, vision_tensors, projector_tensors)

    digest = sha256_file(output)
    info(f"Created: {output}")
    info(f"Size: {output.stat().st_size:,} bytes")
    info(f"SHA-256: {digest}")
    return 0


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except (GeneratorError, KeyboardInterrupt) as exc:
        print(f"[glm5v-mmproj] ERROR: {exc}", file=sys.stderr)
        raise SystemExit(1)