import os
import json
import time
import subprocess

import requests
from huggingface_hub import hf_hub_download

# ----------------------- config (override via Space variables) -----------------------
MODEL_REPO  = os.environ.get("MODEL_REPO", "ffkbblu/pepekberbulu")
GGUF_FILE   = os.environ.get("GGUF_FILE", "VibeThinker-3B-Q4_K_M.gguf")
HF_TOKEN    = os.environ.get("HF_TOKEN")  # optional; model repo is public
PORT        = int(os.environ.get("PORT", "7860"))
LLAMA_PORT  = int(os.environ.get("LLAMA_PORT", "8080"))
# cpu-basic = 2 physical vCPU, but os.cpu_count() reports the HOST core count inside the
# cgroup -> launching -t with that number oversubscribes the 2 vCPU and collapses
# throughput (~0.1 tok/s). Default to 2; override via NUM_THREADS.
def _detect_cpus():
    # cgroup v2
    try:
        q, p = open("/sys/fs/cgroup/cpu.max").read().split()
        if q != "max":
            return max(1, round(int(q) / int(p)))
    except Exception:
        pass
    # cgroup v1
    try:
        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:
            return max(1, round(q / p))
    except Exception:
        pass
    return os.cpu_count() or 2

DETECTED_CPUS = _detect_cpus()
NUM_THREADS = os.environ.get("NUM_THREADS") or str(DETECTED_CPUS)
BIN         = "/opt/llamabin"
CHAT_TEMPLATE_FILE = os.environ.get("CHAT_TEMPLATE_FILE", "/home/user/app/chat_template.jinja")
LLAMA       = f"http://127.0.0.1:{LLAMA_PORT}"
# Generous 32768 context window; raise N_CTX via a Space variable if needed (slower prefill on CPU).
N_CTX       = os.environ.get("N_CTX", "32768")


def _env():
    e = os.environ.copy()
    e["LD_LIBRARY_PATH"] = BIN + ":" + e.get("LD_LIBRARY_PATH", "")
    return e


# ----------------------- step 1: download the GGUF (no conversion) -----------------------
def ensure_gguf():
    print(f"[init] downloading {GGUF_FILE} from {MODEL_REPO} ...", flush=True)
    path = hf_hub_download(MODEL_REPO, GGUF_FILE, repo_type="model", token=HF_TOKEN)
    print(f"[init] model ready: {path}", flush=True)
    return path


# ----------------------- step 2: launch internal llama-server -----------------------
def start_llama(model_path):
    cmd = [
        os.path.join(BIN, "llama-server"),
        "-m", model_path, "--host", "127.0.0.1", "--port", str(LLAMA_PORT),
        "-c", N_CTX, "-t", NUM_THREADS, "-b", "256", "--no-mmap",
        "--parallel", os.environ.get("PARALLEL", "1"),
        "--jinja",
    ]
    if CHAT_TEMPLATE_FILE and os.path.exists(CHAT_TEMPLATE_FILE):
        cmd += ["--chat-template-file", CHAT_TEMPLATE_FILE]
    # Separate the <think> chain-of-thought into reasoning_content so the answer stays clean.
    cmd += ["--reasoning-format", "auto", "-fa", "on"]
    print(f"[init] cpu_count={os.cpu_count()} threads={NUM_THREADS}", flush=True)
    extra = os.environ.get("LLAMA_EXTRA_ARGS", "").split()
    if extra:
        cmd += extra
    print("[init] starting internal llama-server: " + " ".join(cmd), flush=True)
    subprocess.Popen(cmd, env=_env())
    for _ in range(900):
        try:
            r = requests.get(LLAMA + "/health", timeout=3)
            if r.status_code == 200 and r.json().get("status") == "ok":
                print("[init] internal llama-server is healthy.", flush=True)
                return
        except Exception:
            pass
        time.sleep(1)
    raise RuntimeError("internal llama-server did not become healthy in time")


