main.py

import httpx
import os
from fastapi import FastAPI, Request, WebSocket, WebSocketDisconnect
from fastapi.responses import RedirectResponse
from starlette.middleware.base import BaseHTTPMiddleware
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from contextlib import asynccontextmanager

from server.queue import GenerationQueue
from server.nai_api import generate_image, encode_vibes, augment_image
from server.tagger import tagger_predictor
from server.cloud_storage import cloud_storage
from server.auth import (
        is_authenticated, verify_credentials, create_session,
                login_page_html, auth_enabled, SESSION_COOKIE, SESSION_MAX_AGE,
                )


gen_queue = GenerationQueue()
tagger_queue = GenerationQueue()

@asynccontextmanager
async def lifespan(app: FastAPI):
    gen_queue.start(handle_task)
    tagger_queue.start(handle_tagger_task)
    yield    

app = FastAPI(title="NAI Studio", lifespan=lifespan)

app.mount("/static", StaticFiles(directory="static"), name="static")
app.mount("/data", StaticFiles(directory="data"), name="data")
templates = Jinja2Templates(directory="templates")


# ── Auth Middleware ──────────────────────────────
class AuthMiddleware(BaseHTTPMiddleware):
    # Paths that don't require authentication
    PUBLIC_PATHS = {"/login", "/static/", "/favicon.ico"}

    async def dispatch(self, request: Request, call_next):
        if not auth_enabled():
            return await call_next(request)

        path = request.url.path

        # Allow public paths
        for pp in self.PUBLIC_PATHS:
            if path == pp or path.startswith(pp):
                return await call_next(request)

        # Check auth
        if not is_authenticated(request):
            if path == "/" or path.startswith("/api/") or path.startswith("/data/"):
                return RedirectResponse(url="/login", status_code=302)
            # For WebSocket, just reject
            return RedirectResponse(url="/login", status_code=302)

        return await call_next(request)

app.add_middleware(AuthMiddleware)


# ── Queue task handler ──────────────────────────
async def handle_task(task: dict, queue: GenerationQueue):
    client_id = task["client_id"]
    task_type = task["type"]

    try:
        if task_type == "generate":
            await queue.notify_client(client_id, {
                "type": "status", "status": "generating",
                "message": "正在生成图片...",
            })
            result = await generate_image(task["params"])
            queue.last_seeds[client_id] = result.get("seed", 0)

            # Cloud save if requested
            cloud_path = None
            if task["params"].get("save_to_cloud", False) and result.get("image"):
                try:
                    import asyncio
                    loop = asyncio.get_event_loop()
                    seed = result.get("seed", 0)
                    cloud_path = await loop.run_in_executor(
                        None, cloud_storage.save_image, result["image"], seed
                    )
                    # Store search text and upload in background
                    if cloud_path:
                        # Build search text from params
                        p = task["params"]
                        search_parts = []
                        for k in ["prompt", "quality_prompt", "negative_prompt", "model", "sampler"]:
                            if p.get(k):
                                search_parts.append(str(p[k]))
                        if result.get("seed"):
                            search_parts.append(str(result["seed"]))
                        if p.get("steps"):
                            search_parts.append(str(p["steps"]))
                        # Character prompts
                        char_data = p.get("character", {})
                        for ch in char_data.get("characters", []):
                            if ch.get("prompt"):
                                search_parts.append(ch["prompt"])
                        search_text = " ".join(search_parts)
                        cloud_storage.set_search_text(cloud_path, search_text)
                        loop.run_in_executor(None, cloud_storage.upload_file, cloud_path)
                except Exception as e:
                    print(f"[Cloud] Save failed: {e}")

            await queue.notify_client(client_id, {
                "type": "result", "status": "completed",
                "image": result.get("image"),
                "info": result.get("info"),
                "seed": result.get("seed"),
                "cloud_path": cloud_path,
            })

        elif task_type == "augment":
            await queue.notify_client(client_id, {
                "type": "status", "status": "augmenting",
                "message": "正在处理图片...",
            })
            result = await augment_image(task["params"])
            await queue.notify_client(client_id, {
                "type": "augment_result", "status": "completed",
                "images": result["images"],
            })

        elif task_type == "vibe_encode":
            await queue.notify_client(client_id, {
                "type": "status", "status": "encoding",
                "message": "正在编码Vibe参考图片...",
            })
            result = await encode_vibes(task["params"])
            await queue.notify_client(client_id, {
                "type": "vibe_encode_result", "status": "completed",
                "model": task["params"]["model"],
                "results": result,
            })

    except Exception as e:
        if task_type == "generate":
            error_type = "error"
        elif task_type == "augment":
            error_type = "augment_error"
        else:
            error_type = "vibe_encode_error"
        await queue.notify_client(client_id, {
            "type": error_type, "status": "error",
            "message": str(e),
        })


