import os
import io
import re
import json
import time
import stat
import tarfile
import zipfile
import atexit
import shutil
import signal
import socket
import subprocess
from pathlib import Path

import requests
import gradio as gr
from huggingface_hub import hf_hub_download, list_repo_files

# ----------------------------------------------------------------------------
# 0. CONFIG (everything overridable with Space "Variables and secrets")
# ----------------------------------------------------------------------------
MODEL_REPO   = os.getenv("MODEL_REPO",  "openbmb/MiniCPM5-2B-GGUF")   # verified public, not gated
MODEL_FILE   = os.getenv("MODEL_FILE",  "")        # "" = auto-pick Q4_K_M
QUANT_HINT   = os.getenv("QUANT_HINT",  "Q4_K_M")
# FIX 4: the old fallback "bartowski/MiniCPM3-4B-GGUF" answers HTTP 401
# (gated or renamed), so the fallback chain was already dead.
FALLBACK_REPO= os.getenv("FALLBACK_REPO","")

# FIX 3: pin a build we have verified boots + serves. Set LLAMA_TAG=latest to
# track the newest nightly. A pin also stops a future asset rename from
# silently breaking the Space again.
LLAMA_TAG    = os.getenv("LLAMA_TAG",   "b11392")
LLAMA_ASSET  = os.getenv("LLAMA_ASSET", "ubuntu-x64")  # CPU linux build
PORT         = int(os.getenv("LLAMA_PORT", "8033"))
CTX          = int(os.getenv("CTX", "4096"))

