#!/usr/bin/env python3
"""
model_manager.py — Multi-model manager for local LLM API

Directory layout (new):
  model/<ModelName>/               ← per-model directory
  model/<ModelName>/chunks/        ← split GGUF shards (100MB each)

Commands:
  list                          List all configured models + disk status
  quants <name|repo>            Query HuggingFace API — all GGUF files + sizes
  use <name> [quant]            Add model to active list (max 2 run at once)
  use <name-quant>              Combined format — e.g. use Qwen3.5-4B-IQ4_NL
  use <name-quant.gguf>         Full filename — e.g. use Qwen3.5-4B-IQ4_NL.gguf
  remove <name>                 Remove model from active list
  info                          Show full install.json config

Usage:
  python model_manager.py list
  python model_manager.py quants Qwen3.5-4B
  python model_manager.py quants unsloth/Phi-4-mini-reasoning-GGUF
  python model_manager.py use Qwen3.5-4B
  python model_manager.py use DeepSeek-R1-8B
  python model_manager.py use Qwen3.5-4B IQ4_NL
  python model_manager.py remove DeepSeek-R1-8B
  python model_manager.py info
"""

import os, sys, json, urllib.request, urllib.error
from typing import Optional

_HERE = os.path.dirname(os.path.abspath(__file__))
_CFG  = os.path.join(_HERE, "install.json")


# ── Config helpers ────────────────────────────────────────────────────────────

def load_cfg() -> dict:
    try:
        with open(_CFG) as f:
            return json.load(f)
    except Exception as e:
        print(f"[error] Cannot read install.json: {e}")
        sys.exit(1)


def save_cfg(cfg: dict) -> None:
    with open(_CFG, "w") as f:
        json.dump(cfg, f, indent=2, ensure_ascii=False)
        f.write("\n")


def get_active_models(cfg: dict) -> list:
    """Return list of active model names. Supports both new (active_models[])
    and legacy (active_model string) config format."""
    v = cfg.get("active_models")
    if isinstance(v, list):
        return [x for x in v if x]
    s = cfg.get("active_model", "")
    return [s] if s else []


def model_dir_for(cfg: dict, name: str) -> str:
    """Return per-model directory: ./model/<name>/"""
    base = cfg.get("model_dir", "./model")
    return os.path.join(base, name)


def gguf_dir(cfg: dict) -> str:
    """Legacy helper — returns old model_gguf dir (for migration checks)."""
    return cfg.get("gguf", {}).get("dir", "./model_gguf")


# ── Disk helpers ──────────────────────────────────────────────────────────────

def _file_status(cfg: dict, model_name: str) -> dict:
    """Return disk info for a model: file exists, size, chunks."""
    mc  = cfg.get("models", {}).get(model_name, {})
    fn  = mc.get("file", "")
    d   = model_dir_for(cfg, model_name)
    fp  = os.path.join(d, fn)
    chk = os.path.join(d, "chunks")

    status = {"file": fn, "exists": False, "size_gb": None, "chunks": 0, "chunk_total_gb": None}

    if os.path.isfile(fp):
        status["exists"] = True
        status["size_gb"] = round(os.path.getsize(fp) / 1024**3, 2)

    if os.path.isdir(chk):
        base = os.path.splitext(fn)[0]
        splits = [f for f in os.listdir(chk) if f.startswith(base) and f.endswith(".gguf")]
        parts  = [f for f in os.listdir(chk) if f.startswith(base) and f.endswith(".part")]
        chunks = splits or parts
        if chunks:
            total = sum(os.path.getsize(os.path.join(chk, f)) for f in chunks)
            status["chunks"] = len(chunks)
            status["chunk_total_gb"] = round(total / 1024**3, 2)

    return status


# ── HuggingFace API ───────────────────────────────────────────────────────────

def hf_list_gguf(repo: str, token: Optional[str] = None) -> list:
    url = f"https://huggingface.co/api/models/{repo}/tree/main"
    req = urllib.request.Request(url)
    req.add_header("Accept", "application/json")
    if token:
        req.add_header("Authorization", f"Bearer {token}")
    try:
        with urllib.request.urlopen(req, timeout=15) as resp:
            data = json.loads(resp.read().decode())
    except urllib.error.HTTPError as e:
        raise RuntimeError(f"HF API error {e.code}: {url}")
    except Exception as e:
        raise RuntimeError(f"Network error: {e}")

    results = []
    for entry in data:
        path = entry.get("path", "")
        if not path.endswith(".gguf"):
            continue
        if os.path.basename(path).startswith("mmproj"):
            continue
        size  = entry.get("size", 0)
        fname = os.path.basename(path)
        quant = _parse_quant(fname)
        bits  = _quant_bits(quant)
        results.append({
            "name":       fname,
            "size_bytes": size,
            "size_gb":    round(size / 1024**3, 2),
            "quant":      quant,
            "bits":       bits,
        })

    results.sort(key=lambda x: x["size_bytes"])
    return results