# ── Tagger queue task handler ───────────────────
async def handle_tagger_task(task: dict, queue: GenerationQueue):
    client_id = task["client_id"]
    try:
        await queue.notify_client(client_id, {
            "type": "status", "status": "tagging",
            "message": "正在加载模型并推理...",
        })
        params = task["params"]
        import asyncio, functools
        loop = asyncio.get_event_loop()
        result = await loop.run_in_executor(
            None,
            functools.partial(
                tagger_predictor.predict,
                params["image"],
                params["model_repo"],
                params.get("general_thresh", 0.35),
                params.get("general_mcut", False),
                params.get("character_thresh", 0.85),
                params.get("character_mcut", False),
            )
        )
        await queue.notify_client(client_id, {
            "type": "tagger_result", "status": "completed",
            **result,
        })
    except Exception as e:
        await queue.notify_client(client_id, {
            "type": "tagger_error", "status": "error",
            "message": str(e),
        })


# ── Routes ──────────────────────────────────────
# ── Auth Routes ─────────────────────────────────
@app.get("/login", response_class=HTMLResponse)
async def login_page(request: Request):
    if is_authenticated(request):
        return RedirectResponse(url="/", status_code=302)
    return HTMLResponse(login_page_html())


@app.post("/login")
async def login_submit(request: Request):
    form = await request.form()
    username = form.get("username", "")
    password = form.get("password", "")

    if verify_credentials(username, password):
        token = create_session()
        response = RedirectResponse(url="/", status_code=302)
        response.set_cookie(
            SESSION_COOKIE, token,
            max_age=SESSION_MAX_AGE,
            httponly=True,
            samesite="lax",
        )
        return response
    else:
        return HTMLResponse(login_page_html("用户名或密码错误"), status_code=401)


@app.get("/logout")
async def logout():
    response = RedirectResponse(url="/login", status_code=302)
    response.delete_cookie(SESSION_COOKIE)
    return response


@app.get("/", response_class=HTMLResponse)
async def index(request: Request):
    return templates.TemplateResponse("index.html", {"request": request})


@app.get("/api/cloud/dates")
async def cloud_dates():
    try:
        dates = cloud_storage.list_dates()
        return {"dates": dates}
    except Exception as e:
        return {"error": str(e)}


@app.get("/api/cloud/images")
async def cloud_images(date: str = ""):
    try:
        images = cloud_storage.list_images(date)
        return {"images": images}
    except Exception as e:
        return {"error": str(e)}


@app.get("/api/cloud/image")
async def cloud_image(path: str):
    try:
        b64 = cloud_storage.get_image_b64(path)
        if b64 is None:
            return {"error": "Not found"}
        return {"image": b64}
    except Exception as e:
        return {"error": str(e)}


@app.post("/api/cloud/delete")
async def cloud_delete(request: Request):
    data = await request.json()
    path = data.get("path", "")
    try:
        import asyncio
        loop = asyncio.get_event_loop()
        await loop.run_in_executor(None, cloud_storage.delete_file, path)
        return {"ok": True}
    except Exception as e:
        return {"error": str(e)}


