# -*- coding: utf-8 -*-
"""
Jev-style 判別式決策互動 POC（繁體中文）
後端：jevos（feder-cr/jev，OpenVINO int8，CPU 可跑）
- 是非題 (noul)：輸入狀態＋問題，回傳 P(yes)
- 多選題 (choice)：輸入狀態＋問題＋候選選項，回傳每選項機率
- 每次回應標示延遲與 output_tokens=0（零解碼，判別 vs 生成的核心差異）

免費部署：Hugging Face Spaces（CPU Basic 免費方案即可）
  1. 在 huggingface.co/new-space 建一個 Gradio Space（公開）
  2. 上傳 app.py＋requirements.txt
  3. 首次啟動會自動下載 jevos（約 660MB）並快取於 /data，之後秒啟
本地運行：見 README.md
"""
import json
import os
import shutil
import subprocess
import threading
import time
import zipfile
import tarfile

import gradio as gr
import requests
import spaces  # ZeroGPU 執行期內建；不可寫進 requirements.txt（會蓋掉內建模組）

BASE = os.path.dirname(os.path.abspath(__file__))
if os.path.isdir("/data") and os.access("/data", os.W_OK):
    DATA = "/data"
    _DATA_PERSISTENT = True
else:
    DATA = os.path.join(BASE, ".poc-data")
    _DATA_PERSISTENT = False
os.makedirs(DATA, exist_ok=True)

BIN_URL = "https://github.com/feder-cr/jev/releases/download/jevos-v3/jev-linux-x64.tar.gz"
MODEL_URL = "https://github.com/feder-cr/jev/releases/download/jevos-v3/jevos-v3-openvino-int8.zip"
API = "http://127.0.0.1:8017/v1/systemone"
_server = None
_lock = threading.Lock()

# 後端狀態（給 UI 顯示用，避免首次下載時使用者以為當機）
_progress = {"phase": "idle", "pct": 0, "msg": ""}
_last_pct_shown = -1


def _set_phase(phase, msg="", pct=None):
    global _last_pct_shown
    _progress["phase"] = phase
    _progress["msg"] = msg
    if pct is not None:
        _progress["pct"] = pct
    print(f"[poc] {msg or phase}", flush=True)


def status_text():
    ph = _progress["phase"]
    if ph == "ready":
        return "✅ 後端就緒（jevos 運行中）"
    if ph == "downloading":
        return f"⏳ 下載模型中 {_progress['pct']}%（首次啟動需下載約 660MB，請稍候）"
    if ph == "extracting":
        return "⏳ 解壓縮模型中…"
    if ph == "starting":
        return "⏳ 啟動 jevos 中（載入模型約需 10 秒）…"
    if ph == "error":
        return f"❌ 後端錯誤：{_progress['msg']}"
    return "⏳ 後端準備中…"