def _parse_quant(filename: str) -> str:
    base  = filename.replace(".gguf", "")
    parts = base.split("-")
    for i in range(len(parts) - 1, 0, -1):
        candidate = "-".join(parts[i:])
        if any(candidate.upper().startswith(p) for p in
               ("Q2","Q3","Q4","Q5","Q6","Q8","IQ","UD","BF","F1")):
            return candidate
    return base.split("-")[-1]


def _quant_bits(quant: str) -> str:
    q = quant.upper()
    if q.startswith("BF16") or q.startswith("F16"):  return "16-bit"
    if q.startswith("Q8")   or q.startswith("UD-Q8"): return "8-bit"
    if q.startswith("Q6")   or q.startswith("UD-Q6"): return "6-bit"
    if q.startswith("Q5")   or q.startswith("UD-Q5"): return "5-bit"
    if q.startswith("Q4")   or q.startswith("IQ4") or q.startswith("UD-Q4"): return "4-bit"
    if q.startswith("Q3")   or q.startswith("IQ3") or q.startswith("UD-Q3"): return "3-bit"
    if q.startswith("Q2")   or q.startswith("IQ2") or q.startswith("UD-Q2"): return "2-bit"
    if q.startswith("IQ1")  or q.startswith("UD-IQ1"): return "1-bit"
    return "?"


# ── Commands ──────────────────────────────────────────────────────────────────

def cmd_list(cfg: dict) -> None:
    models  = cfg.get("models", {})
    actives = get_active_models(cfg)

    print(f"\n{'Model':<20} {'Repo':<45} {'File':<40} {'Disk'}")
    print("─" * 135)

    for name, mc in models.items():
        st    = _file_status(cfg, name)
        mark  = "▶ " if name in actives else "  "
        idx   = actives.index(name) + 1 if name in actives else 0
        label = f"[{idx}]" if name in actives else "   "
        repo  = mc.get("repo", "")
        fn    = mc.get("file", "")
        note  = mc.get("note", "")

        if st["exists"]:
            disk = f"✓ {st['size_gb']}GB"
        elif st["chunks"] > 0:
            disk = f"✓ {st['chunks']} chunks ({st['chunk_total_gb']}GB)"
        else:
            disk = "✗ not downloaded"

        print(f"{mark}{label} {name:<16} {repo:<45} {fn:<40} {disk}")
        if note:
            print(f"  {'':24} {note}")

    model_base = cfg.get("model_dir", "./model")
    print(f"\nModel dir : {os.path.abspath(model_base)}")
    print(f"Active    : {actives}  (max 2 run simultaneously)")
    print(f"\nAdd model    : python model_manager.py use <ModelName>")
    print(f"Remove model : python model_manager.py remove <ModelName>")
    print(f"Change quant : python model_manager.py use <ModelName> <QUANT>")
    print(f"List quants  : python model_manager.py quants <ModelName>")


def cmd_quants(cfg: dict, target: str) -> list:
    models = cfg.get("models", {})

    if "/" in target:
        repo = target
        name = target.split("/")[-1]
    elif target in models:
        repo = models[target].get("repo", "")
        name = target
        if not repo:
            print(f"[error] No repo configured for '{target}'")
            sys.exit(1)
    else:
        matches = [n for n in models if target.lower() in n.lower()]
        if len(matches) == 1:
            name = matches[0]
            repo = models[name].get("repo", "")
        elif len(matches) > 1:
            print(f"[error] Ambiguous: {matches}. Be more specific.")
            sys.exit(1)
        else:
            print(f"[error] Unknown model '{target}'. Run: python model_manager.py list")
            sys.exit(1)

    print(f"\nQuerying HuggingFace: {repo} ...")
    try:
        files = hf_list_gguf(repo)
    except RuntimeError as e:
        print(f"[error] {e}")
        sys.exit(1)

    if not files:
        print("[warn] No .gguf files found in this repo.")
        return []

    active_file = models.get(name, {}).get("file", "") if name in models else ""

    by_bits: dict = {}
    for f in files:
        by_bits.setdefault(f["bits"], []).append(f)

    bit_order = ["1-bit","2-bit","3-bit","4-bit","5-bit","6-bit","8-bit","16-bit","?"]

    print(f"\n{'Bits':<8} {'Quantization':<25} {'Size':>7}   {'File'}")
    print("─" * 90)
    for bits in bit_order:
        if bits not in by_bits:
            continue
        for f in by_bits[bits]:
            mark = " ◀ active" if f["name"] == active_file else ""
            print(f"  {bits:<6} {f['quant']:<25} {f['size_gb']:>5.2f}GB   {f['name']}{mark}")

    print(f"\nTo use a specific quant (any of these work):")
    print(f"  python model_manager.py use {name} IQ4_NL")
    print(f"  python model_manager.py use {name}-IQ4_NL")
    print(f"  python model_manager.py use {name}-IQ4_NL.gguf")
    return files