@app.post("/api/cloud/sync")
async def cloud_sync():
    try:
        import asyncio
        loop = asyncio.get_event_loop()
        await loop.run_in_executor(None, cloud_storage.sync_download)
        return {"ok": True}
    except Exception as e:
        return {"error": str(e)}


@app.get("/api/cloud/search")
async def cloud_search(q: str = "", date: str = ""):
    try:
        images = cloud_storage.search_images(q, date)
        return {"images": images}
    except Exception as e:
        return {"error": str(e)}


@app.post("/api/cloud/rebuild-index")
async def cloud_rebuild_index():
    try:
        import asyncio
        loop = asyncio.get_event_loop()
        await loop.run_in_executor(None, cloud_storage.build_search_index_from_files)
        return {"ok": True}
    except Exception as e:
        return {"error": str(e)}


@app.get("/api/anlas")
async def get_anlas():
    api_key = os.environ.get("NAI_API_KEY", "")
    if not api_key:
        return {"error": "NAI_API_KEY not set"}
    try:
        headers = {
            "Authorization": f"Bearer {api_key}",
        }
        async with httpx.AsyncClient(timeout=10.0) as client:
            resp = await client.get(
                "https://image.novelai.net/user/data",
                headers=headers,
            )
        data = resp.json()
        steps = data["subscription"]["trainingStepsLeft"]
        anlas = steps["fixedTrainingStepsLeft"] + steps["purchasedTrainingSteps"]
        return {"anlas": anlas}
    except Exception as e:
        return {"error": f"获取失败: {str(e)}"}


@app.websocket("/ws/{client_id}")
async def websocket_endpoint(websocket: WebSocket, client_id: str):
    # Check auth for WebSocket
    if auth_enabled():
        token = websocket.cookies.get(SESSION_COOKIE, "")
        from server.auth import validate_session
        if not validate_session(token):
            await websocket.close(code=4001, reason="Unauthorized")
            return
    await websocket.accept()
    gen_queue.clients[client_id] = websocket
    tagger_queue.clients[client_id] = websocket
    try:
        while True:
            data = await websocket.receive_json()
            action = data.get("action")

            if action == "generate":
                params = data.get("params", {})
                params["vibe"] = data.get("vibe", {})
                params["extra_input"] = data.get("extra_input")
                params["character"] = data.get("character", {})
                params["reference"] = data.get("reference")
                params["advanced"] = data.get("advanced", {})
                task_id, queue_size = await gen_queue.add_task(
                    client_id, "generate", params
                )
                await websocket.send_json({
                    "type": "queued", "task_id": task_id,
                    "position": queue_size,
                    "message": f"已加入队列，当前位置: {queue_size}",
                })

            elif action == "vibe_encode":
                encode_params = {
                    "model": data.get("model", ""),
                    "images": data.get("images", []),
                }
                task_id, queue_size = await gen_queue.add_task(
                    client_id, "vibe_encode", encode_params
                )
                await websocket.send_json({
                    "type": "queued", "task_id": task_id,
                    "position": queue_size,
                    "message": f"Vibe编码已加入队列，当前位置: {queue_size}",
                })

            elif action == "tagger":
                tagger_params = data.get("params", {})
                task_id, queue_size = await tagger_queue.add_task(
                    client_id, "tagger", tagger_params
                )
                await websocket.send_json({
                    "type": "queued", "task_id": task_id,
                    "position": queue_size,
                    "message": f"Tagger推理已加入队列，当前位置: {queue_size}",
                })

            elif action == "augment":
                aug_params = data.get("params", {})
                task_id, queue_size = await gen_queue.add_task(
                    client_id, "augment", aug_params
                )
                await websocket.send_json({
                    "type": "queued", "task_id": task_id,
                    "position": queue_size,
                    "message": f"定向修图已加入队列，当前位置: {queue_size}",
                })

            elif action == "get_last_seed":
                seed = gen_queue.get_last_seed(client_id)
                await websocket.send_json({"type": "last_seed", "seed": seed})

    except WebSocketDisconnect:
        gen_queue.clients.pop(client_id, None)
        tagger_queue.clients.pop(client_id, None)


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)