#!/usr/bin/env python3
"""
bge 嵌入服务故障切换代理 (2026-09-05)
======================================
zhiyid 只认 localhost:8000 — 本代理常驻 8000：
  * 远端优先: 192.168.5.104:8000 (ZSB GPU / DML bge-m3)
  * 远端故障 → 自动切本地 127.0.0.1:8001 (fallback bge)
  * 远端恢复 → 自动切回（每次请求前探测 /health）

启动: python3 bge-failover-proxy.py
端口: 8000 (env PROXY_PORT)
远端: 192.168.5.104:8000 (env REMOTE_EMBED_URL)
本地: 127.0.0.1:8001 (env LOCAL_EMBED_URL)
"""
import json
import logging
import os
import socket
import subprocess
import threading
import time
from http.server import HTTPServer, BaseHTTPRequestHandler, ThreadingHTTPServer
from urllib import request as urlrequest
from urllib.error import URLError

PROXY_PORT = int(os.environ.get("PROXY_PORT", "8000"))
REMOTE_EMBED_URL = os.environ.get("REMOTE_EMBED_URL", "http://192.168.5.104:8000")
LOCAL_EMBED_URL = os.environ.get("LOCAL_EMBED_URL", "http://127.0.0.1:8001")
HEALTH_PATH = "/health"
EMBED_PATH = "/v1/embeddings"
HEALTH_TIMEOUT = float(os.environ.get("HEALTH_TIMEOUT", "1.5"))
REQUEST_TIMEOUT = float(os.environ.get("REQUEST_TIMEOUT", "60"))

logging.basicConfig(
    level=logging.INFO,
    format="[bge-proxy] %(asctime)s %(message)s",
    datefmt="%H:%M:%S",
)
log = logging.getLogger("bge-proxy")

_state_lock = threading.Lock()
# 标记文件：proxy 拉起的 local 需在 proxy 重启后仍记得（避免重启后 local 永不自动停）
LOCAL_STAMP = os.path.expanduser("~/.hermes/run/bge_proxy_local.stamp")
_state = {"remote_ok": None, "last_check": 0.0, "using": None,
         "local_started_by_proxy": os.path.exists(LOCAL_STAMP),  # 重启后从标记恢复
         "remote_ok_since": None,
         "remote_retry_after": 0.0}  # embedding 实发失败后的冷却截止（秒，300s）

# 按需拉起本地 fallback 的参数（2026-09-05: bge-embed 不再常驻，104 故障时才拉起）
LOCAL_START_TIMEOUT = float(os.environ.get("LOCAL_START_TIMEOUT", "30"))   # 拉起后等待就绪上限
LOCAL_STOP_GRACE = float(os.environ.get("LOCAL_STOP_GRACE", "60"))         # 远端稳定多久后停本地
BGE_EMBED_SERVICE = os.environ.get("BGE_EMBED_SERVICE", "bge-embed.service")


def _probe_remote() -> bool:
    """探测远端 /health，1.5s 超时。"""
    try:
        req = urlrequest.Request(REMOTE_EMBED_URL + HEALTH_PATH)
        with urlrequest.urlopen(req, timeout=HEALTH_TIMEOUT) as r:
            if r.status != 200:
                return False
            body = r.read(200).decode("utf-8", "ignore")
            return "ok" in body
    except Exception:
        return False


def _local_healthy() -> bool:
    """探测本地 fallback (8001) /health。"""
    try:
        req = urlrequest.Request(LOCAL_EMBED_URL + HEALTH_PATH)
        with urlrequest.urlopen(req, timeout=1.5) as r:
            if r.status != 200:
                return False
            return "ok" in r.read(200).decode("utf-8", "ignore")
    except Exception:
        return False


def _ensure_local() -> bool:
    """远端不可用时确保本地可用：未启动则 systemctl --user 拉起 bge-embed。

    返回 True = 本地可用；False = 本地拉起失败/超时（转发将 502）。
    """
    if _local_healthy():
        return True
    with _state_lock:
        if not _state["local_started_by_proxy"]:
            log.info("远端不可用 → 拉起 %s (按需)", BGE_EMBED_SERVICE)
            try:
                subprocess.Popen(
                    ["systemctl", "--user", "start", BGE_EMBED_SERVICE],
                    stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
                )
                _state["local_started_by_proxy"] = True
                try:
                    os.makedirs(os.path.dirname(LOCAL_STAMP), exist_ok=True)
                    open(LOCAL_STAMP, "w").close()
                except Exception:
                    pass
            except Exception as e:
                log.error("拉起 %s 失败: %s", BGE_EMBED_SERVICE, e)
    # 轮询等待模型加载（int8 onnx ~10-20s）
    deadline = time.time() + LOCAL_START_TIMEOUT
    while time.time() < deadline:
        if _local_healthy():
            return True
        time.sleep(1)
    log.error("本地 fallback %s 拉起超时 (%ss)", BGE_EMBED_SERVICE, LOCAL_START_TIMEOUT)
    return False


def _maybe_stop_local(remote_ok: bool):
    """远端已恢复且稳定 → 停掉 proxy 拉起的本地 fallback，回到省电模式。"""
    now = time.time()
    with _state_lock:
        if remote_ok and _state["local_started_by_proxy"]:
            if _state["remote_ok_since"] is None:
                _state["remote_ok_since"] = now
            elif now - _state["remote_ok_since"] >= LOCAL_STOP_GRACE:
                log.info("远端已稳定 %ss → 停 %s (释放内存)", LOCAL_STOP_GRACE, BGE_EMBED_SERVICE)
                try:
                    subprocess.Popen(
                        ["systemctl", "--user", "stop", BGE_EMBED_SERVICE],
                        stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
                    )
                    _state["local_started_by_proxy"] = False
                    _state["remote_ok_since"] = None
                    try:
                        os.remove(LOCAL_STAMP)
                    except OSError:
                        pass
                except Exception as e:
                    log.error("停止 %s 失败: %s", BGE_EMBED_SERVICE, e)
        elif not remote_ok:
            _state["remote_ok_since"] = None


