import os
import sys
import re
import gc
import json
import time
import subprocess
import numpy as np

# ============================================================================
# USER CONFIG
# ============================================================================

WORK_DIR = "/content/gguf_workspace"
os.makedirs(WORK_DIR, exist_ok=True)

# --- YOUR MODEL (input) ------------------------------------------------------
# Put your trained model GGUF here (Q8_0 in your case):
INPUT_FILE = os.path.join(WORK_DIR, "Qwopus3.5-27B-v3-Uncensored-FernflowerAI.Q8_0.gguf")

# --- Unsloth reference quant (used only to extract profile) ------------------
REF_FILE = os.path.join(WORK_DIR, "Qwen3.5-27B-UD-IQ2_XXS.gguf")
REF_URL = (
    "https://huggingface.co/unsloth/Qwen3.5-27B-GGUF/blob/main/Qwen3.5-27B-UD-IQ2_XXS.gguf?download=true"
)

# --- iMatrix (recommended) ---------------------------------------------------
IMATRIX_FILE = os.path.join(WORK_DIR, "imatrix_unsloth.gguf_file")
IMATRIX_URL = (
    "https://huggingface.co/unsloth/Qwen3.5-27B-GGUF/resolve/main/imatrix_unsloth.gguf_file?download=true"
)

# --- Output ------------------------------------------------------------------
TARGET_BASE_TYPE = "Q2_K"
OUTPUT_TAG = "UD"  # marker in output filename

# --- Store the extracted quant schema in JSON --------------------------------
PROFILE_JSON = os.path.join(WORK_DIR, "unsloth_ud_iq4_xs_profile.json")

# After JSON is built, delete Unsloth reference quant GGUF to save disk
DELETE_REF_AFTER_PROFILE = True

# Keep JSON compact by default (do NOT store huge name->type map)
# Turn on only if you really want the full audit dump (can be big).
STORE_RAW_TENSOR_MAP = False

# --- llama.cpp binaries ------------------------------------------------------
LLAMA_TAG = "b8192"
LLAMA_URL = (
    "https://github.com/ggml-org/llama.cpp/releases/download/"
    f"{LLAMA_TAG}/llama-{LLAMA_TAG}-bin-ubuntu-x64.tar.gz"
)
LLAMA_TAR = f"llama-{LLAMA_TAG}-bin-ubuntu-x64.tar.gz"
LLAMA_BIN_DIR = os.path.join(WORK_DIR, "llama-bin")

# Verbosity
PRINT_PROFILE_SUMMARY = True
VERIFY_AFTER_QUANT = True


# ============================================================================
# UTILS
# ============================================================================

def ensure_python_pkg(pkg: str):
    subprocess.check_call(
        [sys.executable, "-m", "pip", "install", "-q", pkg],
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
    )

def file_size_gb(path):
    return os.path.getsize(path) / (1 << 30)

def disk_free_gb(path=WORK_DIR):
    st = os.statvfs(path)
    return (st.f_bavail * st.f_frsize) / (1 << 30)

def ram_used_gb():
    try:
        with open("/proc/self/status", "r") as f:
            for line in f:
                if line.startswith("VmRSS:"):
                    return int(line.split()[1]) / (1 << 20)
    except Exception:
        pass
    return -1

def safe_remove(path):
    try:
        if os.path.isfile(path):
            os.remove(path)
    except Exception:
        pass

def download_file(url, out_path, label, timeout=6*3600):
    if os.path.isfile(out_path) and os.path.getsize(out_path) > 1024:
        print(f"[{label}] Found: {out_path} ({file_size_gb(out_path):.2f} GB)")
        return out_path

    print(f"[{label}] Downloading -> {out_path}")
    os.makedirs(os.path.dirname(out_path), exist_ok=True)

    for cmd in [
        ["wget", "-O", out_path, "-c", "--show-progress", url],
        ["curl", "-L", "-o", out_path, url],
    ]:
        try:
            subprocess.check_call(cmd, timeout=timeout)
            if os.path.isfile(out_path) and os.path.getsize(out_path) > 1024:
                print(f"[{label}] Downloaded: {out_path} ({file_size_gb(out_path):.2f} GB)")
                return out_path
        except Exception:
            continue

    print(f"[{label}] ERROR: download failed")
    return None