def _dl(url, dst):
    """下載檔案（支援斷點續傳＋重試），顯示進度百分比。"""
    global _last_pct_shown
    if os.path.exists(dst):
        return dst
    name = url.split("/")[-1]
    tmp = dst + ".part"
    _last_pct_shown = -1
    for attempt in range(1, 4):
        try:
            start = os.path.getsize(tmp) if os.path.exists(tmp) else 0
            if start > 0:
                print(f"[poc] 續傳 {name}（從 {start // 1024 // 1024}MB 繼續）", flush=True)
            else:
                _set_phase("downloading", f"開始下載 {name}")
            headers = {"Range": f"bytes={start}-"} if start else {}
            with requests.get(url, headers=headers, stream=True, timeout=30) as r:
                r.raise_for_status()
                total = int(r.headers.get("content-length", 0) or 0)
                if start and r.status_code != 206:
                    # 伺服器不支援續傳，從頭來
                    print("[poc] 伺服器不支援續傳，從頭下載", flush=True)
                    start, total, mode = 0, total, "wb"
                else:
                    mode = "ab" if start else "wb"
                full_total = total + start
                downloaded = start
                with open(tmp, mode) as f:
                    for chunk in r.iter_content(chunk_size=1024 * 256):
                        if not chunk:
                            continue
                        f.write(chunk)
                        downloaded += len(chunk)
                        if full_total:
                            pct = min(100, downloaded * 100 // full_total)
                            _progress["pct"] = pct
                            if pct // 10 != _last_pct_shown // 10:
                                _last_pct_shown = pct
                                print(f"[poc] 下載進度 {pct}%", flush=True)
            os.rename(tmp, dst)
            _progress["pct"] = 100
            print("[poc] 下載完成", flush=True)
            return dst
        except (requests.RequestException, OSError) as e:
            print(f"[poc] 下載中斷（嘗試 {attempt}/3）：{type(e).__name__}，5 秒後續傳重試", flush=True)
            time.sleep(5)
    raise RuntimeError(f"下載失敗，已重試 3 次：{name}")


def backend_ready():
    """不阻塞地檢查 jevos 是否已在服務。"""
    try:
        requests.post(API, json={"model": "jev-latest", "state": "ping",
                                 "questions": {"q": {"type": "noul", "instructions": "ok?"}}},
                      timeout=3)
        return True
    except Exception:
        return False


def ensure_backend():
    """下載、解壓、啟動 jevos（冪等，可重複呼叫）。"""
    global _server
    with _lock:
        if _server is not None and _server.poll() is None:
            return
        _ensure_backend_locked()


def _ensure_backend_locked():
    global _server
    # 先看是不是已經有 server 在跑（例如本機手動啟動的）→ 有就直接用，不下載
    if backend_ready():
        _set_phase("ready", "jevos 已在運行")
        return
    try:
        jev_dir, binary = _find_or_prepare_backend()
        # 保險：確保可執行（git/網頁上傳可能掉執行權限）
        try:
            os.chmod(binary, 0o755)
        except Exception:
            pass
        _ensure_model_weights(os.path.join(jev_dir, "model"))
        _set_phase("starting", "啟動 jevos…")
        _server = subprocess.Popen([binary, "serve"], cwd=jev_dir,
                                   stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        for _ in range(90):
            time.sleep(1)
            if backend_ready():
                _set_phase("ready", "jevos 就緒")
                return
        raise RuntimeError("jevos 啟動逾時（90 秒）")
    except Exception as e:
        _set_phase("error", f"{type(e).__name__}: {e}")
        raise


def _ensure_model_weights(model_dir):
    """確保 openvino_model.bin 存在；缺失時從 release 下載 zip 並只解出該檔（支援續傳）。"""
    bin_path = os.path.join(model_dir, "openvino_model.bin")
    if os.path.isfile(bin_path):
        return
    print("[poc] 缺少 openvino_model.bin，從 release 下載模型包並解出該檔…", flush=True)
    _set_phase("downloading", "下載模型包（約 630MB，只取權重檔，支援續傳）…")
    zf = _dl(MODEL_URL, os.path.join(DATA, "jevos-v3-openvino-int8.zip"))
    _set_phase("extracting", "解出 openvino_model.bin…")
    os.makedirs(model_dir, exist_ok=True)
    with zipfile.ZipFile(zf) as z:
        for name in z.namelist():
            if name.endswith("openvino_model.bin"):
                with z.open(name) as src, open(bin_path, "wb") as dst:
                    shutil.copyfileobj(src, dst, length=1024 * 1024)
                print("[poc] openvino_model.bin 就緒", flush=True)
                return
    raise RuntimeError("模型包內找不到 openvino_model.bin")


def _find_or_prepare_backend():
    """找後端檔案，順序：repo 內已解壓 → 快取 → repo 內壓縮檔 → 線上下載。
    回傳 (jev_dir, binary)。jev_dir 下保證有執行檔與 model/（巢狀佈局會建連結）。"""
    # 1. 已解壓好的（手動上傳到 repo 根目錄，或之前快取的）
    for base in (BASE, DATA):
        # 1a. 標準佈局：<base>/jev/{jev(執行檔), model/, *.so}
        d = os.path.join(base, "jev")
        b = os.path.join(d, "jev")
        if os.path.isfile(b) and os.path.isdir(os.path.join(d, "model")):
            print(f"[poc] 使用現成後端：{b}", flush=True)
            return d, b
        # 1b. 巢狀佈局：<base>/jev/jev/{jev(執行檔), *.so}，model/ 在 <base>/jev/model/
        #     （tarball 內層目錄被原樣上傳時會長這樣；執行檔在內層，./model 要能被找到）
        inner = os.path.join(d, "jev")
        b2 = os.path.join(inner, "jev")
        outer_model = os.path.join(d, "model")
        if os.path.isfile(b2) and os.path.isdir(outer_model):
            link = os.path.join(inner, "model")
            if not os.path.exists(link):
                try:
                    os.symlink(outer_model, link)
                    print("[poc] 巢狀佈局：已建立 model 連結", flush=True)
                except Exception as e:
                    print(f"[poc] 建 model 連結失敗：{e}", flush=True)
            print(f"[poc] 使用現成後端（巢狀）：{b2}", flush=True)
            return inner, b2
    # 2. repo 根目錄有壓縮檔（手動上傳）→ 解到 DATA
    jev_dir = os.path.join(DATA, "jev")
    tgz = os.path.join(BASE, "jev-linux-x64.tar.gz")
    zf = os.path.join(BASE, "jevos-v3-openvino-int8.zip")
    if os.path.exists(tgz) and os.path.exists(zf):
        print("[poc] 使用 repo 內附壓縮檔解壓…", flush=True)
        _set_phase("extracting", "解壓縮 repo 內附檔案…")
        os.makedirs(DATA, exist_ok=True)
        with tarfile.open(tgz) as tf:
            tf.extractall(DATA, filter="data")
        print("[poc] 主程式解壓完成，解壓縮模型（約 630MB，需要幾分鐘）…", flush=True)
        _extract_zip_progress(zf, jev_dir)
        print("[poc] 模型解壓完成", flush=True)
        return jev_dir, os.path.join(jev_dir, "jev")
    # 3. 線上下載（Spaces 執行期可連外網）
    os.makedirs(DATA, exist_ok=True)
    tgz = _dl(BIN_URL, os.path.join(DATA, "jev-linux-x64.tar.gz"))
    _set_phase("extracting", "解壓縮 jevos 主程式…")
    with tarfile.open(tgz) as tf:
        tf.extractall(DATA, filter="data")
    zf = _dl(MODEL_URL, os.path.join(DATA, "jevos-v3-openvino-int8.zip"))
    _set_phase("extracting", "解壓縮模型（約 630MB，需要幾分鐘）…")
    _extract_zip_progress(zf, jev_dir)
    print("[poc] 模型解壓完成", flush=True)
    return jev_dir, os.path.join(jev_dir, "jev")


def _extract_zip_progress(zf, jev_dir):
    """解壓 zip 並印進度（避免長時間無輸出被誤判卡死）。"""
    with zipfile.ZipFile(zf) as z:
        members = z.namelist()
        total = len(members)
        for i, m in enumerate(members):
            z.extract(m, jev_dir)
            if (i + 1) % 25 == 0 or (i + 1) == total:
                print(f"[poc] 模型解壓縮 {i + 1}/{total}", flush=True)


def decide(state, mode, question, options_text):
    if not (state or "").strip() or not (question or "").strip():
        return "請先輸入狀態文字與問題", "", "", status_text()
    # 後端還沒好就不要卡住：直接告訴使用者目前進度
    if not backend_ready():
        ensure_backend_if_needed()
        return status_text() + "，請稍後再按一次「送出判別」。", "", "", status_text()
    ensure_backend()
    t0 = time.perf_counter()
    if mode == "是非題":
        payload = {"model": "jev-latest", "state": state,
                   "questions": {"q1": {"type": "noul", "instructions": question}}}
    else:
        criteria = {}
        for line in options_text.strip().split("\n"):
            line = line.strip()
            if not line:
                continue
            if ":" in line:
                name, desc = line.split(":", 1)
            elif "：" in line:
                name, desc = line.split("：", 1)
            else:
                name, desc = line, ""
            name = name.strip()
            desc = desc.strip() or name  # 無描述時用名稱代替（choice 需要 criteria 描述）
            criteria[name] = desc
        if not criteria:
            return "請至少輸入一個候選選項", "", "", status_text()
        payload = {"model": "jev-latest", "state": state,
                   "questions": {"q1": {"type": "choice", "instructions": question,
                                       "criteria": criteria}}}
    try:
        r = requests.post(API, json=payload, timeout=120)
        r.raise_for_status()
        data = r.json()
    except Exception as e:
        return f"後端呼叫失敗：{type(e).__name__}（{e}）", "", "", status_text()
    ms = (time.perf_counter() - t0) * 1000
    try:
        ans = data["answers"]["q1"]
    except (KeyError, TypeError):
        return f"後端回傳格式異常：{str(data)[:300]}", "", "", status_text()
    usage = data.get("usage", {})
    if ans["type"] == "noul":
        p = ans["noul"]
        bars = f"P(yes) = {p:.3f}\n" + "█" * int(p * 40) + f" {p*100:.1f}%"
        summary = f"答案：{'是' if p >= 0.5 else '否'}（信心 {p if p>=0.5 else 1-p:.1%}）"
    else:
        probs = ans["probabilities"]
        bars = "\n".join(f"{k}: {'█' * int(v*40)} {v*100:.1f}%"
                         for k, v in sorted(probs.items(), key=lambda x: -x[1]))
        summary = (f"答案：{ans['choice']}（信心 {ans.get('confidence', 0):.1%}）\n"
                   f"各選項機率：{json.dumps(probs, ensure_ascii=False)}")
    meta = (f"延遲 {ms:.0f} ms ｜ input_tokens={usage.get('input_tokens','?')} ｜ "
            f"output_tokens={usage.get('output_tokens','?')}（零解碼）")
    return summary, bars, meta, status_text()


def ensure_backend_if_needed():
    """確保背景下載執行緒有在跑（不阻塞）。"""
    t = threading.Thread(target=ensure_backend, daemon=True)
    t.start()


PRESETS = {
    "客服訊息分類": ("客人說：「這個月帳單怎麼多了一筆我沒買的東西？」",
                    "多選題", "這封訊息應該派給哪個團隊？",
                    "billing: 帳單、付款、退款\nshipping: 寄送、包裹遺失損壞\ntech: 產品故障不會動"),
    "是否為退款請求": ("I want my money back for the broken headphones.",
                    "是非題", "這是一則退款請求嗎？", ""),
    "內容審核": ("這家餐廳難吃到我想給一星，老闆還敢嗆客人。",
                    "是非題", "這則評論含有針對店家的惡意攻擊嗎？", ""),
}


def load_preset(name):
    s, m, q, o = PRESETS[name]
    return s, m, q, o


@spaces.GPU(duration=10)
def gpu_compat_check():
    """ZeroGPU 硬體要求至少一個 @spaces.GPU 函數才能啟動。
    jevos 為純 CPU 運算（OpenVINO），此函數僅為通過啟動檢查的佔位，實際不參與判別、不消耗有意義的 GPU 配額。"""
    return "ZeroGPU 相容檢查通過（判別運算實際於 CPU 執行）"


CSS = """
.gradio-container { max-width: 900px !important; }
.bar textarea { font-family: monospace !important; white-space: pre !important; }
"""

with gr.Blocks(title="Jev-style 判別式決策 POC") as demo:
    gr.Markdown("# ⚡ Jev-style 判別式決策 POC\n"
                "輸入一段**狀態文字**＋**問題**，模型直接輸出**機率分佈**——不生成任何文字（`output_tokens=0`）。"
                "這就是 System 1 判別引擎：傳統 LLM Agent 每步 1–3 秒還可能寫壞 JSON，這裡一次前向傳播就給出完整機率分佈。\n\n"
                "> ⏰ **首次啟動**需要下載模型（約 660MB），請看下方「後端狀態」等它就緒再按送出，之後啟動都是秒開。")
    status_box = gr.Textbox(label="後端狀態", value="⏳ 後端準備中…", interactive=False)
    gpu_out = gr.Textbox(label="ZeroGPU 相容檢查", value="", interactive=False)
    with gr.Row():
        preset = gr.Dropdown(list(PRESETS.keys()), label="範例一鍵填入", value="客服訊息分類")
        refresh_btn = gr.Button("重新整理狀態", size="sm")
        gpu_btn = gr.Button("GPU 相容檢查", size="sm")
    state = gr.Textbox(label="狀態文字（state）", lines=3)
    mode = gr.Radio(["是非題", "多選題"], label="題型", value="多選題")
    question = gr.Textbox(label="問題（instructions）", lines=1)
    options = gr.Textbox(label="候選選項（多選題用，每行「名稱: 描述」）", lines=4)
    btn = gr.Button("送出判別", variant="primary")
    out_summary = gr.Textbox(label="判別結果")
    out_bars = gr.Textbox(label="機率分佈", elem_classes="bar", lines=6)
    out_meta = gr.Textbox(label="效能", lines=1)

    preset.change(load_preset, preset, [state, mode, question, options])
    refresh_btn.click(lambda: status_text(), outputs=status_box)
    gpu_btn.click(gpu_compat_check, outputs=gpu_out)
    btn.click(decide, [state, mode, question, options],
              [out_summary, out_bars, out_meta, status_box])
    demo.load(load_preset, preset, [state, mode, question, options])
    demo.load(lambda: status_text(), outputs=status_box)

if __name__ == "__main__":
    print("[poc] app.py v20261003-zerogpu (multi-source backend + @spaces.GPU)", flush=True)
    print(f"[poc] 資料目錄：{DATA}（{'重啟後保留，可斷點續傳' if _DATA_PERSISTENT else '重啟後消失，請勿重啟'}）", flush=True)
    # Spaces 啟動時先備好後端（背景執行，不擋 UI）
    ensure_backend_if_needed()
    demo.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)),
                css=CSS)