#!/usr/bin/python3
# models_data.py: the data layer behind RecapMuse's "Models" view (on-device local model manager).
#
# Enumerates the local AI models actually installed on this Mac (speech-to-text, speaker diarization,
# and local LLMs), reports each one's size on disk, the total reclaimable space, the hardware, and a
# recommendation. Also exposes a curated catalog of models that CAN be downloaded. The Swift shell
# injects the JSON via window.setState({models:...}); the Models view renders it.
#
# Pure stdlib so the system /usr/bin/python3 can run it (same pattern as preflight.py / people_data.py).
#
#   models_data.py --json                 # inventory (default; reads only)
#   models_data.py --remove <id>          # delete an installed model to reclaim space (re-downloadable)
#   models_data.py --catalog              # just the downloadable catalog
#
# SAFETY: --remove only ever runs when the user clicks "Remove" in the app. Every model it can remove
# is re-downloadable from its source, so removal reclaims space without destroying anything permanent.
import sys, os, json, subprocess, shutil

HOME = os.path.expanduser("~")
HF_HUB = os.path.join(HOME, ".cache/huggingface/hub")
OLLAMA_DIR = os.path.join(HOME, ".ollama/models")


def human(n):
    for unit in ("B", "KB", "MB", "GB", "TB"):
        if n < 1024 or unit == "TB":
            return f"{n:.0f} {unit}" if unit in ("B", "KB") else f"{n:.1f} {unit}"
        n /= 1024.0


def dir_size(path):
    total = 0
    for root, _dirs, files in os.walk(path):
        for f in files:
            try:
                total += os.lstat(os.path.join(root, f)).st_size
            except OSError:
                pass
    return total


def hardware():
    def sysctl(key):
        try:
            return subprocess.run(["/usr/sbin/sysctl", "-n", key], capture_output=True,
                                  text=True, timeout=5).stdout.strip()
        except Exception:
            return ""
    chip = sysctl("machdep.cpu.brand_string") or "Apple Silicon"
    try:
        ram_gb = round(int(sysctl("hw.memsize")) / 1073741824)
    except ValueError:
        ram_gb = 0
    try:
        st = os.statvfs("/")
        free = st.f_bavail * st.f_frsize
    except OSError:
        free = 0
    return {"chip": chip, "ram_gb": ram_gb, "disk_free_bytes": free, "disk_free_human": human(free)}


def classify_hf(org, model):
    m = model.lower()
    if "parakeet" in m:
        return ("Speech to text", "MLX", "Fast, accurate, English")
    if "whisper" in m:
        return ("Speech to text", "MLX", "Multilingual")
    if "wespeaker" in m or "embed" in m:
        return ("Speaker voiceprints", "pyannote", "Voice embedding")
    if "segmentation" in m or "diariz" in m or "speaker" in m:
        return ("Speaker diarization", "pyannote", "Who-spoke-when")
    return ("Model", org, "")


def installed_models():
    out = []
    # 1) Hugging Face hub cache (Parakeet, Whisper, pyannote, ...)
    if os.path.isdir(HF_HUB):
        for d in sorted(os.listdir(HF_HUB)):
            if not d.startswith("models--"):
                continue
            path = os.path.join(HF_HUB, d)
            parts = d[len("models--"):].split("--")
            org = parts[0] if parts else ""
            model = "--".join(parts[1:]) if len(parts) > 1 else d
            kind, runtime, tier = classify_hf(org, model)
            size = dir_size(path)
            out.append({
                "id": "hf:" + d, "label": model, "kind": kind, "runtime": runtime, "tier": tier,
                "size_bytes": size, "size_human": human(size), "path": path,
                "installed": True, "removable": True, "source": f"huggingface.co/{org}/{model}",
            })
    # 2) Ollama local LLMs
    try:
        r = subprocess.run(["ollama", "list"], capture_output=True, text=True, timeout=8)
        if r.returncode == 0:
            for line in r.stdout.splitlines()[1:]:
                cols = line.split()
                if len(cols) >= 3:
                    name = cols[0]
                    size_str = cols[2] + " " + cols[3] if len(cols) >= 4 and cols[3] in ("GB", "MB", "KB") else cols[2]
                    out.append({
                        "id": "ollama:" + name, "label": name, "kind": "Local LLM (chat / summaries)",
                        "runtime": "Ollama", "tier": "Reasoning", "size_bytes": _parse_size(size_str),
                        "size_human": size_str, "path": OLLAMA_DIR, "installed": True, "removable": True,
                        "source": "ollama.com/library",
                    })
    except Exception:
        pass
    # 3) Apple on-device speech (SpeechAnalyzer, macOS 26) — system-managed, drives the live transcript
    out.append({
        "id": "apple:speechanalyzer", "label": "Apple Speech (on-device)",
        "kind": "Speech to text", "runtime": "Apple Neural Engine", "tier": "Live, zero-setup",
        "size_bytes": 0, "size_human": "system", "path": "", "installed": True, "removable": False,
        "source": "Built into macOS 26", "note": "Powers the live transcript. Managed by macOS.",
    })
    return out