def get_llama_quantize():
    # Try existing
    for root, _, files in os.walk(LLAMA_BIN_DIR):
        if "llama-quantize" in files:
            p = os.path.join(root, "llama-quantize")
            if os.access(p, os.X_OK):
                print(f"[BIN] Found: {p}")
                return p

    # Download
    tar_path = os.path.join(WORK_DIR, LLAMA_TAR)
    ok = download_file(LLAMA_URL, tar_path, "BIN", timeout=1200)
    if not ok:
        return None

    os.makedirs(LLAMA_BIN_DIR, exist_ok=True)

    extracted = False
    for strip in [True, False]:
        try:
            cmd = ["tar", "xzf", tar_path, "-C", LLAMA_BIN_DIR]
            if strip:
                cmd.append("--strip-components=1")
            subprocess.check_call(cmd)
            extracted = True
            break
        except subprocess.CalledProcessError:
            continue

    safe_remove(tar_path)

    if not extracted:
        print("[BIN] ERROR: extraction failed.")
        return None

    for root, _, files in os.walk(LLAMA_BIN_DIR):
        if "llama-quantize" in files:
            p = os.path.join(root, "llama-quantize")
            os.chmod(p, 0o755)
            print(f"[BIN] Ready: {p}")
            return p

    print("[BIN] ERROR: llama-quantize not found.")
    return None

def normalize_output_basename(input_path):
    base = os.path.splitext(os.path.basename(input_path))[0]
    known_suffixes = [
        ".Q8_0", "-Q8_0", ".F16", "-F16", ".BF16", "-BF16", ".F32", "-F32",
        ".Q4_0", "-Q4_0", ".Q4_K_M", "-Q4_K_M", ".Q5_K_M", "-Q5_K_M",
        ".IQ4_XS", "-IQ4_XS", ".IQ4_NL", "-IQ4_NL",
    ]
    for suf in known_suffixes:
        if base.endswith(suf):
            base = base[:-len(suf)]
            break
    return base

def regex_escape_literal(s: str) -> str:
    return re.sub(r'([.^$*+?{}\[\]\\|()])', r'\\\1', s)

def normalize_qtype(s):
    if s is None:
        return None
    s = str(s).strip()
    s = s.replace("GGML_TYPE_", "").replace("ggml_type_", "").replace("GGUF_TYPE_", "")
    return s.upper()

def seems_quantized_from_types(type_list):
    # If any tensor type starts with Q or IQ, treat as already-quantized.
    for t in type_list:
        if not t:
            continue
        u = str(t).upper()
        if u.startswith("Q") or u.startswith("IQ") or u.startswith("TQ") or u.startswith("KQ"):
            return True
    return False

# ============================================================================
# GGUF INSPECTION
# ============================================================================

def read_gguf_tensors_basic(gguf_path):
    """
    Returns list of dicts: {name, type_str, n_bytes, data_offset}
    """
    ensure_python_pkg("gguf")
    from gguf import GGUFReader

    r = GGUFReader(gguf_path)
    out = []
    for t in r.tensors:
        typ = None
        for a in ("tensor_type", "type", "ggml_type", "data_type"):
            if hasattr(t, a):
                typ = getattr(t, a)
                break
        typ_s = normalize_qtype(typ.name if hasattr(typ, "name") else typ)

        n_bytes = getattr(t, "n_bytes", None)
        data_offset = getattr(t, "data_offset", None)

        out.append({
            "name": t.name,
            "type": typ_s,
            "n_bytes": int(n_bytes) if n_bytes is not None else None,
            "data_offset": int(data_offset) if data_offset is not None else None,
        })

    del r
    gc.collect()
    return out

def get_tensor_names(gguf_path):
    return [x["name"] for x in read_gguf_tensors_basic(gguf_path)]

