# (c) City96 || Apache-2.0 (apache.org/licenses/LICENSE-2.0)
import os
import re
import sys
import time
import torch
import logging
import argparse
import subprocess
import huggingface_hub as hf

logging.getLogger().setLevel(logging.DEBUG)

qtypes =[
    # "F16", "BF16",
    "Q8_0", "Q6_K",
    "Q5_K_M", "Q5_K_S", "Q5_1", "Q5_0",
    "Q4_K_M", "Q4_K_S", "Q4_1", "Q4_0",
    "Q3_K_M", "Q3_K_S", "Q2_K"
]

dtype_dict = {
    "F32": torch.float32,
    "F16": torch.float16,
    "BF16": torch.bfloat16,
    "F8_E4M3": getattr(torch, "float8_e4m3fn", "_invalid"),
    "F8_E5M2": getattr(torch, "float8_e5m2", "_invalid"),
}

# this is pretty jank but I want to be able to run it on a blank instance w/o setup
terraform_dict = {
    "repo": "city96/ComfyUI-GGUF",
    "target": "auto_convert",
    "lcpp_repo": "ggerganov/llama.cpp",
    "lcpp_target": "tags/b3962",
}

def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--src", required=True, help="Source model file or huggingface repo name")
    parser.add_argument("--quants", nargs="+", choices=["all", "base", *qtypes], default=["Q8_0"])
    parser.add_argument("--output-dir", default=None, help="Location for output files, defaults to current dir or ComfyUI model dir.")
    parser.add_argument("--temp-dir", default=None, help="Location for temp files, defaults to [output_dir]/tmp")
    parser.add_argument("--force-update", action="store_true", help="Force update & rebuild entire quantization stack.")
    parser.add_argument("--resume", action="store_true", help="Skip over existing files. Will NOT check for broken/interrupted files.")

    args = parser.parse_args()
    if args.output_dir is None:
        args.output_dir = get_output_dir()
    if args.temp_dir is None:
        args.temp_dir = os.path.join(args.output_dir, "tmp")

    if os.path.isdir(args.temp_dir) and len(os.listdir(args.temp_dir)) > 0:
        raise OSError("Output temp folder not empty!")

    if "all" in args.quants:
        args.quants = ["base", *qtypes]

    return args

def run_cmd(*args, log_error=False):
    logging.debug(f"cmd: {args}")
    try:
        log = subprocess.run(args, capture_output=True, text=True)
    except Exception as e:
        logging.warning(f"{args[0]}, {e}")
        return -1
    if log.returncode != 0 and log_error:
        logging.warning(f"{args[0]}: {log.stdout} {log.stderr}")
    else:
        logging.debug(f"{args[0]}: {repr(log.stdout)} {repr(log.stderr.strip())} RET:{log.returncode}")
    return log.returncode

def setup_utils(force_update=False):
    # get ComfyUI-GGUF if missing, then compile patched llama.cpp if required
    root = os.path.dirname(os.path.abspath(__file__))
    root = os.path.normpath(root)

    if os.path.split(root)[1] != "tools":
        cg_dir = os.path.join(root, "ComfyUI-GGUF")
        if not os.path.isdir(cg_dir):
            logging.warning(f"Running outside tools folder! Cloning to {cg_dir}")
            run_cmd("git", "clone", f"https://github.com/{terraform_dict['repo']}", cg_dir)
            need_update = True
        else:
            need_update = False

        if force_update or need_update:
            if terraform_dict['target']:
                logging.info(f"Attemtping to check out ComfyUI-GGUF branch {terraform_dict['target']}")
                run_cmd("git", "-C", cg_dir, "checkout", terraform_dict['target'])

            logging.info("Attemtping to git pull ComfyUI-GGUF to latest")
            run_cmd("git", "-C", cg_dir, "pull")

        tools_dir = os.path.join(root, "ComfyUI-GGUF", "tools")
        sys.path.append(tools_dir) # to make import(s) work
    else:
        # TODO: Git pull here too?
        logging.warning(f"Assuming latest ComfyUI-GGUF. Please git pull & check out branch {terraform_dict['target']} manually!")
        tools_dir = root

    if not os.path.isdir(tools_dir):
        raise OSError(f"Can't find tools subfoder in ComfyUI-GGUF at {tools_dir}")

    convert_path = os.path.join(tools_dir, "convert.py")
    if not os.path.isfile(convert_path):
        raise OSError(f"Cannot find convert.py at location: {convert_path}")

    lcpp_path = os.path.join(root, "llama.cpp.auto") # avoid messing with regular dir
    if not os.path.isdir(lcpp_path):
        logging.info(f"Attemtping to clone llama.cpp repo to {lcpp_path}")
        run_cmd("git", "clone", f"https://github.com/{terraform_dict['lcpp_repo']}", lcpp_path)
        need_update = True
    else:
        need_update = False

    if force_update or need_update:
        # TODO: check reflog and/or git reset before checkout?
        logging.info(f"Attemtping to check out llama.cpp target {terraform_dict['lcpp_target']}")
        run_cmd("git", "-C", lcpp_path, "checkout", terraform_dict['lcpp_target'])

        # TODO: git reset before patch?
        patch_path = os.path.join(tools_dir, "lcpp.patch")
        # patch (probably) has wrong file endings:
        logging.info("Converting patch file endings")
        with open(patch_path, "rb") as file:
            content = file.read().replace(b"\r\n", b"\n")
        with open(patch_path, "wb") as file:
            file.write(content)

        if run_cmd("git", "-C", lcpp_path, "apply", "--check", "-R", patch_path) != 0:
            logging.info("Attemtping to apply patch to llama.cpp repo")
            run_cmd("git", "-C", lcpp_path, "apply", patch_path)
        else:
            logging.info("Patch already applied")

    # using cmake here as llama.cpp switched to it completely for new versions
    if os.name == "nt":
        bin_path = os.path.join(lcpp_path, "build", "bin", "debug", "llama-quantize.exe")
    else:
        bin_path = os.path.join(lcpp_path, "build", "bin", "llama-quantize")

    if not os.path.isfile(bin_path) or force_update or need_update:
        if run_cmd("cmake", "--version") != 0:
            raise RuntimeError("Can't find cmake! Make sure you have a working build environment set up")

        build_path = os.path.join(lcpp_path, "build")
        os.makedirs(build_path, exist_ok=True)
        logging.info("Attempting to build llama.cpp binary from source")
        run_cmd("cmake", "-B", build_path, lcpp_path)
        run_cmd("cmake", "--build", build_path, "--config", "Debug", "-j4", "--target", "llama-quantize")
        if not os.path.isfile(bin_path):
            raise RuntimeError("Build failed! Rerun with --debug to see error log.")
    else:
        logging.info("Binary already present")

    return bin_path

