web_app.py

"""
web_app.py - ChatGPT 注册机 Web UI 后端

FastAPI 主程序：
- REST API：配置、注册、账号池、代理、结果查看
- WebSocket：注册日志、池管理日志实时推送

启动：
    uv run uvicorn web_app:app --host 0.0.0.0 --port 8080 --reload
    或 python web_app.py
"""

from __future__ import annotations

import asyncio
import base64
import binascii
import hashlib
import hmac
import json
import os
import secrets
import sys
import threading
import time
from pathlib import Path
from typing import Dict, List, Optional, Any

import uvicorn
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, HTTPException, Body, Request, Response
from fastapi.responses import HTMLResponse, JSONResponse, FileResponse, PlainTextResponse
from fastapi.middleware.cors import CORSMiddleware

import register as reg

# ============================================================
# 应用初始化
# ============================================================

app = FastAPI(title="ChatGPT 注册机 Web UI", version="1.0.0")

_CORS_ALLOW_ORIGINS = [
    o.strip() for o in os.environ.get(
        "CORS_ALLOW_ORIGINS",
        "http://localhost:8080,http://127.0.0.1:8080"
    ).split(",") if o.strip()
]

app.add_middleware(
    CORSMiddleware,
    allow_origins=_CORS_ALLOW_ORIGINS,
    allow_credentials=True,
    allow_methods=["GET", "POST", "OPTIONS"],
    allow_headers=["Content-Type", "Authorization"],
)

_BASE_DIR = Path(__file__).resolve().parent
_TEMPLATES_DIR = _BASE_DIR / "templates"

# ============================================================
# 全局任务状态
# ============================================================

# 注册任务
_reg_state: Dict[str, Any] = {
    "running": False,
    "success": 0,
    "fail": 0,
    "total": 0,
    "stop_event": None,
    "start_time": None,
}
_reg_log_queue: asyncio.Queue = asyncio.Queue(maxsize=1000)
_reg_ws_clients: List[WebSocket] = []

# 账号池任务
_pool_state: Dict[str, Any] = {
    "running": False,
    "task": "",
}
_pool_log_queue: asyncio.Queue = asyncio.Queue(maxsize=1000)
_pool_ws_clients: List[WebSocket] = []

_PROBE_RESULT_TTL_SEC = 120

# 事件循环引用
_event_loop: Optional[asyncio.AbstractEventLoop] = None

# 运行时 JWT secret（当配置未提供时使用）
_runtime_jwt_secret: str = ""


def _b64url_encode(raw: bytes) -> str:
    return base64.urlsafe_b64encode(raw).rstrip(b"=").decode("ascii")


def _b64url_decode(data: str) -> bytes:
    pad = "=" * (-len(data) % 4)
    try:
        return base64.urlsafe_b64decode((data + pad).encode("ascii"))
    except (ValueError, binascii.Error) as e:
        raise HTTPException(status_code=401, detail="无效 token") from e


def _get_auth_cfg() -> dict:
    cfg = reg.load_config()
    auth = cfg.get("auth") or {}
    return {
        "enabled": bool(auth.get("enabled", True)),
        "username": str(auth.get("username", "admin") or "admin"),
        "password_hash": str(auth.get("password_hash", "") or "").lower(),
        "jwt_secret": str(auth.get("jwt_secret", "") or ""),
        "jwt_exp_hours": int(auth.get("jwt_exp_hours", 24) or 24),
        "cookie_name": str(auth.get("cookie_name", "cc_auth") or "cc_auth"),
    }


def _get_jwt_secret() -> str:
    auth = _get_auth_cfg()
    return auth["jwt_secret"] or _runtime_jwt_secret


def _jwt_encode(payload: dict) -> str:
    header = {"alg": "HS256", "typ": "JWT"}
    h = _b64url_encode(json.dumps(header, separators=(",", ":")).encode("utf-8"))
    p = _b64url_encode(json.dumps(payload, separators=(",", ":")).encode("utf-8"))
    signing_input = f"{h}.{p}".encode("ascii")
    sig = hmac.new(_get_jwt_secret().encode("utf-8"), signing_input, hashlib.sha256).digest()
    s = _b64url_encode(sig)
    return f"{h}.{p}.{s}"


