import os  # 导入系统模块用于处理环境变量。 
import sys  # 导入 sys 用于获取当前 Python 解释器。 
import time  # 导入时间模块用于记录状态时间。 
import json  # 导入 JSON 模块用于会话 token 载荷。 
import hmac  # 导入 HMAC 模块用于 token 签名比对。 
import base64  # 导入 base64 模块用于 token 编码。 
import hashlib  # 导入哈希模块用于 token 签名。 
import threading  # 导入线程模块用于后台执行任务。 
import subprocess  # 导入 subprocess 用于为每个账号隔离执行购买流程。 
from contextlib import redirect_stdout, redirect_stderr  # 导入输出重定向工具用于捕获日志。 
from typing import Optional, Dict, List  # 导入类型注解辅助类型。 

BASE_DIR = os.path.dirname(os.path.abspath(__file__))  # 当前火山模块目录。 
PARENT_DIR = os.path.dirname(BASE_DIR)  # wwwroot 目录。 
for path in (BASE_DIR, PARENT_DIR):  # 固定导入路径，避免启动目录变化导致导错 config。 
    if path in sys.path: 
        sys.path.remove(path) 
    sys.path.insert(0, path) 

from fastapi import FastAPI, Request  # 导入 FastAPI 与 Request 用于构建接口与读取会话。 
from fastapi.responses import JSONResponse  # 导入响应类型用于返回 JSON。 
from starlette.middleware.sessions import SessionMiddleware  # 导入 Session 中间件用于登录态管理。 
from pydantic import BaseModel  # 导入数据模型用于校验请求体。 

from volcengine import main as volc_main  # 导入火山引擎主流程入口。 
from volcengine import config as _config  # 导入配置模块用于读取会话参数。 
from volcengine import customer_guard  # 导入火山引擎客户归属校验服务。 

MAX_LOG_LINES = 2000  # 日志缓冲区最多保留的行数。 
LOG_BUFFER_BY_USER: Dict[str, List[str]] = {}  # 按账号区分的日志缓冲区字典。 
LOG_LOCK = threading.Lock()  # 日志锁用于多线程安全写入。 
STATE = {"running": False, "error": "", "started_at": None, "finished_at": None, "user": ""}  # 运行状态信息字典。 
STATE_BY_USER: Dict[str, Dict] = {}  # 按账号隔离的运行状态。 
STATE_LOCK = threading.Lock()  # 状态锁用于多线程安全更新。 
TASK_THREAD_BY_USER: Dict[str, threading.Thread] = {}  # 按账号记录后台任务线程。 
TASK_PROCESS_BY_USER: Dict[str, subprocess.Popen] = {}  # 按账号记录购买子进程。 
TASK_TIMEOUT_SECONDS = int(os.getenv("VOLC_TASK_TIMEOUT_SECONDS", "21600") or "21600")  # 任务最大占用时长，默认 6 小时。 

app = FastAPI(title="wawacloud volcengine")  # 创建火山引擎 API 服务实例。 
app.add_middleware(
    SessionMiddleware,  # 使用 Starlette 提供的 Session 中间件。 
    secret_key=_config.APP_SESSION_SECRET,  # 设置会话密钥用于 Cookie 加密。 
    session_cookie=_config.APP_SESSION_COOKIE,  # 设置会话 Cookie 名称。 
    same_site="lax",  # 设置 SameSite 保护跨站请求。 
    https_only=_config.APP_SESSION_HTTPS_ONLY,  # 设置是否仅 HTTPS 发送 Cookie。 
)  # 会话中间件配置结束。 


class RunRequest(BaseModel):  # 运行接口请求体模型定义。 
    regions: List[str] = []  # 选择的地区列表。 
    region_counts: Optional[Dict[str, int]] = None  # 地区到数量映射配置。 
    server_count: Optional[int] = None  # 统一购买数量。 
    ak: Optional[str] = None  # 可选 AK 字段用于覆盖。 
    sk: Optional[str] = None  # 可选 SK 字段用于覆盖。 
    manual_confirm_customer: Optional[bool] = False  # 代理后台掉线时的人工确认放行。 