def get_output_dir():
    root = os.path.dirname(os.path.abspath(__file__))
    root = os.path.normpath(root)
    split = os.path.split(root)
    while split[1]:
        if split[1] == "ComfyUI":
            if os.path.isdir(os.path.join(*split, "models", "unet")): # new
                root = os.path.join(*split, "models", "unet", "gguf")
                logging.info(f"Found ComfyUI, using model folder: {root}")
                return root

            if os.path.isdir(os.path.join(*split, "models", "diffusion_models")): # old
                root = os.path.join(*split, "models", "diffusion_models", "gguf")
                logging.info(f"Found ComfyUI, using model folder: {root}")
                return root

            logging.info("Found ComfyUI, but can't find model folder")
            break

        split = os.path.split(split[0])

    root = os.path.join(root, "models")
    logging.info(f"Defaulting to [script dir]/models: {root}")
    return root

def get_hf_fake_sd(repo, path, device=torch.device("meta")):
    sd = {}
    meta = hf.parse_safetensors_file_metadata(repo, path)
    for key, raw in meta.tensors.items():
        shape = tuple(raw.shape)
        dtype = dtype_dict.get(raw.dtype, torch.float32)
        sd[key] = torch.zeros(shape, dtype=dtype, device=device)
    return sd

def get_hf_file_arch(repo, path):
    pattern = r'(\d+)-of-(\d+)'
    match = re.search(pattern, path)

    if match:
        # we need to load it as multipart
        if int(match.group(1)) != 1:
            return None
        sd = {}
        for k in range(int(match.group(2))):
            shard_path = path.replace(match.group(1), f"{k+1:0{len(match.group(1))}}")
            sd.update(get_hf_fake_sd(repo, shard_path))
    else:
        sd = get_hf_fake_sd(repo, path)

    # this should raise an error on failure
    sd = strip_prefix(sd)
    model_arch = detect_arch(sd)

    # this is for SDXL and SD1.5, I want to overhaul this logic to match sd.cpp eventually
    assert not model_arch.shape_fix, "Model uses shape fix (SDXL/SD1) - unsupported for now."
    return model_arch.arch

def get_hf_valid_files(repo):
    # TODO: probably tweak this?
    MIN_SIZE_GB = 1
    VALID_SRC_EXTS = [".safetensors", ] # ".pt", ".ckpt", ]
    meta = hf.model_info(repo, files_metadata=True)

    valid = {}
    for file in meta.siblings:
        path = file.rfilename
        fname = os.path.basename(path)
        name, ext = os.path.splitext(fname)

        if ext.lower() not in VALID_SRC_EXTS:
            logging.debug(f"Invalid ext: {path} {ext}")
            continue

        if file.size / (1024 ** 3) < MIN_SIZE_GB:
            logging.debug(f"File too small: {path} {file.size}")
            continue

        try:
            arch = get_hf_file_arch(repo, path)
        except Exception as e:
            logging.warning(f"Arch detect fail: {e} ({path})")
        else:
            if arch is not None:
                valid[path] = arch
                logging.info(f"Found '{arch}' model at path {path}")
    return valid