def _jwt_decode(token: str) -> dict:
    parts = token.split(".")
    if len(parts) != 3:
        raise HTTPException(status_code=401, detail="无效 token")
    h, p, s = parts
    signing_input = f"{h}.{p}".encode("ascii")
    expected = hmac.new(_get_jwt_secret().encode("utf-8"), signing_input, hashlib.sha256).digest()
    got = _b64url_decode(s)
    if not hmac.compare_digest(expected, got):
        raise HTTPException(status_code=401, detail="token 签名无效")

    payload_raw = _b64url_decode(p)
    try:
        payload = json.loads(payload_raw.decode("utf-8"))
    except Exception as e:
        raise HTTPException(status_code=401, detail="token payload 无效") from e

    exp = int(payload.get("exp", 0) or 0)
    if exp <= int(time.time()):
        raise HTTPException(status_code=401, detail="token 已过期")
    return payload


def _issue_login_token(username: str) -> str:
    auth = _get_auth_cfg()
    now = int(time.time())
    exp = now + max(1, auth["jwt_exp_hours"]) * 3600
    payload = {"sub": username, "iat": now, "exp": exp}
    return _jwt_encode(payload)


def _token_from_request(request: Request) -> str:
    auth = _get_auth_cfg()
    return (request.cookies.get(auth["cookie_name"]) or "").strip()


def _current_user_from_request(request: Request) -> Optional[dict]:
    auth = _get_auth_cfg()
    if not auth["enabled"]:
        return {"sub": "anonymous"}
    token = _token_from_request(request)
    if not token:
        return None
    try:
        return _jwt_decode(token)
    except HTTPException:
        return None


def _set_auth_cookie(response: Response, token: str, request: Request):
    auth = _get_auth_cfg()
    secure = request.url.scheme == "https"
    response.set_cookie(
        key=auth["cookie_name"],
        value=token,
        httponly=True,
        secure=secure,
        samesite="lax",
        max_age=max(1, auth["jwt_exp_hours"]) * 3600,
        path="/",
    )


def _clear_auth_cookie(response: Response):
    auth = _get_auth_cfg()
    response.delete_cookie(key=auth["cookie_name"], path="/")


@app.on_event("startup")
async def _startup():
    global _event_loop, _runtime_jwt_secret
    _event_loop = asyncio.get_event_loop()
    _runtime_jwt_secret = secrets.token_urlsafe(48)


@app.middleware("http")
async def _security_and_auth_middleware(request: Request, call_next):
    path = request.url.path

    # 允许预检与公开接口
    if request.method == "OPTIONS" or path in {"/", "/api/auth/login", "/api/auth/me", "/api/auth/setup"}:
        response = await call_next(request)
    else:
        auth = _get_auth_cfg()
        if auth["enabled"] and path.startswith("/api/"):
            user = _current_user_from_request(request)
            if not user:
                raise HTTPException(status_code=401, detail="未登录或会话已失效")
        response = await call_next(request)

    # 常见安全响应头
    response.headers["X-Content-Type-Options"] = "nosniff"
    response.headers["X-Frame-Options"] = "DENY"
    response.headers["Referrer-Policy"] = "no-referrer"
    response.headers["Cache-Control"] = "no-store"
    response.headers["Pragma"] = "no-cache"
    response.headers["Expires"] = "0"
    response.headers["Permissions-Policy"] = "geolocation=(), microphone=(), camera=()"
    response.headers["Content-Security-Policy"] = "default-src 'self'; img-src 'self' data:; style-src 'self' 'unsafe-inline'; script-src 'self' 'unsafe-inline'; connect-src 'self' ws: wss:"
    return response


def _ws_current_user(ws: WebSocket) -> Optional[dict]:
    auth = _get_auth_cfg()
    if not auth["enabled"]:
        return {"sub": "anonymous"}
    token = (ws.cookies.get(auth["cookie_name"]) or "").strip()
    if not token:
        return None
    try:
        return _jwt_decode(token)
    except HTTPException:
        return None

def _push_log_sync(queue: asyncio.Queue, msg: str):
    """从同步线程将日志推入 asyncio Queue"""
    if _event_loop and not _event_loop.is_closed():
        try:
            asyncio.run_coroutine_threadsafe(queue.put(msg), _event_loop)
        except Exception:
            pass


def _make_reg_log_cb():
    return lambda msg: _push_log_sync(_reg_log_queue, msg)


def _make_pool_log_cb():
    return lambda msg: _push_log_sync(_pool_log_queue, msg)


def _token_fingerprint(token: str) -> str:
    if not token:
        return ""
    return hashlib.sha256(token.encode("utf-8")).hexdigest()[:12]


def _build_probe_signature(base_url: str, token: str, target_type: str, proxy: str) -> str:
    raw = "|".join([
        str(base_url or "").strip().rstrip("/"),
        str(target_type or "").strip().lower(),
        _token_fingerprint(str(token or "").strip()),
        str(proxy or "").strip(),
    ])
    return hashlib.sha256(raw.encode("utf-8")).hexdigest()