def _parse_size(s):
    try:
        num, unit = s.split()
        mult = {"B": 1, "KB": 1024, "MB": 1024**2, "GB": 1024**3, "TB": 1024**4}.get(unit.upper(), 1)
        return int(float(num) * mult)
    except Exception:
        return 0


# Curated set of models the user can add, with approximate download sizes. installed ones are filtered out.
CATALOG = [
    {"id": "dl:parakeet-tdt-0.6b-v3", "label": "Parakeet TDT 0.6B v3", "kind": "Speech to text",
     "runtime": "MLX", "tier": "Recommended — fastest + most accurate English on Apple Silicon",
     "size_human": "~2.3 GB", "source": "mlx-community/parakeet-tdt-0.6b-v3"},
    {"id": "dl:whisper-large-v3-turbo", "label": "Whisper Large v3 Turbo", "kind": "Speech to text",
     "runtime": "MLX", "tier": "Multilingual (99 languages)", "size_human": "~1.5 GB",
     "source": "mlx-community/whisper-large-v3-turbo"},
    {"id": "dl:whisper-small", "label": "Whisper Small", "kind": "Speech to text", "runtime": "MLX",
     "tier": "Lightweight multilingual", "size_human": "~0.5 GB", "source": "mlx-community/whisper-small-mlx"},
    {"id": "dl:llama3.2-3b", "label": "Llama 3.2 3B", "kind": "Local LLM (chat / summaries)",
     "runtime": "Ollama", "tier": "Small, fast summaries", "size_human": "~2.0 GB", "source": "ollama run llama3.2"},
    {"id": "dl:qwen2.5-7b", "label": "Qwen 2.5 7B", "kind": "Local LLM (chat / summaries)",
     "runtime": "Ollama", "tier": "Balanced summaries + chat", "size_human": "~4.7 GB", "source": "ollama run qwen2.5"},
]


def recommend(hw, models):
    asr = next((m for m in models if "parakeet" in m["label"].lower()), None)
    have_llm = any(m["kind"].startswith("Local LLM") for m in models)
    bits = []
    if asr:
        bits.append("Parakeet is installed and is the best on-device transcriber for your " + hw["chip"] + ".")
    else:
        bits.append("Add Parakeet TDT 0.6B for the fastest, most accurate on-device transcription.")
    bits.append("Apple Speech handles the live transcript with zero setup.")
    if have_llm:
        bits.append("A local LLM is installed, so summaries and chat can run fully on-device.")
    else:
        bits.append("Add a small local LLM (Llama 3.2 3B) to summarize on-device instead of the cloud.")
    return " ".join(bits)