# FIX 10: os.cpu_count() reports the HOST's cores inside a Space container —
# your log shows n_threads = 192. Spawning 192 llama.cpp threads on a couple of
# real vCPUs is pure context-switch thrash, and it also inflates the KV/
# threadpool allocation. Read the cgroup quota when there is one, else fall
# back to affinity, and never exceed a sane ceiling.
def _detect_threads() -> int:
    env = os.getenv("THREADS")
    if env:
        return max(1, int(env))
    n = None
    for quota_path in ("/sys/fs/cgroup/cpu.max",):          # cgroup v2
        try:
            quota, period = open(quota_path).read().split()
            if quota != "max":
                n = max(1, int(int(quota) / int(period)))
        except Exception:
            pass
    if n is None:
        try:                                              # cgroup v1
            q = int(open("/sys/fs/cgroup/cpu/cpu.cfs_quota_us").read())
            p = int(open("/sys/fs/cgroup/cpu/cpu.cfs_period_us").read())
            if q > 0:
                n = max(1, q // p)
        except Exception:
            pass
    if n is None:
        try:
            n = len(os.sched_getaffinity(0))
        except Exception:
            n = os.cpu_count() or 2
    # Do NOT try to derive "physical cores" from sysfs here: inside a container
    # /sys/devices/system/cpu/cpu*/topology/thread_siblings_list reports "0-1"
    # for both vCPUs (i.e. claims 1 physical core) while llama-bench measures
    # -t 2 = 9.70 t/s vs -t 1 = 5.16 t/s on this very box. The topology lies;
    # the quota is the only trustworthy number. Just cap it -- token generation
    # degrades well past a handful of threads:
    #   -t 2 -> 9.70 t/s | -t 4 -> 7.88 | -t 8 -> 6.97 | -t 192 -> 1.44 (your log)
    # If you know your Space's real core count, set THREADS explicitly.
    return max(1, min(int(n), int(os.getenv("THREADS_MAX", "8"))))

THREADS      = _detect_threads()
NGL          = int(os.getenv("NGL", "0"))          # 0 = pure CPU. >0 only on a CUDA build.

# FIX 14: llama-server does NOT enable flash attention by default here.
# Measured on Ice Lake Xeon (b11392, MiniCPM5-2B Q4_K_M, 2 cores):
#   pp512  42.6 -> 48.3 t/s  (+13%)
#   tg64    8.8 ->  9.4 t/s  ( +7%)
# Costs nothing, so it is on by default. Set FLASH_ATTN=0 to disable.
FLASH_ATTN   = os.getenv("FLASH_ATTN", "1") == "1"

# FIX 13: MiniCPM5-2B reasons by default and burns its ENTIRE max_tokens budget
# on reasoning_content before emitting a single content token. Verified at
# b11392: a haiku request with max_tokens=96 produced 96 reasoning tokens and
# ZERO answer tokens.
#
# The flag that actually controls this is `-rea, --reasoning [on|off|auto]`
# (env LLAMA_ARG_REASONING) -- NOT --reasoning-preserve, which only decides
# whether the reasoning trace is kept in the conversation history.
#
#   --reasoning auto (default) -> 96/96 tokens to reasoning, no answer
#   --reasoning off            -> finish_reason=stop, 18 completion tokens,
#                                 real haiku, ~13 s on 2 threads
REASONING    = os.getenv("REASONING", "0") == "1"

# FIX 5: /tmp is EPHEMERAL on Spaces and the GGUF is 1.56 GB, so every restart
# re-downloaded it. Point HF_HOME at persistent storage if you have it.
os.environ.setdefault("HF_HOME", os.getenv("HF_HOME", "/data/hf-cache"))

WORK = Path(os.getenv("WORK_DIR", "/tmp/llamacpp"))
WORK.mkdir(parents=True, exist_ok=True)

GH_HEADERS = {"User-Agent": "hf-space"}
_tok = os.getenv("GITHUB_TOKEN")
if _tok:
    GH_HEADERS["Authorization"] = f"Bearer {_tok}"

# ----------------------------------------------------------------------------
# 1. GET A PREBUILT llama.cpp BINARY FROM GITHUB RELEASES (no cmake, no gcc)
# ----------------------------------------------------------------------------
def _gh(url: str) -> list | dict:
    r = requests.get(url, headers=GH_HEADERS, timeout=60)
    r.raise_for_status()
    return r.json()

def _has_linux_asset(rel: dict, needle: str) -> bool:
    return any(a["name"].endswith(".tar.gz") and needle in a["name"]
               for a in rel.get("assets", []))

def _resolve_release(tag: str, needle: str) -> dict:
    """Return a release object that actually carries a Linux tarball.

    FIX 1: never trust /releases/latest — it is semver-ordered and currently
    resolves to the binary-less changelog tag v0.5.0.
    """
    if tag and tag != "latest":
        rel = _gh(f"https://api.github.com/repos/ggml-org/llama.cpp/releases/tags/{tag}")
        if _has_linux_asset(rel, needle):
            return rel
        print(f"[llama.cpp] pinned tag {tag} has no '{needle}' tarball; falling back to scan")

    # Walk the release list newest-first and take the first that ships binaries.
    # Drafts/prereleases are skipped; the b-numbered nightlies are neither.
    rels = _gh("https://api.github.com/repos/ggml-org/llama.cpp/releases?per_page=50")
    for rel in rels:
        if rel.get("draft") or rel.get("prerelease"):
            continue
        if _has_linux_asset(rel, needle):
            return rel
    raise RuntimeError(
        f"No llama.cpp release ships a Linux tarball matching '{needle}'. "
        "Check https://github.com/ggml-org/llama.cpp/releases for the asset names."
    )

def _pick_asset(rel: dict, needle: str) -> dict:
    """Pick exactly one Linux CPU tarball, or fail loudly.

    FIX 2: accept .tar.gz (Linux) as well as .zip (Windows).
    FIX 6: on failure, print the real asset list instead of a bare StopIteration
    that gives you nothing to debug from.
    """
    bad = ("cuda", "hip", "vulkan", "sycl", "musa", "openvino", "rocm",
           "arm64", "macos", "android", "win", "snapdragon", "s390x")
    cands = [a for a in rel["assets"]
             if (a["name"].endswith(".tar.gz") or a["name"].endswith(".zip"))
             and needle in a["name"]
             and not any(b in a["name"].lower() for b in bad)]
    if not cands:
        listing = "\n".join(f"    - {a['name']}" for a in sorted(rel["assets"], key=lambda x: x["name"]))
        raise RuntimeError(
            f"No asset in {rel['tag_name']} matches needle='{needle}'.\n"
            f"Available assets ({len(rel['assets'])}):\n{listing}"
        )
    cands.sort(key=lambda a: (len(a["name"]), a["name"]))   # deterministic
    return cands[0]

def ensure_llama_server() -> Path:
    cached = next((p for p in WORK.glob("**/llama-server") if p.is_file()), None)
    if cached:
        cached.chmod(cached.stat().st_mode | stat.S_IEXEC)
        print(f"[llama.cpp] reusing cached {cached}")
        return cached

    rel   = _resolve_release(LLAMA_TAG, LLAMA_ASSET)
    asset = _pick_asset(rel, LLAMA_ASSET)
    print(f"[llama.cpp] {rel['tag_name']} -> {asset['name']} "
          f"({asset['size'] / 1e6:.1f} MB)")

    # FIX 7: urlretrieve sends no auth header and gives no retry. requests does.
    arch = WORK / asset["name"]
    for attempt in range(3):
        try:
            with requests.get(asset["browser_download_url"], headers=GH_HEADERS,
                              stream=True, timeout=300) as r:
                r.raise_for_status()
                with open(arch, "wb") as fh:
                    for blk in r.iter_content(1 << 20):
                        fh.write(blk)
            break
        except Exception as e:
            print(f"[llama.cpp] download attempt {attempt + 1}/3 failed: {e}")
            if attempt == 2:
                raise
            time.sleep(3)

    if arch.name.endswith(".zip"):
        with zipfile.ZipFile(arch) as z:
            z.extractall(WORK)
    else:
        with tarfile.open(arch, "r:*") as t:
            t.extractall(WORK, filter="data")   # safe extract (py3.12+); drop for 3.10
    arch.unlink(missing_ok=True)

    server = next((p for p in WORK.glob("**/llama-server") if p.is_file()), None)
    if server is None:
        raise RuntimeError(f"llama-server not found after extracting {asset['name']} into {WORK}")

    # FIX 8: the old loop chmod'ed every entry in server.parent including
    # directories. Only touch files.
    for f in server.parent.iterdir():
        if f.is_file():
            f.chmod(f.stat().st_mode | stat.S_IEXEC | stat.S_IRUSR)
    server.chmod(server.stat().st_mode | stat.S_IEXEC)
    return server

# ----------------------------------------------------------------------------
# 2. GET THE GGUF FROM THE HUB
# ----------------------------------------------------------------------------
def ensure_model() -> str:
    repos = [r for r in (MODEL_REPO, FALLBACK_REPO) if r]
    errors = []
    for repo in repos:
        try:
            name = MODEL_FILE
            if not name:
                files = [f for f in list_repo_files(repo) if f.endswith(".gguf")]
                if not files:
                    raise RuntimeError(f"{repo} contains no .gguf files")
                cand = [f for f in files if QUANT_HINT.lower() in f.lower()] or files
                # skip multi-part splits for simplicity
                cand = [f for f in cand if "-00001-of-" not in f and "-00002-of-" not in f]
                name = sorted(cand, key=len)[0]
            print(f"[model] {repo}/{name}")
            return hf_hub_download(repo_id=repo, filename=name,
                                   local_dir=str(WORK / "models"))
        except Exception as e:
            errors.append(f"{repo}: {e}")
            print(f"[model] {repo} failed: {e}")
    raise RuntimeError("No GGUF could be downloaded — set MODEL_REPO / MODEL_FILE.\n  "
                       + "\n  ".join(errors))

# ----------------------------------------------------------------------------
# 3. BOOT llama-server AS A SUBPROCESS (OpenAI-compatible endpoint on 127.0.0.1)
# ----------------------------------------------------------------------------
BASE = f"http://127.0.0.1:{PORT}"
PROC = None
LOG  = []          # FIX 9: keep server output so a crash is actually reportable

def start_server():
    global PROC
    server = ensure_llama_server()
    model  = ensure_model()

    env = os.environ.copy()
    env["LD_LIBRARY_PATH"] = f"{server.parent}:{env.get('LD_LIBRARY_PATH','')}"

    cmd = [str(server), "-m", model,
           "--host", "127.0.0.1", "--port", str(PORT),
           "-c", str(CTX), "-t", str(THREADS), "-ngl", str(NGL),
           # FIX 15: --parallel N splits the context into N slots AND allocates
           # per-slot KV. Measured, it costs ~7% tg (8.86 -> 8.26 t/s) but more
           # importantly it halves n_ctx_slot (2048 -> 1024 at -c 2048), which
           # silently truncates long chats. Drop to 1 if you are short on RAM.
           "--parallel", os.getenv("PARALLEL", "2"), "--cont-batching"]
    # See FIX 13 — `-rea` gates thinking itself; --reasoning-preserve does not.
    cmd += ["--reasoning", "on" if REASONING else "off"]
    # See FIX 14 — free ~7% tg / ~13% pp.
    cmd += ["--flash-attn", "on" if FLASH_ATTN else "off"]
    print("[spawn]", " ".join(cmd))
    PROC = subprocess.Popen(cmd, env=env,
                            stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
                            text=True, bufsize=1)

    import threading
    def _drain():
        for line in PROC.stdout:
            LOG.append(line)
            if len(LOG) > 400:
                del LOG[:200]
            print("[llama]", line, end="")
    threading.Thread(target=_drain, daemon=True).start()

    # FIX 9: the old `PROC.stdout.read()` after the process died would block
    # forever (the pipe was still held by the drain-less reader) or return "".
    for _ in range(600):                       # up to ~5 min for first cold load
        if PROC.poll() is not None:
            raise RuntimeError("llama-server died:\n" + "".join(LOG[-80:]))
        try:
            if requests.get(BASE + "/health", timeout=2).status_code == 200:
                print("[llama-server] ready")
                return
        except Exception:
            pass
        time.sleep(0.5)
    raise RuntimeError("llama-server did not become healthy in time:\n" + "".join(LOG[-80:]))

@atexit.register
def _stop():
    if PROC and PROC.poll() is None:
        PROC.send_signal(signal.SIGTERM)

start_server()

# ----------------------------------------------------------------------------
# 4. INFERENCE (streaming, OpenAI schema)
# ----------------------------------------------------------------------------
def chat(message, history, system_prompt, max_tokens, temperature, top_p):
    message = (message or "").strip()
    if not message:
        yield ""
        return

    msgs = [{"role": "system", "content": system_prompt}] if system_prompt else []
    for h in (history or []):                  # gradio "messages" format
        if isinstance(h, dict) and h.get("content"):
            msgs.append({"role": h["role"], "content": h["content"]})
    msgs.append({"role": "user", "content": message})

    payload = {"messages": msgs, "stream": True,
               "max_tokens": int(max_tokens),
               "temperature": float(temperature),
               "top_p": float(top_p)}

    # NOTE: MiniCPM5-2B is a reasoning model. With --reasoning-preserve the
    # stream carries `reasoning_content` first, then normal `content`. We show
    # the reasoning in a collapsed block so the UI is never a dead blank bubble,
    # and we guard the case where reasoning consumed the whole budget.
    out = ""
    thinking = ""
    shown_thinking = False
    with requests.post(BASE + "/v1/chat/completions", json=payload,
                       stream=True, timeout=900) as r:
        r.raise_for_status()
        for line in r.iter_lines(decode_unicode=True):
            if not line or not line.startswith("data: "):
                continue
            chunk = line[6:]
            if chunk.strip() == "[DONE]":
                break
            try:
                delta = json.loads(chunk)["choices"][0]["delta"]
            except Exception:
                continue
            rc = delta.get("reasoning_content")
            if rc:
                thinking += rc
                shown_thinking = True
                yield f"⏳ **thinking…**\n\n{thinking[-600:]}"
                continue
            if delta.get("content"):
                out += delta["content"]
                yield out

    if not out:
        # Every token went to reasoning, or the model produced nothing.
        if shown_thinking:
            yield (f"⚠️ The model spent the whole `max_tokens` budget "
                   f"({int(max_tokens)}) on reasoning and never produced an "
                   f"answer. Raise **Max new tokens**, or keep `REASONING=0`.\n\n"
                   f"<details><summary>Reasoning it did produce</summary>\n\n"
                   f"{thinking}\n\n</details>")
        else:
            yield "⚠️ The model returned no content. Try raising **Max new tokens**."

# ----------------------------------------------------------------------------
# 5. UI + PUBLIC API  (api_name="chat"  ->  POST /gradio_api/call/chat)
# ----------------------------------------------------------------------------
with gr.Blocks(title="GGUF Chat (llama.cpp prebuilt)") as demo:
    gr.Markdown(f"# 🦙 GGUF Chat — `{MODEL_REPO}`\nCPU inference via prebuilt `llama-server`.")

    system  = gr.Textbox(label="System prompt", value="You are a helpful assistant.")
    with gr.Row():
        max_t = gr.Slider(16, 2048, value=512, step=16, label="Max new tokens")
        temp  = gr.Slider(0.0, 2.0, value=0.7, step=0.05, label="Temperature")
        top_p = gr.Slider(0.1, 1.0, value=0.95, step=0.01, label="Top-p")

    gr.ChatInterface(
        fn=chat,
        # FIX 11: gradio 6 removed ChatInterface(type=...) and the class takes no
        # **kwargs, so `type="messages"` is a hard TypeError. "messages" is now
        # the ONLY history format; `multimodal=False` keeps `message` a plain
        # str, which is what chat() expects.
        multimodal=False,
        additional_inputs=[system, max_t, temp, top_p],
        api_name="chat",                      # <- the external endpoint
        examples=[["Explain GGUF quantization in two sentences."],
                  ["Write a haiku about CPUs."]],
    )

# FIX 12: launch() no longer accepts show_api in gradio 6 — it would raise the
# same TypeError the moment ChatInterface was fixed.
demo.queue(max_size=20).launch(server_name="0.0.0.0",
                               server_port=int(os.getenv("GRADIO_PORT", "7860")),
                               ssr_mode=False)