def _split_model_quant(name: str, models: dict) -> tuple:
    stem = name[:-5] if name.lower().endswith(".gguf") else name
    for mname in sorted(models.keys(), key=len, reverse=True):
        if stem.lower().startswith(mname.lower()):
            remainder = stem[len(mname):]
            if remainder.startswith("-") and len(remainder) > 1:
                return mname, remainder[1:]
    return name, None


def cmd_use(cfg: dict, model_name: str, quant: Optional[str] = None) -> None:
    """Add model to active_models list (max 2). If already active, update quant only."""
    models = cfg.get("models", {})

    if model_name not in models:
        parsed_model, parsed_quant = _split_model_quant(model_name, models)
        if parsed_model in models:
            model_name = parsed_model
            if parsed_quant and not quant:
                quant = parsed_quant
        else:
            matches = [n for n in models if model_name.lower() in n.lower()]
            if len(matches) == 1:
                model_name = matches[0]
            elif len(matches) > 1:
                print(f"[error] Ambiguous name '{model_name}': {matches}")
                sys.exit(1)
            else:
                print(f"[error] Unknown model '{model_name}'")
                print(f"  Known: {', '.join(models.keys())}")
                sys.exit(1)

    mc   = models[model_name]
    repo = mc.get("repo", "")

    if quant:
        print(f"Resolving {quant} from {repo} ...")
        try:
            files = hf_list_gguf(repo)
        except RuntimeError as e:
            print(f"[error] Cannot query HF API: {e}")
            sys.exit(1)

        q_up  = quant.upper()
        match = next((f for f in files if f["quant"].upper() == q_up), None)
        if not match:
            partials = [f for f in files if q_up in f["quant"].upper()]
            if len(partials) == 1:
                match = partials[0]
            elif partials:
                print(f"[warn] Multiple matches for '{quant}':")
                for p in partials:
                    print(f"  {p['quant']}  {p['size_gb']}GB  {p['name']}")
                sys.exit(1)
            else:
                print(f"[error] Quant '{quant}' not found in {repo}")
                print(f"  Run: python model_manager.py quants {model_name}")
                sys.exit(1)

        old_file = mc.get("file", "")
        mc["file"] = match["name"]
        print(f"  File : {old_file} → {match['name']}  ({match['size_gb']}GB)")

    # Manage active_models list (max 2)
    actives = get_active_models(cfg)
    if model_name not in actives:
        if len(actives) >= 2:
            removed = actives.pop(0)
            print(f"  Removed '{removed}' from active list (max 2 slots)")
        actives.append(model_name)
    # Remove legacy field
    cfg.pop("active_model", None)
    cfg["active_models"] = actives
    save_cfg(cfg)

    st = _file_status(cfg, model_name)
    print(f"\n✓ Active models: {actives}")
    print(f"  Added   : {model_name}")
    print(f"  Repo    : {repo}")
    print(f"  File    : {mc.get('file', '')}")
    if st["exists"]:
        print(f"  Disk    : ✓ {st['size_gb']}GB — ready")
    elif st["chunks"] > 0:
        print(f"  Disk    : ✓ {st['chunks']} split chunks ({st['chunk_total_gb']}GB) — ready")
    else:
        d = model_dir_for(cfg, model_name)
        print(f"  Disk    : ✗ Not downloaded")
        print(f"  Download: huggingface-cli download {repo} {mc.get('file','')} --local-dir {os.path.abspath(d)}/")
    print(f"\nRestart server to apply: python app.py")


