# -*- coding: utf-8 -*-
"""samai-4b Colab 一键引导 v2 — 聊天页(7861) + OpenAI兼容API(7862) 双服务.

用法: 在新 Colab (T4 GPU) 笔记本第一个 cell 运行:
    !wget -qO /content/boot.py https://huggingface.co/tchbcb/samai-4b/resolve/main/bootstrap.py
    !python3 /content/boot.py

流程: 装依赖 -> 从 HF 拉 chat_server.py / openai_api.py -> 下载模型 (~8.3GB) ->
启动两个服务 -> 启动两个 aitun 隧道 -> 打印公网地址.
API key: 1234 (Bearer)
"""
import os
import subprocess
import time

T0 = time.time()
REPO = "tchbcb/samai-4b"  # 公开仓库, 匿名下载即可
API_KEY = "1234"


def log(m):
    print(f"[{time.time()-T0:6.1f}s] {m}", flush=True)


def sh(c, t=600):
    r = subprocess.run(c, shell=True, capture_output=True, text=True, timeout=t)
    return (r.stdout + r.stderr).strip()


log("install deps (flask/aitun/transformers==5.16.1) ...")
print(sh("pip install -q flask aitun huggingface_hub 2>&1 | tail -1"))
v = sh("python3 -c 'import transformers;print(transformers.__version__)'")
log(f"transformers = {v}")
if v and not v.startswith("5.16"):
    print(sh("pip install -q transformers==5.16.1 2>&1 | tail -1"))
    log("pinned transformers==5.16.1 (modeling 代码按此版本修补)")

log("pull servers from HF ...")
for f in ["chat_server.py", "openai_api.py", "control_server.py"]:
    subprocess.run(
        f"curl -sL https://huggingface.co/{REPO}/resolve/main/{f} -o /content/{f}",
        shell=True, check=True)
import py_compile  # noqa: E402
py_compile.compile("/content/chat_server.py", doraise=True)
py_compile.compile("/content/openai_api.py", doraise=True)
py_compile.compile("/content/control_server.py", doraise=True)
log("server scripts OK")

log("download model (~8.3GB, 首次约3-6分钟) ...")
r = subprocess.run(
    ["python3", "-c", f"""
from huggingface_hub import snapshot_download
p = snapshot_download("{REPO}",
                      allow_patterns=["*.json","*.py","*.jinja","*.txt","*.safetensors","*.md"])
print("MODEL_DIR:", p)
"""], capture_output=True, text=True)
print(r.stdout[-400:], r.stderr[-300:] if r.returncode else "")
model_dir = [l.split("MODEL_DIR: ")[1].strip() for l in r.stdout.splitlines()
             if "MODEL_DIR:" in l][-1]

if not os.path.exists("/content/samai-4b-sft"):
    os.symlink(model_dir, "/content/samai-4b-sft")
log("model ready: " + model_dir)

subprocess.run("pkill -f chat_server.py; pkill -f openai_api.py; "
               "pkill -f control_server.py; "
               "pkill -f 'aitun -p 7861'; pkill -f 'aitun -p 7862'; "
               "pkill -f 'aitun -p 5000'", shell=True)
time.sleep(2)
env = f"S4_MODEL_DIR=/content/samai-4b-sft S4_API_KEY={API_KEY}"
subprocess.Popen(f"setsid nohup env {env} python3 /content/chat_server.py"
                 f" > /content/chat_run.log 2>&1 &",
                 shell=True, start_new_session=True)
subprocess.Popen(f"setsid nohup env {env} python3 /content/openai_api.py"
                 f" > /content/api_run.log 2>&1 &",
                 shell=True, start_new_session=True)
subprocess.Popen("setsid nohup python3 /content/control_server.py"
                 " > /content/control_run.log 2>&1 &",
                 shell=True, start_new_session=True)
log("services starting (模型加载约2分钟) ...")

subprocess.Popen("setsid nohup aitun -p 7861 > /content/aitun7861.log 2>&1 &",
                 shell=True, start_new_session=True)
subprocess.Popen("setsid nohup aitun -p 7862 > /content/aitun7862.log 2>&1 &",
                 shell=True, start_new_session=True)
subprocess.Popen("setsid nohup aitun -p 5000 > /content/aitun5000.log 2>&1 &",
                 shell=True, start_new_session=True)
time.sleep(12)
log("=== 隧道 7861 (聊天页) ===")
print(sh("tail -4 /content/aitun7861.log"))
log("=== 隧道 7862 (OpenAI API) ===")
print(sh("tail -4 /content/aitun7862.log"))
log("=== 隧道 5000 (远程控制) ===")
print(sh("tail -4 /content/aitun5000.log"))

for name, port in [("chat", 7861), ("api", 7862)]:
    for i in range(30):
        time.sleep(10)
        h = sh(f"curl -s -m 5 http://127.0.0.1:{port}/health", t=10)
        if h and '"loaded": true' in h:
            log(f"[{name}] ready: {h}")
            break
        if i == 29:
            log(f"[{name}] WARN 未就绪: {h}")

log("=" * 60)
log("BOOTSTRAP_DONE")
log(f"聊天页 (UI)       : 见 aitun7861.log 的公网地址")
log(f"OpenAI 兼容 API : 见 aitun7862.log 的公网地址  (key={API_KEY})")
log(f"远程控制通道     : 见 aitun5000.log 的公网地址 (把三个地址都发给助手)")
log("API 测试: curl <API地址>/v1/chat/completions -H 'Authorization: Bearer 1234' "
    "-H 'Content-Type: application/json' "
    "-d '{\"model\":\"samai-4b\",\"messages\":[{\"role\":\"user\",\"content\":\"你好\"}]}'")