def choose_target() -> tuple[str, bool]:
    """返回 (目标URL, 是否远端)。带 3s 冷却 + 状态缓存 + 按需拉起本地 fallback。"""
    global _state
    now = time.time()
    with _state_lock:
        cache_fresh = _state["remote_ok"] is not None and now - _state["last_check"] < 3.0
        in_cooldown = _state["remote_ok"] is False and now < _state["remote_retry_after"]
        if cache_fresh or in_cooldown:
            remote_ok = _state["remote_ok"]
        else:
            remote_ok = _probe_remote()
            _state["remote_ok"] = remote_ok
            _state["last_check"] = now
    _maybe_stop_local(remote_ok)
    if remote_ok:
        using = "remote(104-GPU)"
        if using != _state["using"]:
            log.info("切换 → %s", using)
            with _state_lock:
                _state["using"] = using
        return REMOTE_EMBED_URL, True
    # 远端不可用 → 确保本地可用
    local_ok = _ensure_local()
    using = "local(106-CPU)" if local_ok else "local(106-CPU, 未就绪)"
    if using != _state["using"]:
        log.info("切换 → %s", using)
        with _state_lock:
            _state["using"] = using
    return LOCAL_EMBED_URL, False


def _forward(url: str, method: str, body: bytes, headers: dict) -> tuple[int, dict, bytes]:
    """把请求转发到目标，返回 (状态码, 响应头, 响应体)。"""
    req = urlrequest.Request(url, data=body if method == "POST" else None, method=method)
    for k, v in headers.items():
        if k.lower() not in ("host", "content-length", "connection", "accept-encoding"):
            req.add_header(k, v)
    try:
        with urlrequest.urlopen(req, timeout=REQUEST_TIMEOUT) as r:
            resp_body = r.read()
            resp_headers = dict(r.headers.items())
            return r.status, resp_headers, resp_body
    except URLError as e:
        return 502, {"Content-Type": "application/json"}, json.dumps(
            {"error": f"proxy forward failed: {e}"}
        ).encode()


class Handler(BaseHTTPRequestHandler):
    protocol_version = "HTTP/1.1"

    def log_message(self, fmt, *args):
        log.info("%s %s" % (self.address_string(), fmt % args))

    def _handle(self):
        length = int(self.headers.get("Content-Length", 0) or 0)
        body = self.rfile.read(length) if length else b""
        target, _remote_ok = choose_target()
        status, resp_headers, resp_body = _forward(
            target + self.path, self.command, body, dict(self.headers)
        )
        # 远端 embedding 实发失败（500/502/503/504，health 探测不出的故障）
        # → 标故障 + 30s 冷却 + 立即降级 local 重试一次
        if status >= 500 and target.startswith(REMOTE_EMBED_URL):
            with _state_lock:
                _state["remote_ok"] = False
                # 300s 冷却: 104 可能出现 health OK 但 embedding 500（实发故障），
                # 短冷却会导致反复 remote→500→local 抖动并可能反复启停 bge-embed
                _state["remote_retry_after"] = time.time() + 300
                _state["using"] = None  # 强制下次 log 切换
            log.info("远端 embedding 实发失败(%s) → 降级 local 重试", status)
            if _ensure_local():
                target = LOCAL_EMBED_URL
                status, resp_headers, resp_body = _forward(
                    target + self.path, self.command, body, dict(self.headers)
                )
        self.send_response(status)
        # 只转发安全响应头；Content-Length 必须用实际长度
        for k, v in resp_headers.items():
            if k.lower() in ("content-type",):
                self.send_header(k, v)
        self.send_header("Content-Length", str(len(resp_body)))
        self.send_header("X-BGE-Proxy", "remote-104" if target.startswith(REMOTE_EMBED_URL) else "local-106")
        self.end_headers()
        if self.command != "HEAD":
            try:
                self.wfile.write(resp_body)
            except (BrokenPipeError, ConnectionResetError):
                # 客户端提前断开（探测超时/取消）——不影响其他请求（多线程）
                pass

    do_GET = _handle
    do_POST = _handle
    do_PUT = _handle
    do_DELETE = _handle


def _startup_probe():
    """启动时打一条日志说明当前选路。"""
    remote_ok = _probe_remote()
    with _state_lock:
        _state["remote_ok"] = remote_ok
        _state["last_check"] = time.time()
        _state["using"] = "remote(104-GPU)" if remote_ok else "local(106-CPU)"
    log.info("启动选路: %s (remote=%s)", _state["using"], remote_ok)


if __name__ == "__main__":
    log.info("bge 故障切换代理启动 0.0.0.0:%s (Threading)", PROXY_PORT)
    log.info("  远端: %s", REMOTE_EMBED_URL)
    log.info("  本地 fallback: %s", LOCAL_EMBED_URL)
    _startup_probe()
    # ThreadingHTTPServer: 单请求卡住不阻塞 /health 探测和其他请求
    server = ThreadingHTTPServer(("0.0.0.0", PROXY_PORT), Handler)
    server.daemon_threads = True
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        log.info("停止")