def cmd_remove(cfg: dict, model_name: str) -> None:
    """Remove a model from the active_models list."""
    actives = get_active_models(cfg)
    if model_name not in actives:
        # Fuzzy match
        matches = [n for n in actives if model_name.lower() in n.lower()]
        if len(matches) == 1:
            model_name = matches[0]
        else:
            print(f"[warn] '{model_name}' is not in active list: {actives}")
            return

    actives.remove(model_name)
    cfg.pop("active_model", None)
    cfg["active_models"] = actives
    save_cfg(cfg)
    print(f"✓ Removed '{model_name}' from active list.")
    print(f"  Active now: {actives}")
    print(f"  Restart server to apply: python app.py")


def cmd_info(cfg: dict) -> None:
    print(json.dumps(cfg, indent=2, ensure_ascii=False))


def cmd_sync(cfg: dict, clean: bool = False) -> str:
    """
    Check disk vs active models. If --clean, delete stale model files.
    Prints READY if all active models are on disk, else DOWNLOAD.
    """
    actives = get_active_models(cfg)
    models  = cfg.get("models", {})

    if not actives:
        print("[sync] No active_models configured")
        print("DOWNLOAD")
        return "DOWNLOAD"

    all_ready = True

    for active in actives:
        if not active or active not in models:
            print(f"[sync] Unknown model '{active}' — skipping")
            all_ready = False
            continue

        mc            = models[active]
        expected_file = mc.get("file", "")
        expected_base = os.path.splitext(expected_file)[0]
        d             = model_dir_for(cfg, active)
        chunk_dir     = os.path.join(d, "chunks")

        # Catalogue everything in this model's dir
        all_files: list = []
        if os.path.isdir(d):
            for f in os.listdir(d):
                if f.endswith(".gguf"):
                    all_files.append((os.path.join(d, f), f == expected_file))

        if os.path.isdir(chunk_dir):
            for f in os.listdir(chunk_dir):
                if f.endswith(".gguf") or f.endswith(".part"):
                    mine = f.startswith(expected_base + "-") or f.startswith(expected_base + ".")
                    all_files.append((os.path.join(chunk_dir, f), mine))

        # Delete stale files (files that don't belong to expected_file)
        stale = [p for p, mine in all_files if not mine]
        if stale and clean:
            for sp in stale:
                try:
                    os.remove(sp)
                    print(f"[sync] Deleted stale: {os.path.basename(sp)}")
                except Exception as e:
                    print(f"[sync] Warning — could not delete {os.path.basename(sp)}: {e}")
        elif stale:
            for sp in stale:
                print(f"[sync] Stale (use --clean to remove): {os.path.basename(sp)}")

        # Check if this model is ready
        active_whole = os.path.join(d, expected_file)
        active_ready = os.path.isfile(active_whole)
        if not active_ready and os.path.isdir(chunk_dir):
            chunks = [
                f for f in os.listdir(chunk_dir)
                if (f.startswith(expected_base + "-") or f.startswith(expected_base + "."))
                and (f.endswith(".gguf") or f.endswith(".part"))
            ]
            active_ready = len(chunks) > 0

        if active_ready:
            st = _file_status(cfg, active)
            if st["exists"]:
                print(f"[sync] {active} → {expected_file} ({st['size_gb']}GB) — ready")
            elif st["chunks"] > 0:
                print(f"[sync] {active} → {st['chunks']} chunks ({st['chunk_total_gb']}GB) — ready")
        else:
            print(f"[sync] {active} → {expected_file} — not on disk")
            all_ready = False

    if all_ready:
        print("READY")
        return "READY"
    else:
        print("DOWNLOAD")
        return "DOWNLOAD"