# ============================================================================
#  Serving: thin PASS-THROUGH proxy to llama-server's native OpenAI endpoints.
# ============================================================================
from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse, JSONResponse, HTMLResponse, Response
from fastapi.middleware.cors import CORSMiddleware

app = FastAPI()
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_credentials=False,
                   allow_methods=["*"], allow_headers=["*"])

DEFAULTS = {"temperature": 0.3, "top_p": 0.9, "top_k": 20, "repeat_penalty": 1.05,
            "max_tokens": int(os.environ.get("MAX_TOKENS", "-1"))}


@app.get("/health")
def health():
    return {"status": "ok"}


@app.get("/v1/models")
def models():
    try:
        return JSONResponse(requests.get(LLAMA + "/v1/models", timeout=15).json())
    except Exception:
        return {"object": "list", "data": [{"id": GGUF_FILE, "object": "model", "owned_by": "ffkbblu"}]}


SYSTEM_PROMPT = os.environ.get("SYSTEM_PROMPT", "").strip()


def _prep(body):
    for k, v in DEFAULTS.items():
        body.setdefault(k, v)

    think = None
    if "enable_thinking" in body:
        think = bool(body.pop("enable_thinking"))
    if "thinking" in body:
        t = str(body.pop("thinking")).lower()
        think = t in ("on", "true", "1", "yes", "smart")
    if think is not None:
        ctk = dict(body.get("chat_template_kwargs") or {})
        ctk["enable_thinking"] = think
        body["chat_template_kwargs"] = ctk

    if SYSTEM_PROMPT:
        msgs = body.get("messages")
        if isinstance(msgs, list):
            msgs = [m for m in msgs
                    if not (isinstance(m, dict) and m.get("role") == "system")]
            msgs.insert(0, {"role": "system", "content": SYSTEM_PROMPT})
            body["messages"] = msgs
    return body


@app.post("/v1/chat/completions")
async def chat_completions(request: Request):
    body = _prep(await request.json())
    stream = bool(body.get("stream", False))
    if stream:
        def gen():
            with requests.post(LLAMA + "/v1/chat/completions", json=body, stream=True, timeout=900) as r:
                for chunk in r.iter_content(chunk_size=None):
                    if chunk:
                        yield chunk
        return StreamingResponse(gen(), media_type="text/event-stream")
    r = requests.post(LLAMA + "/v1/chat/completions", json=body, timeout=900)
    return Response(content=r.content, media_type="application/json", status_code=r.status_code)


@app.post("/v1/completions")
async def completions(request: Request):
    body = _prep(await request.json())
    stream = bool(body.get("stream", False))
    if stream:
        def gen():
            with requests.post(LLAMA + "/v1/completions", json=body, stream=True, timeout=900) as r:
                for chunk in r.iter_content(chunk_size=None):
                    if chunk:
                        yield chunk
        return StreamingResponse(gen(), media_type="text/event-stream")
    r = requests.post(LLAMA + "/v1/completions", json=body, timeout=900)
    return Response(content=r.content, media_type="application/json", status_code=r.status_code)


@app.get("/props")
def props():
    try:
        return JSONResponse(requests.get(LLAMA + "/props", timeout=15).json())
    except Exception as e:
        return {"error": str(e)}


@app.get("/diag")
def diag():
    out = {"os_cpu_count": os.cpu_count(), "detected_cpus": DETECTED_CPUS,
           "num_threads_used": NUM_THREADS, "n_ctx": N_CTX}
    try:
        ci = open("/proc/cpuinfo").read()
        flags = ""
        for line in ci.splitlines():
            if line.startswith("flags"):
                flags = line; break
        out["avx2"]    = " avx2 " in (" "+flags+" ")
        out["avx512f"] = " avx512f " in (" "+flags+" ")
        out["model_name"] = next((l.split(":",1)[1].strip() for l in ci.splitlines() if "model name" in l), "")
        out["proc_count_in_cpuinfo"] = ci.count("processor\t")
    except Exception as e:
        out["cpuinfo_err"] = str(e)
    try:
        pr = requests.get(LLAMA + "/props", timeout=15).json()
        out["llama_system_info"] = pr.get("system_info")
        dp = pr.get("default_generation_settings", {})
        out["llama_n_threads"] = dp.get("n_threads") or pr.get("n_threads")
    except Exception as e:
        out["props_err"] = str(e)
    return out


