server.py

from datetime import datetime
from concurrent.futures import ThreadPoolExecutor
import traceback
import threading
import uuid

from fastapi import FastAPI, Header, HTTPException
from fastapi.staticfiles import StaticFiles

import main as core

CONFIG = core.load_config()
core.apply_config(CONFIG)

API_KEY = CONFIG.get("api_key", "")
MAX_WORKERS = int(CONFIG.get("max_workers", 2))

app = FastAPI()
executor = ThreadPoolExecutor(max_workers=MAX_WORKERS)
tasks = {}
tasks_lock = threading.Lock()
print(f"[INFO] server startup: max_workers={MAX_WORKERS}")
app.mount("/screenshots", StaticFiles(directory="screenshots", check_dir=False), name="screenshots")


def require_api_key(x_api_key: str | None):
    if not API_KEY:
        return
    if not x_api_key or x_api_key != API_KEY:
        raise HTTPException(status_code=401, detail="unauthorized")


def update_task(task_id, **updates):
    with tasks_lock:
        task = tasks.get(task_id)
        if not task:
            return
        task.update(updates)
        task["updated_at"] = datetime.utcnow().isoformat()


def run_task(task_id):
    update_task(task_id, status="running")
    try:
        email, ok, cfg, screenshot = core.register_one_account()
        if ok:
            update_task(task_id, status="success", result=cfg)
        else:
            screenshot_url = f"/screenshots/{screenshot}" if screenshot else None
            update_task(task_id, status="failed", error=f"register failed: {email}", screenshot_url=screenshot_url)
    except Exception as exc:
        print(f"[ERR] task {task_id} failed: {exc}")
        print(traceback.format_exc())
        update_task(task_id, status="failed", error=str(exc))


@app.get("/health")
def health():
    with tasks_lock:
        total = len(tasks)
        running = sum(1 for t in tasks.values() if t["status"] == "running")
        pending = sum(1 for t in tasks.values() if t["status"] == "pending")
    return {
        "status": "ok",
        "total_tasks": total,
        "running": running,
        "pending": pending,
        "max_workers": MAX_WORKERS,
    }


@app.get("/")
def root():
    return {"status": "ok"}


@app.post("/task/create_account")
def create_account(x_api_key: str | None = Header(default=None)):
    require_api_key(x_api_key)
    task_id = uuid.uuid4().hex
    now = datetime.utcnow().isoformat()
    with tasks_lock:
        tasks[task_id] = {
            "status": "pending",
            "created_at": now,
            "updated_at": now,
            "result": None,
            "error": None,
        }
    executor.submit(run_task, task_id)
    return {"task_id": task_id}


@app.get("/task/status")
def task_status(task_id: str, x_api_key: str | None = Header(default=None)):
    require_api_key(x_api_key)
    with tasks_lock:
        task = tasks.get(task_id)
    if not task:
        raise HTTPException(status_code=404, detail="task not found")
    return {"task_id": task_id, "status": task["status"]}


@app.get("/task/info")
def task_info(task_id: str, x_api_key: str | None = Header(default=None)):
    require_api_key(x_api_key)
    with tasks_lock:
        task = tasks.get(task_id)
    if not task:
        raise HTTPException(status_code=404, detail="task not found")
    payload = {
        "task_id": task_id,
        "status": task["status"],
        "created_at": task["created_at"],
        "updated_at": task["updated_at"],
    }
    if task["status"] == "success":
        payload["result"] = task["result"]
    if task["status"] == "failed":
        payload["error"] = task["error"]
        payload["screenshot_url"] = task.get("screenshot_url")
    return payload