imageforge/backend/app/main.py

from __future__ import annotations

import logging
import time
from dataclasses import asdict
from datetime import datetime, timezone
from pathlib import Path

from fastapi import FastAPI, Header, HTTPException, Query, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, PlainTextResponse
from PIL import Image
import uvicorn

from .api.schemas import (
    AdminSettings,
    CancelResponse,
    DashboardStats,
    ExportRequest,
    ExportResponse,
    GenerateRequest,
    GenerateResponse,
    HealthResponse,
    HistoryItem,
    JobInfoResponse,
    MetricsResponse,
    ModelInfo,
    PresetPayload,
    PresetResponse,
    RetryResponse,
)
from .core.config import (
    ADMIN_TOKEN,
    CONTENT_PROFILE,
    CORS_ORIGINS,
    DEFAULT_BACKEND_HOST,
    DEFAULT_BACKEND_PORT,
    OUTPUT_DIR,
    OUTPUT_RETENTION_DAYS,
    REQUEST_MAX_BYTES,
)
from .core.logging import setup_logging
from .core.observability import MetricsStore
from .core.policy import ContentPolicy, PolicyAuditStore
from .core.security import ApiSecurity, Principal
from .jobs.manager import JobManager
from .providers.factory import ProviderRegistry
from .storage.history import PromptHistoryStore
from .storage.maintenance import cleanup_outputs
from .storage.presets import PresetStore
from .storage.settings import SettingsStore

setup_logging()
LOGGER = logging.getLogger(__name__)

provider_registry = ProviderRegistry()
history_store = PromptHistoryStore()
job_manager = JobManager(provider_registry, history_store)
policy_audit = PolicyAuditStore()
api_security = ApiSecurity()
preset_store = PresetStore()
settings_store = SettingsStore()
metrics = MetricsStore()

cleanup_removed = cleanup_outputs(OUTPUT_RETENTION_DAYS)
if cleanup_removed:
    LOGGER.info("Startup cleanup removed %s old output day folder(s)", cleanup_removed)