def build():
    hw = hardware()
    models = installed_models()
    installed_bytes = sum(m["size_bytes"] for m in models)
    reclaimable = sum(m["size_bytes"] for m in models if m.get("removable"))
    # Hide catalog entries already installed: match the model identifier (last path/space token of
    # the catalog source, e.g. "parakeet-tdt-0.6b-v3") against the installed labels.
    import re
    installed_labels = " ".join(m["label"].lower() for m in models)
    def already_have(c):
        key = re.split(r"[/\s]", c["source"].lower())[-1]
        return bool(key) and key in installed_labels
    catalog = [c for c in CATALOG if not already_have(c)]
    # sort installed: biggest first within each kind, kinds in a stable order
    kind_order = {"Speech to text": 0, "Speaker diarization": 1, "Speaker voiceprints": 2,
                  "Local LLM (chat / summaries)": 3, "Model": 4}
    models.sort(key=lambda m: (kind_order.get(m["kind"], 9), -m["size_bytes"]))
    return {
        "version": 1,
        "hardware": hw,
        "totals": {
            "installed_bytes": installed_bytes, "installed_human": human(installed_bytes),
            "reclaimable_bytes": reclaimable, "reclaimable_human": human(reclaimable),
            "count": len([m for m in models if m["size_bytes"] > 0]),
        },
        "recommendation": recommend(hw, models),
        "models": models,
        "catalog": catalog,
    }


def remove(model_id):
    """Remove an installed model to reclaim space. Re-downloadable. Only called from the app's Remove button."""
    data = build()
    target = next((m for m in data["models"] if m["id"] == model_id), None)
    if not target:
        return {"ok": False, "error": "not found"}
    if not target.get("removable"):
        return {"ok": False, "error": "not removable (system-managed)"}
    freed = target["size_bytes"]
    if model_id.startswith("ollama:"):
        name = model_id.split(":", 1)[1]
        r = subprocess.run(["ollama", "rm", name], capture_output=True, text=True, timeout=30)
        ok = r.returncode == 0
        return {"ok": ok, "freed_human": human(freed), "label": target["label"],
                "error": "" if ok else r.stderr.strip()[:200]}
    if model_id.startswith("hf:"):
        path = target["path"]
        # safety: only delete inside the HF hub cache
        if os.path.isdir(path) and os.path.realpath(path).startswith(os.path.realpath(HF_HUB) + os.sep):
            shutil.rmtree(path, ignore_errors=True)
            return {"ok": not os.path.exists(path), "freed_human": human(freed), "label": target["label"]}
        return {"ok": False, "error": "path outside model cache; refused"}
    return {"ok": False, "error": "unknown model source"}


PYENV_PY = os.path.expanduser("~/.pyenv/versions/3.12.8/bin/python")


def add(model_id):
    """Download a catalog model. HF repos via huggingface_hub (pyenv python); Ollama via `ollama pull`.
    Only called from the app's Add button. Long-running (multi-GB); the app runs it in the background."""
    c = next((x for x in CATALOG if x["id"] == model_id), None)
    if not c:
        return {"ok": False, "error": "unknown catalog id"}
    src = c["source"]
    try:
        if src.startswith("ollama"):
            model = src.replace("ollama run ", "").replace("ollama pull ", "").strip()
            r = subprocess.run(["ollama", "pull", model], capture_output=True, text=True, timeout=3600)
            return {"ok": r.returncode == 0, "label": c["label"], "error": "" if r.returncode == 0 else r.stderr.strip()[:200]}
        if os.path.exists(PYENV_PY):
            code = "from huggingface_hub import snapshot_download as s; s(%r)" % src
            r = subprocess.run([PYENV_PY, "-c", code], capture_output=True, text=True, timeout=3600)
            return {"ok": r.returncode == 0, "label": c["label"], "error": "" if r.returncode == 0 else r.stderr.strip()[-200:]}
        return {"ok": False, "error": "no downloader available (huggingface_hub / ollama)"}
    except Exception as e:
        return {"ok": False, "error": str(e)[:200]}


def main(argv):
    if "--remove" in argv:
        i = argv.index("--remove")
        mid = argv[i + 1] if i + 1 < len(argv) else ""
        sys.stdout.write(json.dumps(remove(mid)))
        return
    if "--add" in argv:
        i = argv.index("--add")
        mid = argv[i + 1] if i + 1 < len(argv) else ""
        sys.stdout.write(json.dumps(add(mid)))
        return
    if "--catalog" in argv:
        sys.stdout.write(json.dumps({"catalog": CATALOG}))
        return
    sys.stdout.write(json.dumps(build()))


if __name__ == "__main__":
    main(sys.argv[1:])