@app.get("/", response_class=HTMLResponse)
def index():
    return INDEX_HTML


INDEX_HTML = r"""<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="utf-8"/>
<meta name="viewport" content="width=device-width, initial-scale=1"/>
<title>3B Chat</title>
<style>
  :root { color-scheme: light dark; }
  * { box-sizing: border-box; }
  body { margin:0; font-family: ui-sans-serif,system-ui,-apple-system,Segoe UI,Roboto,sans-serif;
         background:#0b0d12; color:#e7e9ee; display:flex; flex-direction:column; height:100vh; }
  header { padding:12px 18px; border-bottom:1px solid #1e2430; font-weight:600; font-size:15px;
           display:flex; align-items:center; gap:8px; }
  header .dot { width:8px; height:8px; border-radius:50%; background:#33d17a; }
  header small { font-weight:400; opacity:.55; }
  #chat { flex:1; overflow-y:auto; padding:18px; display:flex; flex-direction:column; gap:14px; }
  .msg { max-width:820px; width:100%; margin:0 auto; }
  .who { font-size:12px; opacity:.6; margin-bottom:4px; }
  .bubble { padding:10px 14px; border-radius:12px; white-space:pre-wrap; line-height:1.55;
            font-size:14.5px; word-wrap:break-word; overflow-wrap:anywhere; }
  .user { display:flex; justify-content:flex-end; }
  .user .bubble { background:#1d4ed8; color:#fff; }
  .bot .bubble { background:#161b24; border:1px solid #232a36; }
  pre { background:#0f131b; border:1px solid #232a36; border-radius:8px; padding:10px;
        overflow-x:auto; font-size:13px; }
  footer { border-top:1px solid #1e2430; padding:12px; }
  form { max-width:820px; margin:0 auto; display:flex; gap:8px; }
  textarea { flex:1; resize:none; background:#11151d; color:#e7e9ee; border:1px solid #232a36;
             border-radius:10px; padding:11px 12px; font-size:14.5px; max-height:160px; }
  button { background:#1d4ed8; color:#fff; border:0; border-radius:10px; padding:0 18px;
           font-weight:600; cursor:pointer; }
  button:disabled { opacity:.5; cursor:default; }
  .hint { text-align:center; opacity:.45; font-size:12px; margin-top:8px; }
  .typing { opacity:.5; font-style:italic; }
  .row { max-width:820px; margin:0 auto 8px; display:flex; gap:8px; }
  .row button { background:#232a36; font-weight:500; font-size:12px; padding:4px 10px; }
  .think { max-width:820px; width:100%; margin:0 auto 6px; font-size:13px; opacity:.75;
           background:#0f131b; border:1px dashed #2a3340; border-radius:10px; padding:6px 12px; }
  .think summary { cursor:pointer; user-select:none; }
  .think-body { white-space:pre-wrap; margin-top:6px; line-height:1.5; }
</style>
</head>
<body>
<header><span class="dot"></span> 3B Chat <small>&middot; Q4_K_M &middot; fast chat / reasoning &middot; live GGUF</small>
  <label style="margin-left:auto; font-weight:400; font-size:13px; display:flex; align-items:center; gap:6px;">Mode
    <select id="mode" style="background:#11151d; color:#e7e9ee; border:1px solid #232a36; border-radius:8px; padding:4px 8px; font-size:13px;">
      <option value="off" selected>&#9889; Fast (no thinking)</option>
      <option value="on">&#129504; Smart (thinking)</option>
    </select>
  </label>
</header>
<div id="chat"></div>
<footer>
  <div class="row"><button id="reset" type="button">New chat</button></div>
  <form id="f">
    <textarea id="t" rows="1" placeholder="Ask anything..." autofocus></textarea>
    <button id="send" type="submit">Send</button>
  </form>
  <div class="hint">Pure model &middot; OpenAI-compatible API at <code>/v1/chat/completions</code></div>
</footer>
<script>
const chat=document.getElementById('chat'),form=document.getElementById('f'),ta=document.getElementById('t'),sendBtn=document.getElementById('send');
let history=[];
function esc(s){return s.replace(/&/g,'&amp;').replace(/</g,'&lt;').replace(/>/g,'&gt;');}
function render(t){return esc(t).replace(/```([\s\S]*?)```/g,(m,c)=>'<pre><code>'+c.replace(/^\w*\n/,'')+'</code></pre>');}
function stripThink(t){return t.replace(/<think>[\s\S]*?<\/think>/gi,'').replace(/^[\s\S]*?<\/think>/i, m=>m.includes('<think>')?'':m).trim();}
function addUser(t){const w=document.createElement('div');w.className='msg user';w.innerHTML='<div class="bubble">'+render(t)+'</div>';chat.appendChild(w);chat.scrollTop=chat.scrollHeight;}
function addBot(){const w=document.createElement('div');w.className='msg bot';w.innerHTML='<div class="who">Assistant</div><details class="think" style="display:none"><summary>&#129504; Thinking</summary><div class="think-body"></div></details><div class="bubble"><span class="typing">&hellip;</span></div>';chat.appendChild(w);chat.scrollTop=chat.scrollHeight;return {think:w.querySelector('.think'),thinkBody:w.querySelector('.think-body'),bubble:w.querySelector('.bubble')};}
document.getElementById('reset').onclick=()=>{history=[];chat.innerHTML='';ta.focus();};
ta.addEventListener('input',()=>{ta.style.height='auto';ta.style.height=Math.min(ta.scrollHeight,160)+'px';});
ta.addEventListener('keydown',e=>{if(e.key==='Enter'&&!e.shiftKey){e.preventDefault();form.requestSubmit();}});
form.addEventListener('submit',async e=>{
  e.preventDefault();const text=ta.value.trim();if(!text)return;
  ta.value='';ta.style.height='auto';addUser(text);history.push({role:'user',content:text});
  sendBtn.disabled=true;const {think,thinkBody,bubble}=addBot();let acc='';let rc='';
  try{
    const resp=await fetch('/v1/chat/completions',{method:'POST',headers:{'Content-Type':'application/json'},body:JSON.stringify({messages:history,stream:true,thinking:document.getElementById('mode').value})});
    const reader=resp.body.getReader(),dec=new TextDecoder();let buf='';
    while(true){const {value,done}=await reader.read();if(done)break;buf+=dec.decode(value,{stream:true});let idx;
      while((idx=buf.indexOf('\n\n'))>=0){const line=buf.slice(0,idx).trim();buf=buf.slice(idx+2);
        if(!line.startsWith('data:'))continue;const data=line.slice(5).trim();if(data==='[DONE]')continue;
        try{const o=JSON.parse(data);const dl=o.choices?.[0]?.delta||{};
          const rd=dl.reasoning_content||'';if(rd){rc+=rd;think.style.display='block';thinkBody.textContent=rc;chat.scrollTop=chat.scrollHeight;}
          const d=dl.content||'';if(d){acc+=d;bubble.innerHTML=render(stripThink(acc));chat.scrollTop=chat.scrollHeight;}
        }catch(_){}}}
  }catch(err){acc=acc||('[error] '+err);bubble.innerHTML=render(acc);}
  const clean=stripThink(acc);if(!clean)bubble.innerHTML=render('(no response)');
  history.push({role:'assistant',content:clean});sendBtn.disabled=false;ta.focus();
});
</script>
</body>
</html>"""


# ----------------------- main -----------------------
if __name__ == "__main__":
    model_path = ensure_gguf()
    start_llama(model_path)
    import uvicorn
    print(f"[init] starting public proxy on port {PORT} ...", flush=True)
    uvicorn.run(app, host="0.0.0.0", port=PORT, log_level="info")