def make_base_quant(src, output_dir, temp_dir, final=True, resume=True):
    name, ext = os.path.splitext(os.path.basename(src))
    if ext == ".gguf":
        logging.info("Input file already in gguf, assuming base quant")
        return None, src, None

    name = name.lower() # uncomment to preserve case in all quants
    dst_tmp = os.path.join(temp_dir, f"{name}-{{ftype}}.gguf") # ftype is filled in by convert.py

    tmp_path, model_arch, fix_path = convert_file(src, dst_tmp, interact=False, overwrite=False)
    dst_path = os.path.join(output_dir, os.path.basename(tmp_path))
    if os.path.isfile(dst_path):
        if resume:
            logging.warning("Resuming with interrupted base quant, may be incorrect!")
            return dst_path, tmp_path, fix_path
        raise OSError(f"Output already exists! Clear folder? {dst_path}")

    if fix_path is not None and os.path.isfile(fix_path):
        quant_source = tmp_path
        if final:
            apply_5d_fix(tmp_path, dst_path, fix=fix_path, overwrite=False)
        else:
            dst_path = None
    else:
        fix_path = None
        if final:
            os.rename(tmp_path, dst_path)
            quant_source = dst_path
        else:
            dst_path = None
            quant_source = tmp_path

    return dst_path, quant_source, fix_path

def make_quant(src, output_dir, temp_dir, qtype, quantize_binary, fix_path=None, resume=True):
    name, ext = os.path.splitext(os.path.basename(src))
    assert ext.lower() == ".gguf", "Invalid input file"

    src_qtext = [x for x in ["-F32.gguf", "-F16.gguf", "-BF16.gguf"] if x in src]
    if len(src_qtext) == 1:
        tmp_path = os.path.join(
            temp_dir,
            os.path.basename(src).replace(src_qtext[0], f"-{qtype.upper()}.gguf")
        )
    else:
        tmp_path = os.path.join(
            temp_dir,
            f"{name}-{qtype.upper()}.gguf"
        )
    tmp_path = os.path.abspath(tmp_path)
    dst_path = os.path.join(output_dir, os.path.basename(tmp_path))
    if os.path.isfile(dst_path):
        if resume:
            return dst_path
        raise OSError("Output already exists! Clear folder?")

    r = run_cmd(quantize_binary, src, tmp_path, qtype, log_error=True)
    time.sleep(2) # leave time for file sync?
    if r != 0:
        raise RuntimeError(f"Quantization failed with error code {r}")

    if fix_path is not None:
        apply_5d_fix(tmp_path, dst_path, fix=fix_path, overwrite=False)
        if os.path.isfile(dst_path) and os.path.isfile(tmp_path):
            os.remove(tmp_path)
    else:
        os.rename(tmp_path, dst_path)

    return dst_path

if __name__ == "__main__":
    args = get_args()
    os.makedirs(args.output_dir, exist_ok=True)
    os.makedirs(args.temp_dir, exist_ok=True)
    quantize_binary = setup_utils(args.force_update)

    try:
        from convert import detect_arch, strip_prefix, convert_file
        from fix_5d_tensors import apply_5d_fix
    except [ImportError, ModuleNotFoundError] as e:
        raise ImportError(f"Can't import required utils: {e}")

    if not os.path.isfile(args.src):
        # huggingface repo. TODO: file choice
        if len(args.src.split("/")) != "1":
            raise OSError(f"Invalid huggingface repo or model path {args.src}")
        raise NotImplementedError("HF not yet supported")
        # download then set to temp file
        # hf_repo = "Lightricks/LTX-Video" # "fal/AuraFlow-v0.3"
        # get_hf_valid_files(hf_repo)
        # args.src = ...

    out_files = []

    base_quant, quant_source, fix_path = make_base_quant(
        args.src,
        args.output_dir,
        args.temp_dir,
        final=("base" in args.quants),
        resume=args.resume,
    )
    if "base" in args.quants:
        args.quants = [x for x in args.quants if x not in ["base"]]
    if base_quant is not None:
        out_files.append(base_quant)

    for qtype in args.quants:
        out_files.append(make_quant(
            quant_source,
            args.output_dir,
            args.temp_dir,
            qtype,
            quantize_binary,
            fix_path,
            resume=args.resume,
        ))

    if fix_path is not None and os.path.isfile(fix_path):
        os.remove(fix_path)

    if base_quant != quant_source:
        # make sure our quant source is in the temp folder before removing it
        cc = os.path.commonpath([os.path.normpath(quant_source), os.path.normpath(args.temp_dir)])
        if cc == os.path.normpath(args.temp_dir):
            os.remove(quant_source)

    out_file_str = '\n'.join(out_files)
    logging.info(f"Output file(s): {out_file_str}")