def _make_reg_progress_cb():
    def cb(success: int, fail: int, total: int):
        _reg_state["success"] = success
        _reg_state["fail"] = fail
        _reg_state["total"] = total
    return cb


# ============================================================
# 静态文件 & 首页
# ============================================================

@app.get("/", response_class=HTMLResponse)
async def index():
    html_path = _TEMPLATES_DIR / "index.html"
    if not html_path.exists():
        raise HTTPException(status_code=404, detail="index.html 不存在")
    return HTMLResponse(content=html_path.read_text(encoding="utf-8"))


@app.post("/api/auth/login")
async def auth_login(body: dict = Body(...), request: Request = None):
    auth = _get_auth_cfg()
    username = str(body.get("username", "") or "").strip()
    password = str(body.get("password", "") or "")

    if not auth["enabled"]:
        return {"ok": True, "user": {"username": "anonymous"}, "auth_enabled": False, "need_setup": False}

    if not auth["password_hash"]:
        return {"ok": False, "auth_enabled": True, "need_setup": True}

    pwd_hash = hashlib.sha256(password.encode("utf-8")).hexdigest()
    user_ok = hmac.compare_digest(username, auth["username"])
    pwd_ok = hmac.compare_digest(pwd_hash, auth["password_hash"])
    if not (user_ok and pwd_ok):
        raise HTTPException(status_code=401, detail="用户名或密码错误")

    token = _issue_login_token(username)
    resp = JSONResponse({"ok": True, "user": {"username": username}, "auth_enabled": True, "need_setup": False})
    _set_auth_cookie(resp, token, request)
    return resp


@app.post("/api/auth/setup")
async def auth_setup(body: dict = Body(...), request: Request = None):
    auth = _get_auth_cfg()

    if not auth["enabled"]:
        return {"ok": True, "auth_enabled": False, "need_setup": False}

    if auth["password_hash"]:
        raise HTTPException(status_code=409, detail="密码已设置，请直接登录")

    username = str(body.get("username", auth["username"]) or "").strip() or auth["username"]
    password = str(body.get("password", "") or "")
    password_confirm = str(body.get("password_confirm", "") or "")

    if len(password) < 8:
        raise HTTPException(status_code=400, detail="密码长度至少 8 位")
    if password != password_confirm:
        raise HTTPException(status_code=400, detail="两次密码输入不一致")

    cfg = reg.load_config()
    cfg_auth = cfg.get("auth") or {}
    cfg_auth.update({
        "enabled": True,
        "username": username,
        "password_hash": hashlib.sha256(password.encode("utf-8")).hexdigest(),
    })
    cfg["auth"] = cfg_auth

    if not reg.save_config(cfg):
        raise HTTPException(status_code=500, detail="首次密码保存失败")

    token = _issue_login_token(username)
    resp = JSONResponse({"ok": True, "auth_enabled": True, "need_setup": False, "user": {"username": username}})
    _set_auth_cookie(resp, token, request)
    return resp


@app.post("/api/auth/logout")
async def auth_logout():
    resp = JSONResponse({"ok": True})
    _clear_auth_cookie(resp)
    return resp


@app.get("/api/auth/me")
async def auth_me(request: Request):
    auth = _get_auth_cfg()
    if not auth["enabled"]:
        return {
            "ok": True,
            "authenticated": True,
            "user": {"username": "anonymous"},
            "auth_enabled": False,
            "need_setup": False,
        }

    if not auth["password_hash"]:
        return {
            "ok": True,
            "authenticated": False,
            "auth_enabled": True,
            "need_setup": True,
            "user": {"username": auth["username"]},
        }

    user = _current_user_from_request(request)
    if not user:
        return {"ok": True, "authenticated": False, "auth_enabled": True, "need_setup": False}

    return {
        "ok": True,
        "authenticated": True,
        "auth_enabled": True,
        "need_setup": False,
        "user": {"username": user.get("sub", "")},
    }


# ============================================================
# 配置 API
# ============================================================

@app.get("/api/config")
async def get_config():
    return reg.load_config()


