MpakoImageAIbot/web.py

from __future__ import annotations

import base64
import math
import secrets
from contextlib import asynccontextmanager
from datetime import datetime
from pathlib import Path
from typing import Any
from urllib.parse import urlencode

from fastapi import Depends, FastAPI, HTTPException, Request
from fastapi.responses import FileResponse, HTMLResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates

from config import Settings
from db import Database
from db import Request as HistoryRequest

BASE_DIR = Path(__file__).resolve().parent
TEMPLATES_DIR = BASE_DIR / "web" / "templates"
STATIC_DIR = BASE_DIR / "web" / "static"

STATUSES = ("queued", "running", "done", "failed")
MODES = ("t2i", "i2i")


def fmt_dt(value: str | None) -> str:
    if not value:
        return "—"
    try:
        parsed = datetime.fromisoformat(value)
        return parsed.astimezone().strftime("%d.%m.%Y %H:%M")
    except ValueError:
        return value


def make_auth(settings: Settings):
    async def require_auth(request: Request) -> None:
        if not settings.web_password:
            return
        header = request.headers.get("Authorization", "")
        if not header.startswith("Basic "):
            raise HTTPException(
                status_code=401,
                detail="Требуется авторизация",
                headers={"WWW-Authenticate": 'Basic realm="MpakoImageAIbot"'},
            )
        try:
            decoded = base64.b64decode(header[6:]).decode("utf-8")
            username, _, password = decoded.partition(":")
        except (ValueError, UnicodeDecodeError):
            username = password = ""
        valid_user = secrets.compare_digest(username, settings.web_username)
        valid_password = secrets.compare_digest(password, settings.web_password)
        if not (valid_user and valid_password):
            raise HTTPException(
                status_code=401,
                detail="Требуется авторизация",
                headers={"WWW-Authenticate": 'Basic realm="MpakoImageAIbot"'},
            )

    return require_auth


def create_app(settings: Settings | None = None) -> FastAPI:
    settings = settings or Settings.from_env(require_token=False)
    templates = Jinja2Templates(directory=str(TEMPLATES_DIR))
    templates.env.filters["fmt_dt"] = fmt_dt

    @asynccontextmanager
    async def lifespan(app: FastAPI):
        db = Database(settings.db_path)
        await db.connect()
        app.state.db = db
        yield
        await db.close()

    app = FastAPI(
        title="MpakoImageAIbot history",
        docs_url=None,
        redoc_url=None,
        dependencies=[Depends(make_auth(settings))],
        lifespan=lifespan,
    )
    app.state.settings = settings
    app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static")

    @app.get("/healthz")
    async def healthz() -> dict[str, bool]:
        return {"ok": True}

    @app.get("/", response_class=HTMLResponse)
    async def index(
        request: Request,
        user_id: str = "",
        status: str = "",
        mode: str = "",
        page: int = 1,
    ):
        db: Database = request.app.state.db
        settings: Settings = request.app.state.settings

        error = ""
        uid: int | None = None
        raw_user_id = user_id.strip()
        if raw_user_id:
            if raw_user_id.isdigit():
                uid = int(raw_user_id)
            else:
                error = "user_id должен состоять только из цифр."

        status = status if status in STATUSES else ""
        mode = mode if mode in MODES else ""

        total = await db.count_requests(user_id=uid, status=status or None, mode=mode or None)
        page_size = settings.web_page_size
        pages = max(1, math.ceil(total / page_size))
        page = max(1, min(page, pages))
        offset = (page - 1) * page_size

        rows = await db.list_requests(
            user_id=uid,
            status=status or None,
            mode=mode or None,
            limit=page_size,
            offset=offset,
        )
        users = await db.list_users()
        stats = await db.stats()

        base_filters = {
            key: value
            for key, value in (
                ("user_id", raw_user_id),
                ("status", status),
                ("mode", mode),
            )
            if value
        }

        def page_url(number: int) -> str:
            params = dict(base_filters)
            if number > 1:
                params["page"] = str(number)
            query = urlencode(params)
            return f"/?{query}" if query else "/"

        context: dict[str, Any] = {
            "rows": rows,
            "users": users,
            "stats": stats,
            "statuses": STATUSES,
            "modes": MODES,
            "user_id": raw_user_id,
            "status": status,
            "mode": mode,
            "error": error,
            "page": page,
            "pages": pages,
            "total": total,
            "prev_url": page_url(page - 1) if page > 1 else None,
            "next_url": page_url(page + 1) if page < pages else None,
        }
        return templates.TemplateResponse(
            request=request, name="index.html", context=context
        )

    @app.get("/request/{request_id}", response_class=HTMLResponse)
    async def detail(request: Request, request_id: int):
        db: Database = request.app.state.db
        item = await db.get_request(request_id)
        if item is None:
            raise HTTPException(status_code=404, detail="Запрос не найден")
        return templates.TemplateResponse(
            request=request,
            name="detail.html",
            context={"item": item},
        )

    @app.get("/media/{request_id}/{kind}")
    async def media(request: Request, request_id: int, kind: str):
        if kind not in {"input", "output"}:
            raise HTTPException(status_code=404, detail="Не найдено")
        db: Database = request.app.state.db
        settings: Settings = request.app.state.settings
        item: HistoryRequest | None = await db.get_request(request_id)
        if item is None:
            raise HTTPException(status_code=404, detail="Запрос не найден")

        raw_path = item.input_file_path if kind == "input" else item.output_file_path
        if not raw_path:
            raise HTTPException(status_code=404, detail="Файл не найден")

        storage_root = settings.storage_dir.resolve()
        path = Path(raw_path).resolve()
        try:
            path.relative_to(storage_root)
        except ValueError:
            raise HTTPException(status_code=404, detail="Файл не найден") from None
        if not path.is_file():
            raise HTTPException(status_code=404, detail="Файл не найден")
        return FileResponse(path)

    return app


app = create_app()