def _get_log_buffer(user: str) -> List[str]:  # 获取指定账号的日志缓冲区。 
    key = user or "_system"  # 为空时归档到系统日志。 
    with LOG_LOCK:  # 加锁保证线程安全。 
        return LOG_BUFFER_BY_USER.setdefault(key, [])  # 返回日志列表。 


def _append_log(line: str, user: str = "") -> None:  # 追加日志行到缓冲区。 
    text = (line or "").rstrip()  # 取出文本并去掉行尾换行。 
    if not text:  # 若为空则直接返回。 
        return  # 空行不写入日志缓冲。 
    buffer = _get_log_buffer(user)  # 获取日志列表。 
    with LOG_LOCK:  # 加锁保证线程安全。 
        buffer.append(text)  # 追加日志行到缓冲区。 
        overflow = len(buffer) - MAX_LOG_LINES  # 计算超过最大行数的数量。 
        if overflow > 0:  # 若超出最大行数则裁剪。 
            del buffer[:overflow]  # 删除最旧的日志行保持缓冲大小。 


class _LogWriter:  # 日志写入器用于捕获 stdout/stderr。 
    def __init__(self, user: str = "") -> None:  # 初始化日志写入器。 
        self._cache = ""  # 缓存未完成的行内容。 
        self._lock = threading.Lock()  # 日志写入锁用于多线程安全写入。 
        self._user = user  # 记录日志所属账号。 

    def write(self, data: str) -> None:  # 写入数据到缓冲区的方法。 
        if not data:  # 若数据为空则直接返回。 
            return  # 空数据不处理。 
        with self._lock:  # 加锁保证多线程写入安全。 
            self._cache += data  # 将新数据追加到缓存。 
            while "\n" in self._cache:  # 当缓存中包含换行时拆分行。 
                line, self._cache = self._cache.split("\n", 1)  # 拆分出一行并保留剩余缓存。 
                _append_log(line, self._user)  # 将拆分出的行写入日志缓冲。 

    def flush(self) -> None:  # 标准输出刷新接口占位。 
        with self._lock:  # 加锁保证刷新与写入一致。 
            if self._cache:  # 若缓存仍有内容则写入。 
                _append_log(self._cache, self._user)  # 写入残留缓存内容。 
                self._cache = ""  # 清空缓存内容。 


def _blank_state(user: str = "") -> Dict:  # 创建账号独立状态。 
    return {"running": False, "error": "", "started_at": None, "finished_at": None, "user": user} 


def _state_key(user: str) -> str:  # 生成状态字典键。 
    return user or "_system" 


def _ensure_state_unlocked(user: str) -> Dict:  # 在已持锁时获取账号状态。 
    key = _state_key(user) 
    return STATE_BY_USER.setdefault(key, _blank_state(user)) 


def _set_state(user: str, running=None, error=None, started_at=None, finished_at=None) -> None:  # 更新指定账号状态。 
    with STATE_LOCK:  # 加锁保证状态更新线程安全。 
        state = _ensure_state_unlocked(user)  # 获取账号状态。 
        if running is not None:  # 若传入运行状态则更新。 
            state["running"] = running  # 更新运行状态字段。 
        if error is not None:  # 若传入错误信息则更新。 
            state["error"] = error  # 更新错误信息字段。 
        if started_at is not None:  # 若传入开始时间则更新。 
            state["started_at"] = started_at  # 更新开始时间字段。 
        if finished_at is not None:  # 若传入结束时间则更新。 
            state["finished_at"] = finished_at  # 更新结束时间字段。 
        state["user"] = user  # 状态始终绑定当前账号。 


def _summarize_guard_exception(provider: str, exc: Exception) -> str:  # 压缩客户归属校验异常。 
    text = str(exc) 
    lower = text.lower() 
    if "signature" in lower or "auth" in lower or "access denied" in lower or "invalid" in lower: 
        return f"{provider} AK/SK 认证失败，请检查访问密钥" 
    short = text if len(text) <= 180 else text[:180] + "..." 
    return f"{provider}客户归属校验失败: {short}" 


def _manual_confirm_response(message: str, user: str) -> JSONResponse:  # 返回火山引擎人工确认弹窗提示。 
    _append_log(f"step: customer_guard_need_manual_confirm {message}", user) 
    _set_state(user, running=False, error=message, finished_at=time.time()) 
    return JSONResponse( 
        { 
            "ok": False, 
            "message": message, 
            "manual_confirm_available": True, 
            "manual_confirm_text": "已手动确认用户继续执行？", 
        }, 
        status_code=403, 
    ) 