@app.post("/api/config")
async def save_config(config: dict = Body(...)):
    # 移除 _comment 字段保护
    config.pop("_comment", None)

    # 保护 auth 配置：若前端未带 auth，则沿用旧值，避免误清空
    old_cfg = reg.load_config()
    if "auth" not in config and isinstance(old_cfg.get("auth"), dict):
        config["auth"] = old_cfg["auth"]

    ok = reg.save_config(config)
    if not ok:
        raise HTTPException(status_code=500, detail="保存失败")

    # 如果守护进程已启用，同步更新其运行时配置（下次周期生效）
    pool = config.get("pool", {})
    if _pool_daemon["enabled"] and pool:
        _pool_daemon["config"].update({
            "base_url":     pool.get("base_url",     _pool_daemon["config"].get("base_url", "")),
            "token":        pool.get("token",         _pool_daemon["config"].get("token", "")),
            "target_type":  pool.get("target_type",  _pool_daemon["config"].get("target_type", "codex")),
            "target_count": int(pool.get("target_count", _pool_daemon["config"].get("target_count", 10))),
            "proxy":        pool.get("proxy",         _pool_daemon["config"].get("proxy", "")),
            "use_proxy":    bool(pool.get("use_proxy", _pool_daemon["config"].get("use_proxy", True))),
        })
        _pool_daemon["interval_min"] = max(1, int(pool.get("interval_min", _pool_daemon["interval_min"])))

    return {"ok": True}


# ============================================================
# 注册任务 API
# ============================================================

@app.post("/api/register/start")
async def register_start(body: dict = Body(...)):
    if _reg_state["running"]:
        raise HTTPException(status_code=409, detail="已有注册任务运行中")

    config = reg.load_config()
    count = int(body.get("count", 1))
    workers = int(body.get("workers", config.get("workers", 3)))
    proxy = str(body.get("proxy", "")).strip()

    # 重置状态
    _reg_state.update({
        "running": True,
        "success": 0,
        "fail": 0,
        "total": count,
        "start_time": time.time(),
    })
    stop_event = threading.Event()
    _reg_state["stop_event"] = stop_event

    log_cb = _make_reg_log_cb()
    progress_cb = _make_reg_progress_cb()

    def run_task():
        try:
            result = reg.run_batch_register(
                count=count,
                workers=workers,
                proxy=proxy,
                stop_event=stop_event,
                log_cb=log_cb,
                progress_cb=progress_cb,
                config=config,
            )
            _reg_state.update(result)
        except Exception as e:
            log_cb(f"[ERROR] 注册任务异常: {e}")
        finally:
            _reg_state["running"] = False
            _reg_state["stop_event"] = None
            log_cb("[注册] 任务结束")

    thread = threading.Thread(target=run_task, daemon=True)
    thread.start()

    return {"ok": True, "count": count, "workers": workers}


@app.post("/api/register/stop")
async def register_stop():
    stop_event = _reg_state.get("stop_event")
    if stop_event:
        stop_event.set()
    _reg_state["running"] = False
    return {"ok": True}


@app.get("/api/register/status")
async def register_status():
    elapsed = None
    if _reg_state.get("start_time"):
        elapsed = int(time.time() - _reg_state["start_time"])
    return {
        "running": _reg_state["running"],
        "success": _reg_state["success"],
        "fail": _reg_state["fail"],
        "total": _reg_state["total"],
        "elapsed": elapsed,
    }


# ============================================================
# 注册日志 WebSocket
# ============================================================

@app.websocket("/ws/register/logs")
async def ws_register_logs(ws: WebSocket):
    if not _ws_current_user(ws):
        await ws.close(code=1008)
        return
    await ws.accept()
    _reg_ws_clients.append(ws)
    try:
        while True:
            # 从队列取日志并推送
            try:
                msg = await asyncio.wait_for(_reg_log_queue.get(), timeout=1.0)
                for client in list(_reg_ws_clients):
                    try:
                        await client.send_text(json.dumps({"type": "log", "msg": msg}))
                    except Exception:
                        _reg_ws_clients.discard(client) if hasattr(_reg_ws_clients, 'discard') else None
            except asyncio.TimeoutError:
                # 发送心跳
                try:
                    await ws.send_text(json.dumps({"type": "ping"}))
                except Exception:
                    break
    except WebSocketDisconnect:
        pass
    finally:
        if ws in _reg_ws_clients:
            _reg_ws_clients.remove(ws)


# ============================================================
# 账号池 API
# ============================================================

@app.post("/api/pool/probe")
async def pool_probe(body: dict = Body(...)):
    if _pool_state["running"]:
        raise HTTPException(status_code=409, detail="已有池任务运行中")

    base_url = body.get("base_url", "").strip()
    token = body.get("token", "").strip()
    target_type = body.get("target_type", "codex")
    proxy = body.get("proxy", "").strip()
    use_proxy = bool(body.get("use_proxy", True))
    effective_proxy = proxy if use_proxy else ""

    if not base_url or not token:
        raise HTTPException(status_code=400, detail="base_url 和 token 不能为空")

    _pool_state["running"] = True
    _pool_state["task"] = "probe"
    log_cb = _make_pool_log_cb()

    def run_task():
        try:
            result = reg.run_pool_probe(base_url, token, target_type, effective_proxy, log_cb=log_cb)
            log_cb(f"[Pool] 探测结果: 总={result.get('total')}, 目标={result.get('target')}, 401={result.get('invalid_count')}")
        except Exception as e:
            log_cb(f"[ERROR] 探测异常: {e}")
        finally:
            _pool_state["running"] = False

    threading.Thread(target=run_task, daemon=True).start()
    return {"ok": True, "task": "probe"}


