import os, io, json, time, tarfile, subprocess, urllib.request
from pathlib import Path

import gradio as gr
import spaces
import httpx
from huggingface_hub import hf_hub_download

# ===== 0. CONFIG (override in Space Settings → Variables and secrets) =====
LLAMA_TAG  = os.environ.get("LLAMA_TAG", "")            # pin e.g. "b8407"; "" = latest
MODEL_REPO = os.environ.get("MODEL_REPO", "Abiray/MiniCPM5-1B-GGUF")
MODEL_FILE = os.environ.get("MODEL_FILE", "minicpm5-1b-Q6_K.gguf")
CTX_SIZE   = int(os.environ.get("CTX_SIZE", "8192"))
PARALLEL   = int(os.environ.get("PARALLEL", "2"))
LLAMA_PORT = int(os.environ.get("LLAMA_PORT", "8080"))  # internal only
API_KEY    = os.environ.get("LLAMA_API_KEY", "")        # SET THIS as a Space secret!
SYSTEM_PROMPT = os.environ.get("SYSTEM_PROMPT", "You are a helpful AI assistant.")

# ===== 1. HOLD THE ZEROGPU HARDWARE (~5s of quota, once, at boot) =====
@spaces.GPU(duration=5)
def _hold_zerogpu():
    print("[zerogpu] host claimed; GPU slice released, CPU+RAM stay ours.")
    return True

_hold_zerogpu()

# ===== 2. READ THE "DYNAMIC" CPU/RAM WE ACTUALLY GOT (cgroup truth) =====
def effective_cpus() -> int:
    try:
        quota, period = Path("/sys/fs/cgroup/cpu.max").read_text().split()[:2]
        if quota != "max":
            return max(1, int(quota) // int(period))
    except Exception:
        pass
    try:
        return len(os.sched_getaffinity(0))
    except Exception:
        return os.cpu_count() or 2

def memory_limit_gb():
    try:
        v = Path("/sys/fs/cgroup/memory.max").read_text().strip()
        if v != "max":
            return round(int(v) / 1e9, 1)
    except Exception:
        pass
    return None

CORES = effective_cpus()
print(f"[resources] effective cores: {CORES} | RAM limit: {memory_limit_gb() or '?'} GB")

# ===== 3. FETCH PREBUILT llama.cpp (seconds, zero compilation) =====
FALLBACK_TAG = "b8407"

def resolve_tag() -> str:
    if LLAMA_TAG:
        return LLAMA_TAG
    try:
        req = urllib.request.Request(
            "https://api.github.com/repos/ggml-org/llama.cpp/releases/latest",
            headers={"User-Agent": "hf-space"})
        with urllib.request.urlopen(req, timeout=10) as r:
            return json.load(r)["tag_name"]
    except Exception as e:
        print(f"[llama.cpp] latest-tag lookup failed ({e}); pinned fallback {FALLBACK_TAG}")
        return FALLBACK_TAG

def ensure_llama_server() -> Path:
    tag = resolve_tag()
    root = Path.home() / ".cache" / "llamacpp" / tag
    server = root / "build" / "bin" / "llama-server"
    if not server.exists():
        # llama.cpp Linux binaries are now distributed as .tar.gz (not .zip)
        url = (f"https://github.com/ggml-org/llama.cpp/releases/download/"
               f"{tag}/llama-{tag}-bin-ubuntu-x64.tar.gz")
        print(f"[llama.cpp] downloading {url}")
        req = urllib.request.Request(url, headers={"User-Agent": "hf-space"})
        with urllib.request.urlopen(req, timeout=180) as r:
            blob = r.read()
        root.mkdir(parents=True, exist_ok=True)
        # Extract tar.gz archive
        with tarfile.open(fileobj=io.BytesIO(blob), mode="r:gz") as t:
            t.extractall(root)
        for f in (root / "build" / "bin").iterdir():
            f.chmod(0o755)
    return server

# ===== 4. MODEL =====
print(f"[model] downloading {MODEL_REPO}/{MODEL_FILE}")
model_path = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE)

# ===== 5. LAUNCH llama-server, SIZED TO THE DETECTED HARDWARE =====
server_bin = ensure_llama_server()
UPSTREAM = f"http://127.0.0.1:{LLAMA_PORT}"

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