def _cleanup_stale_task_unlocked(user: str) -> None:  # 在已持有 STATE_LOCK 时清理失效运行状态。 
    state = _ensure_state_unlocked(user)  # 获取账号状态。 
    if not state["running"]:  # 未运行时无需处理。 
        return 
    now = time.time()  # 当前时间戳。 
    started_at = state.get("started_at")  # 读取任务开始时间。 
    timed_out = False  # 是否超过最大占用时间。 
    try: 
        timed_out = bool(started_at and TASK_TIMEOUT_SECONDS > 0 and now - float(started_at) > TASK_TIMEOUT_SECONDS) 
    except (TypeError, ValueError): 
        timed_out = False 
    age_seconds = 0.0  # 任务已运行时长。 
    try: 
        age_seconds = now - float(started_at or now) 
    except (TypeError, ValueError): 
        age_seconds = 0.0 
    key = _state_key(user)  # 获取账号键。 
    thread = TASK_THREAD_BY_USER.get(key)  # 读取账号线程。 
    process = TASK_PROCESS_BY_USER.get(key)  # 读取账号子进程。 
    process_done = process is not None and process.poll() is not None  # 子进程已结束但状态未释放。 
    thread_dead = thread is not None and not thread.is_alive()  # 线程已结束但状态未释放。 
    thread_missing = thread is None and state.get("started_at") and age_seconds > 5  # 有运行状态但 5 秒后仍没有线程引用。 
    if not thread_dead and not thread_missing and not process_done and not timed_out:  # 仍有有效线程运行。 
        return 
    old_user = str(state.get("user") or user or "")  # 记录原运行账号。 
    if old_user:  # 给原账号日志写入自动解锁原因。 
        if timed_out: 
            _append_log(f"step: api_unlock timeout_after={TASK_TIMEOUT_SECONDS}s", old_user) 
        elif process_done: 
            _append_log("step: api_unlock process_done", old_user) 
        else: 
            _append_log("step: api_unlock stale_thread", old_user) 
    if timed_out and process is not None and process.poll() is None:  # 超时时终止仍在运行的子进程。 
        try: 
            process.terminate() 
        except Exception: 
            pass 
    state.update(  # 自动释放卡住的运行状态。 
        {
            "running": False,
            "error": "上一个任务状态已自动释放，请重新启动",
            "finished_at": now,
            "user": "",
        }
    )
    TASK_THREAD_BY_USER.pop(key, None)  # 清空线程引用。 
    TASK_PROCESS_BY_USER.pop(key, None)  # 清空子进程引用。 


def _reserve_task(user: str) -> tuple:  # 原子占用当前账号任务运行权，不阻塞其他账号。 
    with STATE_LOCK:  # 检查与占用必须在同一把锁里完成。 
        _cleanup_stale_task_unlocked(user)  # 启动前先释放该账号已经失效的占用。 
        state = _ensure_state_unlocked(user)  # 获取当前账号状态。 
        if state["running"]:  # 当前账号已有任务在执行时拒绝重复启动。 
            return False, user  # 返回占用失败与当前账号。 
        state.update(  # 在后台线程启动前立即标记运行中，封住同账号重复请求窗口。 
            {
                "running": True,
                "error": "",
                "started_at": time.time(),
                "finished_at": None,
                "user": user,
            }
        )
        return True, user  # 返回占用成功与当前账号。 


def _get_session_user(request: Request) -> str:  # 获取当前登录账号。 
    token_user = _token_user(request)  # 优先使用标签页独立 token，避免同浏览器多账号串号。 
    if token_user: 
        return token_user 
    if not hasattr(request, "session"):  # 若请求对象不含会话则返回空。 
        return ""  # 返回空字符串表示未登录。 
    return str(request.session.get("user") or "")  # 读取会话中的用户名并返回。 


def _require_login(request: Request):  # 登录校验辅助函数。 
    user = _get_session_user(request)  # 获取当前会话用户。 
    if not user:  # 若用户为空则返回未登录响应。 
        return JSONResponse({"ok": False, "message": "未登录或登录已过期"}, status_code=401)  # 返回未登录错误。 
    return None  # 登录通过时返回空。 