def get_tensor_types_map(gguf_path):
    items = read_gguf_tensors_basic(gguf_path)
    return {x["name"]: x["type"] for x in items}

# ============================================================================
# PROFILE EXTRACTION (from Unsloth ref GGUF)
# ============================================================================

def extract_profile_from_ref(ref_path):
    """
    Build a compact rule-set that reproduces Unsloth per-tensor types.

    Output JSON structure:
    {
      "ref_file": "...",
      "token_embd_type": "Q8_0",
      "output_type": "Q6_K",
      "rules": [
          {"pattern": "^blk\\.\\d+\\.attn_q\\.weight$", "type": "Q8_0", "kind": "blk_all_layers"},
          {"pattern": "^blk\\.(?:0|1|2)\\.ffn_down_exps\\.weight$", "type": "IQ4_NL", "kind": "blk_subset_layers"},
          ...
      ],
      "stats": {...}
      ["tensor_types_by_name": {...}]  # optional
    }
    """
    types = get_tensor_types_map(ref_path)

    token_embd_type = types.get("token_embd.weight")
    output_type = types.get("output.weight")

    # group block tensors by (suffix,type)->layers
    blk_re = re.compile(r"^blk\.(\d+)\.(.+)$")
    by_suffix_type_layers = {}  # (suffix, type) -> set(layers)
    globals_map = {}            # name -> type
    max_layer = -1

    for name, typ in types.items():
        m = blk_re.match(name)
        if not m:
            globals_map[name] = typ
            continue
        layer = int(m.group(1))
        suffix = m.group(2)
        max_layer = max(max_layer, layer)
        by_suffix_type_layers.setdefault((suffix, typ), set()).add(layer)

    n_layers = max_layer + 1 if max_layer >= 0 else 0

    rules = []

    # Global exact rules (but skip token_embd/output since we’ll pass dedicated flags)
    for name, typ in sorted(globals_map.items()):
        if name in ("token_embd.weight", "output.weight"):
            continue
        if typ is None:
            continue
        rules.append({
            "pattern": "^" + regex_escape_literal(name) + "$",
            "type": typ,
            "kind": "global_exact",
        })

    # Block rules
    for (suffix, typ), layers in sorted(by_suffix_type_layers.items(), key=lambda x: (x[0][0], str(x[0][1]))):
        if typ is None:
            continue
        layers = sorted(layers)
        suffix_esc = regex_escape_literal(suffix)

        if n_layers and len(layers) == n_layers:
            pat = r"^blk\.\d+\." + suffix_esc + r"$"
            kind = "blk_all_layers"
        else:
            alt = "|".join(str(i) for i in layers)
            pat = r"^blk\.(?:" + alt + r")\." + suffix_esc + r"$"
            kind = "blk_subset_layers"

        rules.append({
            "pattern": pat,
            "type": typ,
            "kind": kind,
        })

    prof = {
        "ref_file": os.path.basename(ref_path),
        "created_unix": int(time.time()),
        "token_embd_type": token_embd_type,
        "output_type": output_type,
        "rules": rules,
        "stats": {
            "n_tensors_total": len(types),
            "n_layers_detected": n_layers,
            "n_rules": len(rules),
        },
    }

    if STORE_RAW_TENSOR_MAP:
        prof["tensor_types_by_name"] = types

    return prof

def load_or_build_profile_json():
    # If profile JSON exists, use it and do NOT download ref again.
    if os.path.isfile(PROFILE_JSON) and os.path.getsize(PROFILE_JSON) > 100:
        with open(PROFILE_JSON, "r", encoding="utf-8") as f:
            prof = json.load(f)
        print(f"[PROFILE] Loaded: {PROFILE_JSON}")
        return prof

    print("[PROFILE] JSON not found; building from Unsloth reference GGUF...")

    if not os.path.isfile(REF_FILE):
        download_file(REF_URL, REF_FILE, "REF", timeout=6*3600)

    prof = extract_profile_from_ref(REF_FILE)

    with open(PROFILE_JSON, "w", encoding="utf-8") as f:
        json.dump(prof, f, ensure_ascii=False, indent=2)

    print(f"[PROFILE] Saved: {PROFILE_JSON}")

    if DELETE_REF_AFTER_PROFILE:
        print(f"[PROFILE] Deleting Unsloth reference quant to save disk: {REF_FILE}")
        safe_remove(REF_FILE)

    return prof