@app.post("/api/pool/clean")
async def pool_clean(body: dict = Body(...)):
    if _pool_state["running"]:
        raise HTTPException(status_code=409, detail="已有池任务运行中")

    base_url = body.get("base_url", "").strip()
    token = body.get("token", "").strip()
    target_type = body.get("target_type", "codex")
    proxy = body.get("proxy", "").strip()
    use_proxy = bool(body.get("use_proxy", True))
    effective_proxy = proxy if use_proxy else ""
    probe_result = body.get("probe_result")
    probe_signature = str(body.get("probe_signature", "") or "")
    probe_ts_raw = body.get("probe_ts")

    if not base_url or not token:
        raise HTTPException(status_code=400, detail="base_url 和 token 不能为空")

    _pool_state["running"] = True
    _pool_state["task"] = "clean"
    log_cb = _make_pool_log_cb()

    def run_task():
        try:
            now_ts = int(time.time())
            expected_signature = _build_probe_signature(base_url, token, target_type, effective_proxy)
            use_cached_probe = False
            if isinstance(probe_result, dict):
                try:
                    probe_ts = int(probe_ts_raw)
                except Exception:
                    probe_ts = 0
                is_fresh = probe_ts > 0 and (now_ts - probe_ts) <= _PROBE_RESULT_TTL_SEC
                signature_ok = bool(probe_signature) and probe_signature == expected_signature
                invalid_list_ok = isinstance(probe_result.get("invalid_401"), list)
                use_cached_probe = is_fresh and signature_ok and invalid_list_ok

            if use_cached_probe:
                log_cb("[Pool] 复用检查结果执行清理，跳过重复探测")
                result = reg.run_pool_clean_with_probe_result(
                    base_url=base_url,
                    token=token,
                    probe_result=probe_result,
                    proxy=effective_proxy,
                    log_cb=log_cb,
                )
            else:
                if isinstance(probe_result, dict):
                    log_cb("[Pool] 检查结果不可用或已过期，回退为重新探测后清理")
                result = reg.run_pool_clean(base_url, token, target_type, effective_proxy, log_cb=log_cb)
            log_cb(f"[Pool] 清理完成: 删除={result.get('deleted')}, 失败={result.get('delete_fail')}")
        except Exception as e:
            log_cb(f"[ERROR] 清理异常: {e}")
        finally:
            _pool_state["running"] = False

    threading.Thread(target=run_task, daemon=True).start()
    return {"ok": True, "task": "clean"}


@app.post("/api/pool/fill")
async def pool_fill(body: dict = Body(...)):
    if _pool_state["running"]:
        raise HTTPException(status_code=409, detail="已有池任务运行中")

    count = int(body.get("count", 1))
    base_url = body.get("base_url", "").strip()
    pool_token = body.get("token", "").strip()
    proxy = body.get("proxy", "").strip()
    use_proxy = bool(body.get("use_proxy", True))
    effective_proxy = proxy if use_proxy else ""
    target_type = body.get("target_type", "codex")
    target_count = int(body.get("target_count", 0))
    config = reg.load_config()

    _pool_state["running"] = True
    _pool_state["task"] = "fill"
    log_cb = _make_pool_log_cb()
    stop_event = threading.Event()
    _pool_state["stop_event"] = stop_event

    def progress_cb(s, f, t):
        pass

    def run_task():
        try:
            result = reg.run_pool_fill(
                fill_count=count,
                base_url=base_url,
                pool_token=pool_token,
                stop_event=stop_event,
                log_cb=log_cb,
                progress_cb=progress_cb,
                config=config,
                proxy=effective_proxy,
                target_count=target_count,
                target_type=target_type,
            )
            log_cb(f"[Pool] 补号完成: 成功={result.get('success')}, 失败={result.get('fail')}")
        except Exception as e:
            log_cb(f"[ERROR] 补号异常: {e}")
        finally:
            _pool_state["running"] = False
            _pool_state.pop("stop_event", None)

    threading.Thread(target=run_task, daemon=True).start()
    return {"ok": True, "task": "fill", "count": count}