def _serialize_region_counts(region_counts: Dict[str, int]) -> str:
    items: List[str] = []
    for region, count in region_counts.items():
        region_value = str(region or "").strip()
        if not region_value:
            continue
        try:
            count_value = int(count)
        except (TypeError, ValueError):
            count_value = 1
        count_value = max(1, min(2, count_value))
        items.append(f"{region_value}={count_value}")
    return ",".join(items)


def _build_task_env(payload: RunRequest) -> Dict[str, str]:
    env = os.environ.copy()
    ak_value = (payload.ak or "").strip()
    sk_value = (payload.sk or "").strip()
    if ak_value and sk_value:
        env["VOLC_AK"] = ak_value
        env["VOLC_SK"] = sk_value
    else:
        env.pop("VOLC_AK", None)
        env.pop("VOLC_SK", None)

    if payload.region_counts:
        env["VOLC_REGION_COUNTS"] = _serialize_region_counts(payload.region_counts)
        env.pop("VOLC_REGIONS", None)
    elif payload.regions:
        regions = [str(item).strip() for item in payload.regions if str(item).strip()]
        if regions:
            env["VOLC_REGIONS"] = ",".join(regions)
        else:
            env.pop("VOLC_REGIONS", None)
        env.pop("VOLC_REGION_COUNTS", None)
    else:
        env.pop("VOLC_REGION_COUNTS", None)
        env.pop("VOLC_REGIONS", None)

    if payload.server_count is not None:
        try:
            count_value = int(payload.server_count)
        except (TypeError, ValueError):
            count_value = 0
        if count_value > 0:
            env["VOLC_SERVER_COUNT"] = str(count_value)
        else:
            env.pop("VOLC_SERVER_COUNT", None)
    else:
        env.pop("VOLC_SERVER_COUNT", None)
    return env


def _b64_decode(value: str) -> bytes:  # URL 安全 base64 解码。 
    padding = "=" * (-len(value) % 4)  # 补齐 base64 填充。 
    return base64.urlsafe_b64decode((value + padding).encode("ascii")) 


def _b64_encode(raw: bytes) -> str:  # URL 安全 base64 编码。 
    return base64.urlsafe_b64encode(raw).decode("ascii").rstrip("=") 


def _token_signature(payload_b64: str) -> str:  # 计算标签页 token 签名。 
    secret = str(_config.APP_SESSION_SECRET or "CHANGE_ME_TO_RANDOM_SECRET").encode("utf-8") 
    digest = hmac.new(secret, payload_b64.encode("ascii"), hashlib.sha256).digest() 
    return _b64_encode(digest) 


def _token_user(request: Request) -> str:  # 从请求头读取标签页 token 对应账号。 
    token = request.headers.get("x-wawa-session") or request.headers.get("x-wawacloud-session") or "" 
    token = token.strip() 
    if not token or "." not in token: 
        return "" 
    payload_b64, signature = token.rsplit(".", 1) 
    if not payload_b64 or not signature: 
        return "" 
    if not hmac.compare_digest(signature, _token_signature(payload_b64)): 
        return "" 
    try: 
        payload = json.loads(_b64_decode(payload_b64).decode("utf-8")) 
    except Exception: 
        return "" 
    try: 
        if int(payload.get("exp") or 0) < int(time.time()): 
            return "" 
    except (TypeError, ValueError): 
        return "" 
    return str(payload.get("user") or "").strip() 