# ============================================================================
# QUANTIZE
# ============================================================================

def filter_rules_to_existing_tensors(rules, tensor_names):
    """
    Keep only rules whose regex matches at least one tensor name in your input model.
    This reduces command size and avoids useless rules.
    """
    kept = []
    names = tensor_names

    for r in rules:
        pat = r.get("pattern")
        typ = r.get("type")
        if not pat or not typ:
            continue
        try:
            rx = re.compile(pat)
        except re.error:
            continue
        if any(rx.search(n) for n in names):
            kept.append(r)

    return kept

def quantize_with_profile(input_path, output_path, profile, imatrix_path=None):
    bin_path = get_llama_quantize()
    if not bin_path:
        raise RuntimeError("llama-quantize not available")

    if not os.path.isfile(input_path):
        raise FileNotFoundError(f"Input not found: {input_path}")

    # Disk check (quantization creates output + temp buffers)
    in_gb = file_size_gb(input_path)
    free_gb = disk_free_gb()
    need_gb = max(12.0, in_gb * 0.9)
    if free_gb < need_gb:
        raise RuntimeError(f"Not enough disk. Need ~{need_gb:.1f} GB free. Free={free_gb:.1f} GB")

    # Decide if allow-requantize is needed (input already quantized)
    input_types_map = get_tensor_types_map(input_path)
    allow_requantize = seems_quantized_from_types(list(input_types_map.values()))

    input_names = list(input_types_map.keys())

    rules = profile.get("rules", [])
    rules = filter_rules_to_existing_tensors(rules, input_names)

    env = os.environ.copy()
    bin_dir = os.path.dirname(bin_path)
    ld = env.get("LD_LIBRARY_PATH", "")
    env["LD_LIBRARY_PATH"] = f"{bin_dir}:{ld}" if ld else bin_dir

    cmd = [bin_path]

    if allow_requantize:
        cmd += ["--allow-requantize"]

    if imatrix_path and os.path.isfile(imatrix_path):
        cmd += ["--imatrix", imatrix_path]
        cmd += ["--exclude-weights", "output.weight"]

    # Apply dedicated types only if corresponding tensors exist in input
    if "token_embd.weight" in input_types_map and profile.get("token_embd_type"):
        cmd += ["--token-embedding-type", str(profile["token_embd_type"])]
    if "output.weight" in input_types_map and profile.get("output_type"):
        cmd += ["--output-tensor-type", str(profile["output_type"])]

    for r in rules:
        cmd += ["--tensor-type", f"{r['pattern']}={r['type']}"]

    cmd += [input_path, output_path, TARGET_BASE_TYPE]

    print("\n" + "=" * 78)
    print("QUANTIZATION")
    print("=" * 78)
    print(f"[IN ] {input_path} ({in_gb:.2f} GB)")
    print(f"[OUT] {output_path}")
    print(f"[BASE] {TARGET_BASE_TYPE}")
    print(f"[RULES] {len(rules)} (filtered to your model)")
    print(f"[IMX] {imatrix_path if (imatrix_path and os.path.isfile(imatrix_path)) else 'none'}")
    print(f"[FLAG] allow_requantize = {allow_requantize}")
    print(f"[SYS] free disk = {free_gb:.1f} GB | ram used = {ram_used_gb():.2f} GB")

    print("[CMD] (showing head)")
    print("  " + " \\\n  ".join(cmd[:40]) + (" \\\n  ... (many --tensor-type rules) ..." if len(cmd) > 40 else ""))

    t0 = time.time()
    proc = subprocess.Popen(
        cmd,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        env=env,
        bufsize=1,
    )
    for line in proc.stdout:
        line = line.rstrip()
        if line:
            print(" ", line)
    proc.wait()

    if proc.returncode != 0:
        raise RuntimeError(f"llama-quantize failed (code={proc.returncode})")

    dt = time.time() - t0
    out_gb = file_size_gb(output_path)
    print(f"[DONE] {dt/60:.1f} min | size {in_gb:.2f} GB -> {out_gb:.2f} GB")