@app.post("/api/pool/overview")
async def pool_overview_api(body: dict = Body(...)):
    base_url = body.get("base_url", "").strip()
    token = body.get("token", "").strip()
    target_type = body.get("target_type", "codex")
    proxy = body.get("proxy", "").strip()
    use_proxy = bool(body.get("use_proxy", True))
    effective_proxy = proxy if use_proxy else ""

    if not base_url or not token:
        raise HTTPException(status_code=400, detail="base_url 和 token 不能为空")

    result = await asyncio.get_event_loop().run_in_executor(
        None, lambda: reg.get_pool_overview(base_url, token, target_type, effective_proxy)
    )
    return result


@app.get("/api/pool/accounts")
async def pool_accounts(base_url: str, token: str,
                        target_type: str = "codex", proxy: str = "", use_proxy: bool = True):
    if not base_url or not token:
        raise HTTPException(status_code=400, detail="base_url 和 token 不能为空")
    effective_proxy = proxy if use_proxy else ""
    result = await asyncio.get_event_loop().run_in_executor(
        None, lambda: reg.get_pool_accounts(base_url, token, target_type, effective_proxy)
    )
    return result


@app.post("/api/pool/sync")
async def pool_sync(body: dict = Body(...)):
    base_url = body.get("base_url", "").strip()
    token = body.get("token", "").strip()
    target_type = body.get("target_type", "codex")
    proxy = body.get("proxy", "").strip()
    use_proxy = bool(body.get("use_proxy", True))
    effective_proxy = proxy if use_proxy else ""
    target_count = int(body.get("target_count", 0))

    if not base_url or not token:
        raise HTTPException(status_code=400, detail="base_url 和 token 不能为空")

    log_cb = _make_pool_log_cb()
    result = await asyncio.get_event_loop().run_in_executor(
        None, lambda: reg.sync_local_remote(base_url, token, target_type, proxy=effective_proxy, log_cb=log_cb, target_count=target_count)
    )
    return result


@app.get("/api/pool/task-status")
async def pool_task_status():
    return {
        "running": _pool_state["running"],
        "task": _pool_state.get("task", ""),
    }


@app.post("/api/pool/inspect")
async def pool_inspect(body: dict = Body(...)):
    """同步探测账号池，返回检查结果（供前端渲染确认卡）"""
    base_url = body.get("base_url", "").strip()
    token = body.get("token", "").strip()
    target_type = body.get("target_type", "codex")
    target_count = int(body.get("target_count", 100))
    proxy = body.get("proxy", "").strip()
    use_proxy = bool(body.get("use_proxy", True))
    effective_proxy = proxy if use_proxy else ""

    if not base_url or not token:
        raise HTTPException(status_code=400, detail="base_url 和 token 不能为空")

    log_cb = _make_pool_log_cb()
    log_cb("[Pool] 开始检查失效账号...")
    loop = asyncio.get_event_loop()
    result = await loop.run_in_executor(
        None,
        lambda: reg.run_pool_probe(base_url, token, target_type, effective_proxy, log_cb=log_cb),
    )

    if not result.get("ok"):
        raise HTTPException(status_code=500, detail=result.get("error", "探测失败"))

    log_cb(f"[Pool] 检查结束: 总={result.get('total')}, 目标={result.get('target')}, 401={result.get('invalid_count')}")
    probe_ts = int(time.time())
    probe_signature = _build_probe_signature(base_url, token, target_type, effective_proxy)
    valid = result["target"] - result["invalid_count"]
    gap = max(0, target_count - valid)
    return {
        **result,
        "valid_count": valid,
        "gap": gap,
        "target_count": target_count,
        "probe_ts": probe_ts,
        "probe_signature": probe_signature,
    }


# ============================================================
# 池管理日志 WebSocket
# ============================================================

@app.websocket("/ws/pool/logs")
async def ws_pool_logs(ws: WebSocket):
    if not _ws_current_user(ws):
        await ws.close(code=1008)
        return
    await ws.accept()
    _pool_ws_clients.append(ws)
    try:
        while True:
            try:
                msg = await asyncio.wait_for(_pool_log_queue.get(), timeout=1.0)
                for client in list(_pool_ws_clients):
                    try:
                        await client.send_text(json.dumps({"type": "log", "msg": msg}))
                    except Exception:
                        pass
            except asyncio.TimeoutError:
                try:
                    await ws.send_text(json.dumps({"type": "ping"}))
                except Exception:
                    break
    except WebSocketDisconnect:
        pass
    finally:
        if ws in _pool_ws_clients:
            _pool_ws_clients.remove(ws)


# ============================================================
# 代理管理 API
# ============================================================