def _run_task(user: str, payload: RunRequest) -> None:  # 后台执行任务的函数。 
    key = _state_key(user)  # 获取账号键。 
    process = None  # 子进程对象。 
    try:  # 捕获执行过程中的异常。 
        env = _build_task_env(payload)  # 为该账号构建独立环境变量，避免并发任务互相串 AK/地区。 
        base_dir = os.path.dirname(os.path.abspath(__file__))  # 当前火山模块目录。 
        parent_dir = os.path.dirname(base_dir)  # wwwroot 目录，用于 python -m 导入 volcengine 包。 
        python_path = env.get("PYTHONPATH", "")  # 读取已有 PYTHONPATH。 
        env["PYTHONPATH"] = os.pathsep.join([parent_dir, base_dir, python_path]) if python_path else os.pathsep.join([parent_dir, base_dir]) 
        env["PYTHONUNBUFFERED"] = "1"  # 强制子进程实时输出日志，避免购买结束后才一次性刷新。 
        command = [sys.executable, "-u", "-m", "volcengine.main"]  # 使用当前服务解释器启动隔离购买进程。 
        _append_log(f"step: api_start regions={payload.regions or []}", user)  # 输出启动日志。 
        process = subprocess.Popen(  # 启动独立子进程执行购买流程。 
            command,
            cwd=parent_dir,
            env=env,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            bufsize=1,
        ) 
        with STATE_LOCK:  # 保存子进程引用用于状态健康检查。 
            TASK_PROCESS_BY_USER[key] = process 
        if process.stdout:  # 持续读取子进程输出写入当前账号日志。 
            for line in process.stdout: 
                _append_log(line, user) 
        exit_code = process.wait()  # 等待子进程结束。 
        if exit_code != 0:  # 若返回非 0 则记录异常状态。 
            _set_state(user, error=f"exit_code={exit_code}")  # 更新错误状态。 
            _append_log(f"step: api_error exit_code={exit_code}", user)  # 输出错误日志。 
        else:  # 正常完成。 
            _append_log("step: api_done", user)  # 输出完成日志。 
    except Exception as exc:  # 捕获执行异常并记录。 
        _append_log(f"step: api_exception {exc}", user)  # 输出异常日志。 
        _set_state(user, error=str(exc))  # 保存错误信息到状态。 
    finally:  # 无论成功失败均更新结束状态。 
        _set_state(user, running=False, finished_at=time.time())  # 设置任务结束状态。 
        with STATE_LOCK:  # 清理当前线程引用。 
            if TASK_THREAD_BY_USER.get(key) is threading.current_thread(): 
                TASK_THREAD_BY_USER.pop(key, None) 
            if process is not None and TASK_PROCESS_BY_USER.get(key) is process: 
                TASK_PROCESS_BY_USER.pop(key, None) 


@app.post("/api/run")  # 定义启动任务的接口路由。 
def run_task(payload: RunRequest, request: Request):  # 启动任务接口函数。 
    auth_error = _require_login(request)  # 校验登录状态。 
    if auth_error:  # 未登录则返回错误响应。 
        return auth_error  # 返回未登录响应。 
    user = _get_session_user(request)  # 获取当前账号。 
    regions = payload.regions or []  # 获取地区列表。 
    region_counts = payload.region_counts or {}  # 获取地区数量配置。 
    if not regions and not region_counts:  # 若地区列表为空则返回错误。 
        return JSONResponse({"ok": False, "message": "请至少选择一个地区"}, status_code=400)  # 返回错误响应。 
    ak = (payload.ak or "").strip()  # 获取并清理 AK 字符串。 
    sk = (payload.sk or "").strip()  # 获取并清理 SK 字符串。 
    if (ak and not sk) or (sk and not ak):  # 若 AK/SK 只填一个则返回错误。 
        return JSONResponse({"ok": False, "message": "AK 和 SK 需要同时填写"}, status_code=400)  # 返回错误响应。 
    if ak and sk and payload.manual_confirm_customer:  # 人工确认放行时仍先验证火山 AK/SK 有效。 
        try: 
            identity = customer_guard.get_account_identity_for_credentials(ak, sk) 
        except Exception as exc: 
            guard_message = _summarize_guard_exception("火山引擎", exc) 
            _append_log(f"step: customer_guard_manual_confirm_rejected {guard_message}", user) 
            _set_state(user, running=False, error=guard_message, finished_at=time.time()) 
            return JSONResponse({"ok": False, "message": guard_message}, status_code=403) 
        account_id = identity.get("account_id") or identity.get("identity_id") or "unknown" 
        _append_log(f"step: customer_guard_manual_confirm_ok account_id={account_id}", user) 
    elif ak and sk:  # 客户传入 AK/SK 时先确认属于代理后台客户。 
        try: 
            guard_ok, guard_message, _guard_detail = customer_guard.validate_customer_ownership(ak, sk) 
        except customer_guard.PartnerCustomerLookupError as exc: 
            guard_message = f"火山引擎代理后台客户查询失败，可人工确认后继续: {exc}" 
            return _manual_confirm_response(guard_message, user) 
        except Exception as exc: 
            guard_message = _summarize_guard_exception("火山引擎", exc) 
            _append_log(f"step: customer_guard_failed {guard_message}", user) 
            _set_state(user, running=False, error=guard_message, finished_at=time.time()) 
            return JSONResponse({"ok": False, "message": guard_message}, status_code=403) 
        if not guard_ok: 
            return _manual_confirm_response(guard_message, user) 
        _append_log(f"step: customer_guard_ok {guard_message}", user) 
    reserved, running_user = _reserve_task(user)  # 在启动后台线程前原子占用任务运行权。 
    if not reserved:  # 若已有任务运行则拒绝启动。 
        message = "本账号任务正在运行，请等待执行完毕" 
        return JSONResponse({"ok": False, "message": message, "running_user": running_user}, status_code=409)  # 返回冲突响应。 
    thread = threading.Thread(target=_run_task, args=(user, payload), daemon=True)  # 创建后台线程执行任务。 
    try:  # 防止线程启动失败后状态一直卡在运行中。 
        with STATE_LOCK:  # 先保存线程引用，避免启动后立即查询时误判。 
            TASK_THREAD_BY_USER[_state_key(user)] = thread 
        thread.start()  # 启动后台线程。 
    except Exception as exc:  # 捕获极少数线程启动异常。 
        with STATE_LOCK:  # 回滚线程引用。 
            TASK_THREAD_BY_USER.pop(_state_key(user), None) 
        _set_state(user, running=False, error=str(exc), finished_at=time.time())  # 回滚运行状态。 
        return JSONResponse({"ok": False, "message": f"任务启动失败: {exc}"}, status_code=500)  # 返回启动失败。 
    return {"ok": True, "message": "任务已启动"}  # 返回启动成功响应。 