def _download_one(cfg: dict, active: str) -> bool:
    """Download a single model. Returns True on success."""
    models = cfg.get("models", {})
    if not active or active not in models:
        print(f"[download] Unknown model '{active}'")
        return False

    mc   = models[active]
    repo = mc.get("repo", "")
    file = mc.get("file", "")
    d    = model_dir_for(cfg, active)

    if not repo or not file:
        print(f"[download] Model '{active}' has no repo/file configured")
        return False

    os.makedirs(d, exist_ok=True)
    out_path  = os.path.join(d, file)
    base      = os.path.splitext(file)[0]
    chunk_dir = os.path.join(d, "chunks")

    if os.path.isfile(out_path):
        sz = round(os.path.getsize(out_path) / 1024**3, 2)
        print(f"[download] {active}: Already on disk — {file} ({sz}GB)")
        return True

    if os.path.isdir(chunk_dir):
        chunks = [
            f for f in os.listdir(chunk_dir)
            if (f.startswith(base + "-") or f.startswith(base + "."))
            and (f.endswith(".gguf") or f.endswith(".part"))
        ]
        if chunks:
            print(f"[download] {active}: Chunks already on disk ({len(chunks)} files) — skip")
            return True

    note = mc.get("note", "")
    print(f"[download] Downloading {active}: {file}" + (f"  ({note})" if note else ""))
    print(f"[download]   Repo : {repo}")
    print(f"[download]   Dest : {os.path.abspath(d)}/")
    sys.stdout.flush()

    try:
        from huggingface_hub import hf_hub_download
    except ImportError:
        import subprocess
        print("[download] Installing huggingface-hub ...", flush=True)
        subprocess.run(
            [sys.executable, "-m", "pip", "install", "-q", "--no-cache-dir", "huggingface-hub"],
            capture_output=True)
        from huggingface_hub import hf_hub_download

    try:
        path = hf_hub_download(repo_id=repo, filename=file, local_dir=d)
        sz   = round(os.path.getsize(path) / 1024**3, 2)
        print(f"[download] {active}: Done — {file} ({sz}GB)")
        return True
    except Exception as e:
        print(f"[error] {active}: Download failed: {e}", file=sys.stderr)
        print(f"  Manual: huggingface-cli download {repo} {file} --local-dir {os.path.abspath(d)}/",
              file=sys.stderr)
        return False


def cmd_download(cfg: dict) -> None:
    """Download all active models that are not yet on disk."""
    actives = get_active_models(cfg)
    if not actives:
        print("[error] No active_models in install.json")
        sys.exit(1)

    any_failed = False
    for active in actives:
        ok = _download_one(cfg, active)
        if not ok:
            any_failed = True

    if any_failed:
        sys.exit(1)


# ── Flask-callable helpers (imported by app.py) ───────────────────────────────

def get_models_status(cfg: Optional[dict] = None) -> dict:
    if cfg is None:
        cfg = load_cfg()
    models  = cfg.get("models", {})
    actives = get_active_models(cfg)
    result  = {}
    for name, mc in models.items():
        st = _file_status(cfg, name)
        result[name] = {
            "repo":           mc.get("repo", ""),
            "file":           mc.get("file", ""),
            "note":           mc.get("note", ""),
            "active":         name in actives,
            "active_slot":    actives.index(name) + 1 if name in actives else None,
            "on_disk":        st["exists"] or st["chunks"] > 0,
            "size_gb":        st["size_gb"],
            "chunks":         st["chunks"],
            "chunk_total_gb": st["chunk_total_gb"],
        }
    base = cfg.get("model_dir", "./model")
    return {
        "active_models": actives,
        "models":        result,
        "model_dir":     os.path.abspath(base),
    }


def get_quants_for_model(name_or_repo: str, cfg: Optional[dict] = None) -> list:
    if cfg is None:
        cfg = load_cfg()
    models = cfg.get("models", {})
    if "/" in name_or_repo:
        repo = name_or_repo
    elif name_or_repo in models:
        repo = models[name_or_repo].get("repo", "")
    else:
        matches = [n for n in models if name_or_repo.lower() in n.lower()]
        repo = models[matches[0]].get("repo", "") if matches else name_or_repo
    return hf_list_gguf(repo)


# ── CLI entry point ───────────────────────────────────────────────────────────

def main():
    args = sys.argv[1:]
    if not args:
        print(__doc__)
        sys.exit(0)

    cmd = args[0].lower()
    cfg = load_cfg()

    if cmd == "list":
        cmd_list(cfg)

    elif cmd == "quants":
        if len(args) < 2:
            print("[error] Usage: python model_manager.py quants <ModelName|repo>")
            sys.exit(1)
        cmd_quants(cfg, args[1])

    elif cmd == "use":
        if len(args) < 2:
            print("[error] Usage: python model_manager.py use <ModelName> [QUANT]")
            sys.exit(1)
        quant = args[2] if len(args) >= 3 else None
        cmd_use(cfg, args[1], quant)

    elif cmd == "remove":
        if len(args) < 2:
            print("[error] Usage: python model_manager.py remove <ModelName>")
            sys.exit(1)
        cmd_remove(cfg, args[1])

    elif cmd == "info":
        cmd_info(cfg)

    elif cmd == "sync":
        clean = "--clean" in args
        cmd_sync(cfg, clean=clean)

    elif cmd == "download":
        cmd_download(cfg)

    else:
        print(f"[error] Unknown command: {cmd}")
        print("  Commands: list | quants | use | remove | info | sync | download")
        sys.exit(1)


if __name__ == "__main__":
    main()