@app.get("/api/proxy/fetch")
async def proxy_fetch():
    config = reg.load_config()
    fallback_proxy = config.get("proxy", "")
    proxies = await asyncio.get_event_loop().run_in_executor(
        None, lambda: reg.fetch_free_proxies(proxy=fallback_proxy)
    )
    return {"ok": True, "proxies": proxies, "count": len(proxies)}


@app.post("/api/proxy/test")
async def proxy_test(body: dict = Body(...)):
    proxies = body.get("proxies", [])
    target_url = body.get("target_url", "https://httpbin.org/ip")
    timeout = int(body.get("timeout", 5))

    if not proxies:
        raise HTTPException(status_code=400, detail="proxies 不能为空")
    if len(proxies) > 50:
        proxies = proxies[:50]

    results = await asyncio.get_event_loop().run_in_executor(
        None,
        lambda: reg.test_proxies_concurrent(proxies, target_url, timeout, max_workers=20),
    )
    return {"ok": True, "results": results}


# ============================================================
# 结果查看 API
# ============================================================

@app.get("/api/results")
async def get_results():
    config = reg.load_config()
    accounts = reg.read_registered_accounts(config)
    return {"ok": True, "accounts": accounts, "count": len(accounts)}


@app.get("/api/tokens")
async def get_tokens():
    config = reg.load_config()
    tokens = reg.list_codex_tokens(config)
    return {"ok": True, "tokens": tokens, "count": len(tokens)}


@app.get("/api/tokens/ak")
async def get_ak():
    config = reg.load_config()
    content = reg.read_token_file("ak.txt", config)
    return PlainTextResponse(content=content)


@app.get("/api/tokens/rk")
async def get_rk():
    config = reg.load_config()
    content = reg.read_token_file("rk.txt", config)
    return PlainTextResponse(content=content)


@app.get("/api/tokens/download/ak")
async def download_ak():
    config = reg.load_config()
    ak_file = config.get("ak_file", "ak.txt")
    if not os.path.isabs(ak_file):
        ak_file = str(_BASE_DIR / ak_file)
    if not os.path.exists(ak_file):
        raise HTTPException(status_code=404, detail="ak.txt 不存在")
    return FileResponse(ak_file, filename="ak.txt", media_type="text/plain")


@app.get("/api/tokens/download/rk")
async def download_rk():
    config = reg.load_config()
    rk_file = config.get("rk_file", "rk.txt")
    if not os.path.isabs(rk_file):
        rk_file = str(_BASE_DIR / rk_file)
    if not os.path.exists(rk_file):
        raise HTTPException(status_code=404, detail="rk.txt 不存在")
    return FileResponse(rk_file, filename="rk.txt", media_type="text/plain")


# ============================================================
# 号池守护进程
# ============================================================

_pool_daemon: Dict[str, Any] = {
    "enabled": False,
    "interval_min": 30,
    "next_run_ts": None,
    "last_run_ts": None,
    "running_now": False,
    "config": {},  # {base_url, token, target_type, target_count, proxy}
}
_pool_daemon_timer: Optional[threading.Timer] = None


def _run_daemon_once():
    """守护进程单次执行"""
    global _pool_daemon_timer

    if not _pool_daemon["enabled"]:
        return

    _pool_daemon["running_now"] = True
    log_cb = _make_pool_log_cb()
    try:
        cfg = _pool_daemon["config"]
        proxy = reg._proxy_pool.get_best(cfg.get("proxy", "")) if cfg.get("use_proxy", True) else ""
        stop_event = threading.Event()
        reg.run_pool_maintain_cycle(
            base_url=cfg.get("base_url", ""),
            token=cfg.get("token", ""),
            target_type=cfg.get("target_type", "codex"),
            target_count=int(cfg.get("target_count", 10)),
            stop_event=stop_event,
            log_cb=log_cb,
            config=reg.load_config(),
            proxy=proxy,
        )
    except Exception as e:
        log_cb(f"[Daemon] 执行异常: {e}")
    finally:
        _pool_daemon["running_now"] = False
        _pool_daemon["last_run_ts"] = time.time()
        if _pool_daemon["enabled"]:
            interval_sec = _pool_daemon["interval_min"] * 60
            _pool_daemon["next_run_ts"] = time.time() + interval_sec
            _pool_daemon_timer = threading.Timer(interval_sec, _run_daemon_once)
            _pool_daemon_timer.daemon = True
            _pool_daemon_timer.start()