app = FastAPI(title="ImageForge Backend", version="0.3.0")
allow_origins = [item.strip() for item in CORS_ORIGINS.split(",")] if CORS_ORIGINS != "*" else ["*"]
app.add_middleware(
    CORSMiddleware,
    allow_origins=allow_origins,
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.middleware("http")
async def request_limits_middleware(request: Request, call_next):
    length = request.headers.get("content-length")
    if length and int(length) > REQUEST_MAX_BYTES:
        raise HTTPException(status_code=413, detail="Request too large")
    start = time.perf_counter()
    response = await call_next(request)
    elapsed_ms = (time.perf_counter() - start) * 1000.0
    metrics.incr("http_requests_total")
    metrics.observe_ms("http_request", elapsed_ms)
    response.headers["X-Content-Type-Options"] = "nosniff"
    response.headers["X-Frame-Options"] = "DENY"
    response.headers["Referrer-Policy"] = "no-referrer"
    response.headers["Content-Security-Policy"] = "default-src 'self'"
    return response


def _principal(
    http_request: Request,
    x_imageforge_api_key: str | None,
    minimum_role: str,
) -> Principal:
    api_security.limit = int(settings_store.get().get("rate_limit_per_minute", api_security.limit))
    principal = api_security.authenticate(
        api_key=x_imageforge_api_key,
        client_id=http_request.client.host if http_request.client else "unknown",
    )
    api_security.require_role(principal, minimum_role)
    return principal


@app.get("/health", response_model=HealthResponse)
def health() -> HealthResponse:
    return HealthResponse(status="ok", timestamp=datetime.now(timezone.utc))


@app.get("/ready", response_model=HealthResponse)
def ready() -> HealthResponse:
    try:
        OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
        probe = OUTPUT_DIR / ".ready_probe"
        probe.write_text("ok", encoding="utf-8")
        probe.unlink(missing_ok=True)
    except Exception as exc:  # noqa: BLE001
        raise HTTPException(status_code=503, detail=f"Output directory not writable: {exc}") from exc
    return HealthResponse(status="ready", timestamp=datetime.now(timezone.utc))


@app.get("/metrics", response_model=MetricsResponse)
def metrics_json(
    http_request: Request,
    x_imageforge_api_key: str | None = Header(default=None, alias="X-ImageForge-Api-Key"),
) -> MetricsResponse:
    _principal(http_request, x_imageforge_api_key, "viewer")
    return MetricsResponse(metrics=metrics.snapshot())


@app.get("/metrics/prometheus")
def metrics_prom(
    http_request: Request,
    x_imageforge_api_key: str | None = Header(default=None, alias="X-ImageForge-Api-Key"),
) -> PlainTextResponse:
    _principal(http_request, x_imageforge_api_key, "viewer")
    return PlainTextResponse(metrics.to_prometheus())


@app.get("/models", response_model=list[ModelInfo])
def models(
    http_request: Request,
    x_imageforge_api_key: str | None = Header(default=None, alias="X-ImageForge-Api-Key"),
) -> list[ModelInfo]:
    _principal(http_request, x_imageforge_api_key, "viewer")
    runtime = settings_store.get()
    adult_enabled = bool(runtime.get("adult_enabled", False))
    return [
        ModelInfo(
            id=provider.id,
            name=provider.name,
            description=provider.description,
            available=provider.is_available() if provider.id != "zimageturbo" else (adult_enabled and provider.is_available()),
        )
        for provider in provider_registry.list()
    ]


@app.post("/generate", response_model=GenerateResponse)
def generate(
    request: GenerateRequest,
    http_request: Request,
    x_imageforge_api_key: str | None = Header(default=None, alias="X-ImageForge-Api-Key"),
    x_imageforge_admin_token: str | None = Header(default=None, alias="X-ImageForge-Admin-Token"),
) -> GenerateResponse:
    _principal(http_request, x_imageforge_api_key, "operator")
    try:
        provider_registry.get(request.model)
    except KeyError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc

    runtime = settings_store.get()
    if request.model == "zimageturbo" and not bool(runtime.get("adult_enabled", False)):
        raise HTTPException(status_code=400, detail="Adult provider is disabled in admin settings")
    policy = ContentPolicy(runtime.get("content_profile", CONTENT_PROFILE))

    admin_override_applied = bool(
        request.admin_override and ADMIN_TOKEN and x_imageforge_admin_token == ADMIN_TOKEN
    )
    decision = policy.evaluate(
        f"{request.prompt}\n{request.negative_prompt}",
        admin_override=admin_override_applied,
    )
    policy_audit.write(
        prompt=request.prompt,
        negative_prompt=request.negative_prompt,
        profile=runtime.get("content_profile", CONTENT_PROFILE),
        decision=decision,
        client_ip=http_request.client.host if http_request.client else "unknown",
        model=request.model,
        admin_override_requested=request.admin_override,
        admin_override_applied=admin_override_applied,
    )
    if not decision.allowed:
        raise HTTPException(status_code=400, detail=f"Blocked by policy: {decision.reason}")

    metrics.incr("jobs_submitted_total")
    job_id = job_manager.submit(request)
    return GenerateResponse(job_id=job_id)


@app.get("/jobs/{job_id}", response_model=JobInfoResponse)
def job_status(
    job_id: str,
    http_request: Request,
    x_imageforge_api_key: str | None = Header(default=None, alias="X-ImageForge-Api-Key"),
) -> JobInfoResponse:
    _principal(http_request, x_imageforge_api_key, "viewer")
    state = job_manager.get(job_id)
    if state is None:
        raise HTTPException(status_code=404, detail="Job not found")
    return _to_job_response(state)


@app.get("/jobs", response_model=list[JobInfoResponse])
def list_jobs(
    http_request: Request,
    x_imageforge_api_key: str | None = Header(default=None, alias="X-ImageForge-Api-Key"),
) -> list[JobInfoResponse]:
    _principal(http_request, x_imageforge_api_key, "viewer")
    return [_to_job_response(state) for state in job_manager.list()]


@app.get("/image")
def get_image(path: str = Query(...)) -> FileResponse:
    try:
        source = Path(path).resolve(strict=True)
    except FileNotFoundError as exc:
        raise HTTPException(status_code=404, detail="Image not found") from exc

    output_root = OUTPUT_DIR.resolve()
    if source != output_root and output_root not in source.parents:
        raise HTTPException(status_code=403, detail="Forbidden source path")
    if not source.is_file():
        raise HTTPException(status_code=404, detail="Image not found")
    return FileResponse(source)


@app.post("/jobs/{job_id}/retry", response_model=RetryResponse)
def retry_job(
    job_id: str,
    http_request: Request,
    x_imageforge_api_key: str | None = Header(default=None, alias="X-ImageForge-Api-Key"),
) -> RetryResponse:
    _principal(http_request, x_imageforge_api_key, "operator")
    new_id = job_manager.retry(job_id)
    if not new_id:
        raise HTTPException(status_code=404, detail="Job not found")
    metrics.incr("jobs_retry_total")
    return RetryResponse(old_job_id=job_id, new_job_id=new_id)


def _to_job_response(state) -> JobInfoResponse:  # noqa: ANN001
    return JobInfoResponse(
        job_id=state.job_id,
        status=state.status,
        progress=state.progress,
        message=state.message,
        created_at=state.created_at,
        updated_at=state.updated_at,
        image_paths=state.image_paths,
        output_dir=state.output_dir,
        error=state.error,
    )


@app.post("/jobs/{job_id}/cancel", response_model=CancelResponse)
def cancel(
    job_id: str,
    http_request: Request,
    x_imageforge_api_key: str | None = Header(default=None, alias="X-ImageForge-Api-Key"),
) -> CancelResponse:
    _principal(http_request, x_imageforge_api_key, "operator")
    success = job_manager.cancel(job_id)
    if not success:
        raise HTTPException(status_code=404, detail="Job not found")
    metrics.incr("jobs_cancel_total")
    return CancelResponse(success=True)


@app.get("/history", response_model=list[HistoryItem])
def history(
    http_request: Request,
    x_imageforge_api_key: str | None = Header(default=None, alias="X-ImageForge-Api-Key"),
) -> list[HistoryItem]:
    _principal(http_request, x_imageforge_api_key, "viewer")
    items = history_store.list()
    parsed: list[HistoryItem] = []
    for item in items:
        try:
            parsed.append(
                HistoryItem(
                    prompt=item.get("prompt", ""),
                    negative_prompt=item.get("negative_prompt", ""),
                    timestamp=datetime.fromisoformat(item.get("timestamp", "")),
                )
            )
        except Exception:  # noqa: BLE001
            continue
    return parsed


@app.get("/dashboard/stats", response_model=DashboardStats)
def dashboard_stats(
    http_request: Request,
    x_imageforge_api_key: str | None = Header(default=None, alias="X-ImageForge-Api-Key"),
) -> DashboardStats:
    _principal(http_request, x_imageforge_api_key, "viewer")
    return DashboardStats(**job_manager.stats())


@app.get("/presets", response_model=list[PresetResponse])
def list_presets(
    http_request: Request,
    x_imageforge_api_key: str | None = Header(default=None, alias="X-ImageForge-Api-Key"),
) -> list[PresetResponse]:
    _principal(http_request, x_imageforge_api_key, "viewer")
    return [PresetResponse(**asdict(preset)) for preset in preset_store.list()]


@app.post("/presets", response_model=PresetResponse)
def upsert_preset(
    payload: PresetPayload,
    http_request: Request,
    x_imageforge_api_key: str | None = Header(default=None, alias="X-ImageForge-Api-Key"),
) -> PresetResponse:
    _principal(http_request, x_imageforge_api_key, "operator")
    preset = preset_store.upsert(payload.model_dump())
    return PresetResponse(**asdict(preset))


@app.delete("/presets/{name}", response_model=CancelResponse)
def delete_preset(
    name: str,
    http_request: Request,
    x_imageforge_api_key: str | None = Header(default=None, alias="X-ImageForge-Api-Key"),
) -> CancelResponse:
    _principal(http_request, x_imageforge_api_key, "admin")
    ok = preset_store.delete(name)
    if not ok:
        raise HTTPException(status_code=404, detail="Preset not found")
    return CancelResponse(success=True)


@app.post("/export", response_model=ExportResponse)
def export_image(
    payload: ExportRequest,
    http_request: Request,
    x_imageforge_api_key: str | None = Header(default=None, alias="X-ImageForge-Api-Key"),
) -> ExportResponse:
    _principal(http_request, x_imageforge_api_key, "operator")
    source = Path(payload.source_path).resolve()
    output_root = OUTPUT_DIR.resolve()
    if output_root not in source.parents:
        raise HTTPException(status_code=403, detail="Forbidden source path")
    if not source.exists() or not source.is_file():
        raise HTTPException(status_code=404, detail="Source image not found")
    img = Image.open(source).convert("RGB")
    if payload.max_width or payload.max_height:
        max_w = payload.max_width or img.width
        max_h = payload.max_height or img.height
        img.thumbnail((max_w, max_h))
    out_path = source.with_name(f"{source.stem}_export.{payload.format}")
    save_kwargs: dict[str, int] = {}
    if payload.format in {"jpg", "webp"}:
        save_kwargs["quality"] = payload.quality
    img.save(out_path, format=payload.format.upper(), **save_kwargs)
    return ExportResponse(output_path=str(out_path.resolve()))


@app.get("/admin/settings", response_model=AdminSettings)
def get_admin_settings(
    http_request: Request,
    x_imageforge_api_key: str | None = Header(default=None, alias="X-ImageForge-Api-Key"),
) -> AdminSettings:
    principal = _principal(http_request, x_imageforge_api_key, "admin")
    current = settings_store.get()
    metrics.incr(f"admin_settings_read_by_{principal.role}")
    return AdminSettings(**current)


@app.put("/admin/settings", response_model=AdminSettings)
def put_admin_settings(
    payload: AdminSettings,
    http_request: Request,
    x_imageforge_api_key: str | None = Header(default=None, alias="X-ImageForge-Api-Key"),
) -> AdminSettings:
    principal = _principal(http_request, x_imageforge_api_key, "admin")
    current = settings_store.update(payload.model_dump(), actor=principal.client_id)
    return AdminSettings(**current)


@app.post("/admin/cleanup", response_model=DashboardStats)
def cleanup_endpoint(
    http_request: Request,
    x_imageforge_api_key: str | None = Header(default=None, alias="X-ImageForge-Api-Key"),
) -> DashboardStats:
    _principal(http_request, x_imageforge_api_key, "admin")
    runtime = settings_store.get()
    cleanup_outputs(int(runtime.get("output_retention_days", OUTPUT_RETENTION_DAYS)))
    return DashboardStats(**job_manager.stats())


@app.get("/files/{relative_path:path}")
def output_file(
    relative_path: str,
    http_request: Request,
    x_imageforge_api_key: str | None = Header(default=None, alias="X-ImageForge-Api-Key"),
):
    _principal(http_request, x_imageforge_api_key, "viewer")
    path = (Path.cwd() / relative_path).resolve()
    output_root = OUTPUT_DIR.resolve()
    if output_root not in path.parents and path != output_root:
        raise HTTPException(status_code=403, detail="Forbidden")
    if not path.exists() or not path.is_file():
        raise HTTPException(status_code=404, detail="File not found")
    return FileResponse(path)


def run() -> None:
    uvicorn.run(
        "backend.app.main:app",
        host=DEFAULT_BACKEND_HOST,
        port=DEFAULT_BACKEND_PORT,
        reload=False,
        log_level="info",
    )


if __name__ == "__main__":
    LOGGER.info("Starting ImageForge backend")
    run()