@app.get("/api/status")  # 定义状态查询接口路由。 
def get_status(request: Request):  # 状态查询接口函数。 
    auth_error = _require_login(request)  # 校验登录状态。 
    if auth_error:  # 未登录则返回错误响应。 
        return auth_error  # 返回未登录响应。 
    user = _get_session_user(request)  # 获取当前账号。 
    with STATE_LOCK:  # 加锁读取状态信息。 
        _cleanup_stale_task_unlocked(user)  # 查询状态前先清理该账号失效占用。 
        return dict(_ensure_state_unlocked(user))  # 返回当前账号自己的状态信息。 


@app.get("/api/logs")  # 定义日志拉取接口路由。 
def get_logs(request: Request, cursor: int = 0):  # 日志拉取接口函数。 
    auth_error = _require_login(request)  # 校验登录状态。 
    if auth_error:  # 未登录则返回错误响应。 
        return auth_error  # 返回未登录响应。 
    user = _get_session_user(request)  # 获取当前账号。 
    buffer = _get_log_buffer(user)  # 获取当前账号的日志缓冲。 
    try:  # 捕获游标解析异常。 
        cursor_value = int(cursor)  # 转换游标为整数。 
    except (TypeError, ValueError):  # 若转换失败则重置。 
        cursor_value = 0  # 重置游标为 0。 
    if cursor_value < 0:  # 若游标为负数则修正。 
        cursor_value = 0  # 修正游标为 0。 
    with LOG_LOCK:  # 加锁读取日志缓冲区。 
        lines = buffer[cursor_value:]  # 取出指定游标后的日志行。 
        next_cursor = len(buffer)  # 计算下一个游标位置。 
    return {"lines": lines, "next_cursor": next_cursor}  # 返回日志数据与下一个游标。 


@app.post("/api/clear_logs")  # 定义清空日志接口路由。 
def clear_logs(request: Request):  # 清空日志接口函数。 
    auth_error = _require_login(request)  # 校验登录状态。 
    if auth_error:  # 未登录则返回错误响应。 
        return auth_error  # 返回未登录响应。 
    user = _get_session_user(request)  # 获取当前账号。 
    buffer = _get_log_buffer(user)  # 获取当前账号的日志缓冲。 
    with LOG_LOCK:  # 加锁清空日志缓冲。 
        buffer.clear()  # 清空日志列表。 
    return {"ok": True}  # 返回清空成功响应。 
