backend/app.py

from __future__ import annotations

import base64
import os
import platform
import tempfile
from pathlib import Path
from typing import Any, Dict

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from fastapi.staticfiles import StaticFiles

from .config import AppConfig, load_config, save_config
from .models import ConfigUpdate, I2IRequest, InpaintRequest, T2IRequest
from .services.novelai import generate_i2i, generate_inpaint, generate_t2i


def _ensure_x64(n: int) -> int:
    if n <= 64:
        return 64
    if n % 64 == 0:
        return n
    return ((n // 64) + 1) * 64 if (n / 64) % 1 >= 0.5 else (n // 64) * 64


def _as_data_uri(b64: str) -> str:
    return f"data:image/png;base64,{b64}"


app = FastAPI(title="New NAI", version="1.0.0")

# CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=False,
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.get("/api/health")
def health() -> Dict[str, Any]:
    return {"status": "ok"}


@app.get("/api/config")
def get_config() -> Dict[str, Any]:
    cfg: AppConfig = load_config()
    return cfg.model_dump()


@app.put("/api/config")
def update_config(update: ConfigUpdate) -> Dict[str, Any]:
    cfg = load_config()
    data = update.model_dump(exclude_none=True)
    for k, v in data.items():
        setattr(cfg, k, v)
    save_config(cfg)
    return cfg.model_dump()


@app.post("/api/open-dir")
def api_open_dir(payload: Dict[str, Any]) -> Dict[str, Any]:
    """
    打开指定目录；若未传 path，则打开配置中的 output_dir。
    Windows 使用 os.startfile，macOS 用 open，Linux 用 xdg-open。
    """
    path = (payload or {}).get("path") or ""
    if not path:
        cfg = load_config()
        path = cfg.output_dir
    if not path:
        raise HTTPException(status_code=400, detail="未提供路径且配置中未设置 output_dir")

    p = Path(path)
    try:
        p.mkdir(parents=True, exist_ok=True)
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"创建目录失败: {e}") from e

    try:
        if hasattr(os, "startfile"):
            os.startfile(str(p))  # Windows
        elif platform.system() == "Darwin":
            import subprocess
            subprocess.run(["open", str(p)], check=False)
        else:
            import subprocess
            subprocess.run(["xdg-open", str(p)], check=False)
        return {"ok": True, "path": str(p)}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"无法打开目录: {e}") from e


@app.post("/api/generate/t2i")
def api_t2i(req: T2IRequest):
    cfg = load_config()
    if not cfg.key:
        raise HTTPException(status_code=400, detail="尚未配置 key，请先在配置中设置 key。")

    width = _ensure_x64(req.width or 768)
    height = _ensure_x64(req.height or 768)

    try:
        b64, saved = generate_t2i(
            cfg,
            prompt=req.prompt,
            negative=req.negative or "",
            width=width,
            height=height,
            scale=req.scale,
            steps=req.steps,
            sampler=req.sampler,
            noise_schedule=req.noise_schedule,
            seed=req.seed,
            variety=req.variety,
            decrisp=req.decrisp,
            cfg_rescale=req.cfg_rescale,
        )
        return {"image_base64": _as_data_uri(b64), "saved_path": saved}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e)) from e


@app.post("/api/generate/i2i")
def api_i2i(req: I2IRequest):
    cfg = load_config()
    if not cfg.key:
        raise HTTPException(status_code=400, detail="尚未配置 key，请先在配置中设置 key。")

    width = _ensure_x64(req.width or 768)
    height = _ensure_x64(req.height or 768)

    try:
        b64, saved = generate_i2i(
            cfg,
            positive=req.positive,
            negative=req.negative or "",
            image_base64=req.image_base64,
            width=width,
            height=height,
            scale=req.scale,
            steps=req.steps,
            sampler=req.sampler,
            noise_schedule=req.noise_schedule,
            strength=req.strength or 0.5,
            noise=req.noise or 0.0,
            seed=req.seed,
            variety=req.variety,
            decrisp=req.decrisp,
            cfg_rescale=req.cfg_rescale,
        )
        return {"image_base64": _as_data_uri(b64), "saved_path": saved}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e)) from e


@app.post("/api/generate/inpaint")
def api_inpaint(req: InpaintRequest):
    cfg = load_config()
    if not cfg.key:
        raise HTTPException(status_code=400, detail="尚未配置 key，请先在配置中设置 key。")

    width = _ensure_x64(req.width or 768)
    height = _ensure_x64(req.height or 768)

    try:
        b64, saved = generate_inpaint(
            cfg,
            positive=req.positive,
            negative=req.negative or "",
            image_base64=req.image_base64,
            mask_base64=req.mask_base64,
            add_original_image=req.add_original_image,
            width=width,
            height=height,
            scale=req.scale,
            steps=req.steps,
            sampler=req.sampler,
            noise_schedule=req.noise_schedule,
            strength=req.strength or 0.5,
            noise=req.noise or 0.0,
            seed=req.seed,
            variety=req.variety,
            decrisp=req.decrisp,
            cfg_rescale=req.cfg_rescale,
        )
        return {"image_base64": _as_data_uri(b64), "saved_path": saved}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e)) from e


# 前端静态资源（仅保留必要 UI，无教程/仓库链接）
_frontend_dir = Path(__file__).resolve().parent.parent / "frontend"
# 提示音静态资源映射：/ring -> 项目根/ring 目录（例如 G:\NOVELAI\New NAI\ring）
_ring_dir = Path(__file__).resolve().parent.parent / "ring"
app.mount("/ring", StaticFiles(directory=_ring_dir, html=False), name="ring")
app.mount("/", StaticFiles(directory=_frontend_dir, html=True), name="frontend")