cmd = [
    str(server_bin), "-m", model_path,
    "--host", "127.0.0.1", "--port", str(LLAMA_PORT),
    "-t", str(CORES), "-tb", str(CORES),   # every core the cgroup gives us
    "-c", str(CTX_SIZE),
    "-np", str(PARALLEL),                  # real concurrent slots, no global lock
    "-fa",                                 # flash attention → much smaller KV cache
    "--cache-reuse", "256",                # KV reuse across chat turns
]
if API_KEY:
    cmd += ["--api-key", API_KEY]

proc = subprocess.Popen(cmd, env=env)

def wait_ready(timeout=600):
    t0 = time.time()
    while time.time() - t0 < timeout:
        if proc.poll() is not None:
            raise RuntimeError(f"llama-server died early (code {proc.returncode})")
        try:
            if httpx.get(f"{UPSTREAM}/health", timeout=2).status_code == 200:
                print("[llama.cpp] server ready")
                return
        except Exception:
            pass
        time.sleep(1)
    raise TimeoutError("llama-server not ready in time")

wait_ready()

# ===== 6. GRADIO UI (same backend as the API) =====
from openai import OpenAI
oai = OpenAI(base_url=f"{UPSTREAM}/v1", api_key=API_KEY or "not-needed")

def _normalize(history):
    msgs = [{"role": "system", "content": SYSTEM_PROMPT}]
    for item in history:
        if isinstance(item, dict):
            msgs.append({"role": item["role"], "content": item["content"]})
        else:
            user, bot = item
            msgs.append({"role": "user", "content": user})
            if bot:
                msgs.append({"role": "assistant", "content": bot})
    return msgs

def chat(message, history):
    stream = oai.chat.completions.create(
        model="local",
        messages=_normalize(history) + [{"role": "user", "content": message}],
        stream=True, temperature=0.7, max_tokens=1024,
    )
    out = ""
    for chunk in stream:
        delta = chunk.choices[0].delta.content
        if delta:
            out += delta
            yield out

with gr.Blocks() as demo:
    gr.Markdown(
        f"# ⚡ llama.cpp on ZeroGPU's free CPU\n"
        f"`{MODEL_REPO}` · {CORES} threads · ctx {CTX_SIZE} · {PARALLEL} parallel slots\n\n"
        f"**External API:** `POST /v1/chat/completions` on this Space's `*.hf.space` URL."
    )
    gr.ChatInterface(fn=chat, examples=["Who are you?",
        "Write a python script to reverse a string.", "Explain quantum computing."])

demo = demo.queue()

# ===== 7. PUBLIC /v1 PROXY — the supported way, no monkey-patch =====
from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import Response, StreamingResponse

app = FastAPI(title="llama-cpp-proxy")
app.add_middleware(CORSMiddleware, allow_origins=["*"],
                   allow_methods=["*"], allow_headers=["*"])
_http = httpx.AsyncClient(timeout=httpx.Timeout(None))

@app.api_route("/v1/{path:path}", methods=["GET", "POST"])
async def v1_proxy(path: str, request: Request):
    body = await request.body()
    headers = {k: request.headers[k] for k in ("authorization", "content-type")
               if k in request.headers}
    wants_stream = False
    if body:
        try:
            wants_stream = bool(json.loads(body).get("stream"))
        except Exception:
            pass

    url = f"{UPSTREAM}/v1/{path}"
    if request.method == "POST" and wants_stream:
        req = _http.build_request("POST", url, content=body, headers=headers)
        resp = await _http.send(req, stream=True)
        async def gen():
            try:
                async for chunk in resp.aiter_raw():
                    yield chunk
            finally:
                await resp.aclose()
        return StreamingResponse(gen(), status_code=resp.status_code,
                                 media_type="text/event-stream")

    resp = await _http.request(request.method, url, content=body or None, headers=headers)
    return Response(content=resp.content, status_code=resp.status_code,
                    media_type=resp.headers.get("content-type", "application/json"))

app = gr.mount_gradio_app(app, demo, path="/")   # mounted LAST so /v1/* wins

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("PORT", "7860")))