def verify_output(output_path):
    print("[VERIFY] Checking output GGUF parse + quick zero-fill sample...")
    try:
        ensure_python_pkg("gguf")
        from gguf import GGUFReader

        r = GGUFReader(output_path)
        tensors = list(r.tensors)
        print(f"[VERIFY] Tensors: {len(tensors)}")

        checked = 0
        zeroish = 0

        with open(output_path, "rb") as f:
            for t in tensors[:60]:
                n_bytes = getattr(t, "n_bytes", None)
                off = getattr(t, "data_offset", None)
                if n_bytes is None or off is None:
                    continue
                n = min(4096, int(n_bytes))
                if n < 128:
                    continue
                f.seek(int(off))
                raw = f.read(n)
                arr = np.frombuffer(raw, dtype=np.uint8)
                if arr.size and arr.sum() == 0:
                    zeroish += 1
                checked += 1

        del r
        gc.collect()

        if zeroish:
            print(f"[VERIFY] WARNING: {zeroish}/{checked} sampled tensors looked zero-filled.")
        else:
            print(f"[VERIFY] OK: sampled {checked} tensors.")
        return True
    except Exception as e:
        print(f"[VERIFY] ERROR: {e}")
        return False


# ============================================================================
# MAIN
# ============================================================================

def main():
    print("=" * 78)
    print("UNSLOTH UD-IQ4_XS PROFILE -> JSON -> DELETE REF -> QUANTIZE YOUR MODEL")
    print("=" * 78)
    print(f"[INFO] Work dir:  {WORK_DIR}")
    print(f"[INFO] Input:     {INPUT_FILE}")
    print(f"[INFO] Profile:   {PROFILE_JSON}")
    print(f"[INFO] Target:    {TARGET_BASE_TYPE}")
    print(f"[INFO] Disk free: {disk_free_gb():.2f} GB | RAM used: {ram_used_gb():.2f} GB")

    if not os.path.isfile(INPUT_FILE):
        sys.exit(f"[ERROR] Your input GGUF not found: {INPUT_FILE}")

    # iMatrix is small enough to keep; download if missing
    download_file(IMATRIX_URL, IMATRIX_FILE, "IMATRIX", timeout=1800)

    # Build/load JSON profile (this step deletes REF_FILE if it had to download it)
    profile = load_or_build_profile_json()

    if PRINT_PROFILE_SUMMARY:
        print("[PROFILE] Summary:")
        print(json.dumps(profile.get("stats", {}), indent=2, ensure_ascii=False))
        print(f"[PROFILE] token_embd_type: {profile.get('token_embd_type')}")
        print(f"[PROFILE] output_type:     {profile.get('output_type')}")
        print(f"[PROFILE] rules_total:     {len(profile.get('rules', []))}")

    base = normalize_output_basename(INPUT_FILE)
    output_path = os.path.join(WORK_DIR, f"{base}-{OUTPUT_TAG}-{TARGET_BASE_TYPE}.gguf")

    quantize_with_profile(INPUT_FILE, output_path, profile, IMATRIX_FILE)

    if VERIFY_AFTER_QUANT:
        print("\n" + "=" * 78)
        verify_output(output_path)

    print("\n" + "=" * 78)
    print("COMPLETE")
    print("=" * 78)
    print(f"[OUT] {output_path}")
    print(f"[OUT] Size: {file_size_gb(output_path):.2f} GB")
    print(f"[SYS] Free disk: {disk_free_gb():.2f} GB")

    print("\n[FS]")
    for fn in sorted(os.listdir(WORK_DIR)):
        fp = os.path.join(WORK_DIR, fn)
        if os.path.isfile(fp):
            try:
                print(f"  {fn}  ({file_size_gb(fp):.2f} GB)")
            except Exception:
                print(f"  {fn}")

if __name__ == "__main__":
    main()