@app.post("/api/pool/daemon/start")
async def pool_daemon_start(body: dict = Body(...)):
    global _pool_daemon_timer

    base_url = body.get("base_url", "").strip()
    token = body.get("token", "").strip()
    if not base_url or not token:
        raise HTTPException(status_code=400, detail="base_url 和 token 不能为空")

    # 停止旧 timer
    if _pool_daemon_timer and _pool_daemon_timer.is_alive():
        _pool_daemon_timer.cancel()

    interval_min = max(1, int(body.get("interval_min", 30)))
    _pool_daemon.update({
        "enabled": True,
        "interval_min": interval_min,
        "config": {
            "base_url": base_url,
            "token": token,
            "target_type": body.get("target_type", "codex"),
            "target_count": int(body.get("target_count", 10)),
            "proxy": body.get("proxy", "").strip(),
            "use_proxy": bool(body.get("use_proxy", True)),
        },
    })

    # 立即执行一次（在后台线程）
    _pool_daemon["next_run_ts"] = None
    t = threading.Thread(target=_run_daemon_once, daemon=True)
    t.start()

    return {"ok": True, "interval_min": interval_min}


@app.post("/api/pool/daemon/stop")
async def pool_daemon_stop():
    global _pool_daemon_timer

    _pool_daemon["enabled"] = False
    _pool_daemon["next_run_ts"] = None
    if _pool_daemon_timer and _pool_daemon_timer.is_alive():
        _pool_daemon_timer.cancel()
        _pool_daemon_timer = None

    return {"ok": True}


@app.get("/api/pool/daemon/status")
async def pool_daemon_status():
    remaining = None
    if _pool_daemon["next_run_ts"]:
        remaining = max(0, int(_pool_daemon["next_run_ts"] - time.time()))
    return {
        "enabled": _pool_daemon["enabled"],
        "running_now": _pool_daemon["running_now"],
        "interval_min": _pool_daemon["interval_min"],
        "next_run_ts": _pool_daemon["next_run_ts"],
        "last_run_ts": _pool_daemon["last_run_ts"],
        "remaining_sec": remaining,
        "config": _pool_daemon["config"],
    }


@app.post("/api/pool/daemon/run-once")
async def pool_daemon_run_once(body: dict = Body(default={})):
    """立即触发一次维护周期（不影响守护进程定时器）"""
    cfg = _pool_daemon["config"]
    # 允许临时覆盖配置
    if body.get("base_url"):
        cfg = {**cfg, **body}

    base_url = cfg.get("base_url", "").strip()
    token = cfg.get("token", "").strip()
    if not base_url or not token:
        raise HTTPException(status_code=400, detail="请先配置 base_url 和 token")

    if _pool_daemon["running_now"]:
        raise HTTPException(status_code=409, detail="守护进程正在运行中")

    log_cb = _make_pool_log_cb()

    def run_task():
        _pool_daemon["running_now"] = True
        try:
            proxy = reg._proxy_pool.get_best(cfg.get("proxy", "")) if cfg.get("use_proxy", True) else ""
            stop_event = threading.Event()
            reg.run_pool_maintain_cycle(
                base_url=cfg.get("base_url", ""),
                token=cfg.get("token", ""),
                target_type=cfg.get("target_type", "codex"),
                target_count=int(cfg.get("target_count", 10)),
                stop_event=stop_event,
                log_cb=log_cb,
                config=reg.load_config(),
                proxy=proxy,
            )
        except Exception as e:
            log_cb(f"[Daemon] 执行异常: {e}")
        finally:
            _pool_daemon["running_now"] = False
            _pool_daemon["last_run_ts"] = time.time()

    threading.Thread(target=run_task, daemon=True).start()
    return {"ok": True}


# ============================================================
# 代理池端点
# ============================================================

@app.post("/api/proxy/pool/update")
async def proxy_pool_update(body: dict = Body(...)):
    """接收前端测试结果，更新内存代理池"""
    results = body.get("results", [])
    if not isinstance(results, list):
        raise HTTPException(status_code=400, detail="results 必须是列表")
    reg._proxy_pool.update(results)
    best = reg._proxy_pool.get_best()
    return {"ok": True, "count": len(results), "best": best}


@app.get("/api/proxy/active")
async def proxy_active():
    """返回当前最优代理"""
    config = reg.load_config()
    fallback = config.get("proxy", "")
    best = reg._proxy_pool.get_best(fallback)
    source = "free_pool" if reg._proxy_pool.get_best() else ("user_config" if fallback else "direct")
    return {
        "proxy": best,
        "source": source,
        "pool_size": len(reg._proxy_pool.get_all()),
    }

@app.get("/healthyz")
async def proxy_active():
    
    return {
        "healthyz": "ok"
    }

# ============================================================
# 主入口
# ============================================================

if __name__ == "__main__":
    uvicorn.run(
        "web_app:app",
        host="0.0.0.0",
        port=8080,
        reload=False,
        log_level="info",
    )