app.py

from __future__ import annotations

import io
import logging
import os
from urllib.parse import urlparse

import httpx

from fastapi import FastAPI, File, Form, HTTPException, Request, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, RedirectResponse, Response, StreamingResponse
from pydantic import BaseModel, Field
from starlette.background import BackgroundTask
from starlette.concurrency import run_in_threadpool

from PIL import Image
from features.doubao_watermark.service import run_remove_doubao
from services.douyin.service import DouyinParseError, resolve_douyin_share
from services.inpainting import service as inpainting_service
from services.ocr.routes import ocr_router, ocr_service
from services.youtube.service import (
    YouTubeTranscriptError,
    fetch_signed_caption,
    fetch_youtube_transcript,
    fetch_youtube_video,
)


logger = logging.getLogger(__name__)

DEPLOYMENT_TARGET = os.getenv("DEPLOYMENT_TARGET", "local").strip().lower()
SUB2API_ORIGIN_URL = os.getenv("SUB2API_ORIGIN_URL", "http://101.43.24.194").rstrip("/")
SUB2API_ORIGIN_HOST = os.getenv("SUB2API_ORIGIN_HOST", "101.43.24.194")
SUB2API_ORIGIN_TOKEN = os.getenv("SUB2API_ORIGIN_TOKEN", "")
XTOKEN_IMAGES_BASE_URL = os.getenv(
    "XTOKEN_IMAGES_BASE_URL", "https://api.xtokenmirror.com/v1/images"
).rstrip("/")
XTOKEN_IMAGES_MAX_BODY_BYTES = 32 * 1024 * 1024
OCR_WARMUP_ON_STARTUP = os.getenv("OCR_WARMUP_ON_STARTUP", "0") == "1"
HOP_BY_HOP_HEADERS = {
    "connection",
    "keep-alive",
    "proxy-authenticate",
    "proxy-authorization",
    "te",
    "trailer",
    "transfer-encoding",
    "upgrade",
}


class DouyinParseRequest(BaseModel):
    share_text: str = Field(min_length=1, max_length=4096)


class YouTubeTranscriptRequest(BaseModel):
    url: str = Field(min_length=1, max_length=2048)
    languages: list[str] | None = Field(default=None, max_length=10)


class YouTubeInfoRequest(YouTubeTranscriptRequest):
    include_transcript: bool = True


class YouTubeCaptionRequest(BaseModel):
    url: str = Field(min_length=1, max_length=8192)

app = FastAPI(title="Image Services API")
app.include_router(ocr_router)
app.add_middleware(
    CORSMiddleware,
    allow_origins=[
        "https://image.goudaner.fun",
        "https://watermark-blush.vercel.app",
        "https://watermark-clipdalles-projects.vercel.app",
        "https://watermark-clipdalle-clipdalles-projects.vercel.app",
    ],
    allow_origin_regex=r"^https?://(localhost|127\.0\.0\.1)(:\d+)?$",
    allow_methods=["GET", "POST", "OPTIONS"],
    allow_headers=["*"],
    expose_headers=[
        "X-Inpaint-Mode",
        "X-Inpaint-Backend",
        "X-Inpaint-Elapsed",
        "X-Doubao-Confidence",
        "X-Doubao-Coverage",
        "X-Doubao-Detection",
        "X-Doubao-BBox",
    ],
)


async def _close_proxy_response(response: httpx.Response, client: httpx.AsyncClient) -> None:
    try:
        await response.aclose()
    finally:
        await client.aclose()


async def _proxy_sub2api(request: Request, upstream_path: str) -> StreamingResponse:
    if not SUB2API_ORIGIN_TOKEN:
        raise HTTPException(status_code=503, detail="Sub2API origin token is not configured")

    target_url = httpx.URL(f"{SUB2API_ORIGIN_URL}{upstream_path}")
    if request.url.query:
        target_url = target_url.copy_with(query=request.url.query.encode("utf-8"))
    request_headers = {
        name: value
        for name, value in request.headers.items()
        if name.lower() not in HOP_BY_HOP_HEADERS | {"host", "content-length"}
    }
    request_headers["host"] = SUB2API_ORIGIN_HOST
    request_headers["x-sub2api-origin-token"] = SUB2API_ORIGIN_TOKEN

    timeout = httpx.Timeout(connect=20.0, read=None, write=120.0, pool=20.0)
    client = httpx.AsyncClient(timeout=timeout, follow_redirects=False)
    upstream_request = client.build_request(
        request.method,
        target_url,
        headers=request_headers,
        content=request.stream(),
    )

    try:
        upstream_response = await client.send(upstream_request, stream=True)
    except httpx.RequestError as exc:
        await client.aclose()
        logger.warning("Sub2API upstream request failed: %s", exc)
        raise HTTPException(
            status_code=502,
            detail=f"Sub2API upstream unavailable: {exc}",
        ) from exc

    response_headers = {
        name: value
        for name, value in upstream_response.headers.items()
        if name.lower() not in HOP_BY_HOP_HEADERS
    }
    return StreamingResponse(
        upstream_response.aiter_raw(),
        status_code=upstream_response.status_code,
        headers=response_headers,
        background=BackgroundTask(_close_proxy_response, upstream_response, client),
    )


@app.api_route(
    "/v1/{path:path}",
    methods=["GET", "POST", "PUT", "PATCH", "DELETE", "OPTIONS", "HEAD"],
)
async def proxy_sub2api(request: Request, path: str) -> StreamingResponse:
    return await _proxy_sub2api(request, f"/v1/{path}")


async def _proxy_xtoken_images(request: Request, operation: str) -> StreamingResponse:
    if operation not in {"generations", "edits"}:
        raise HTTPException(status_code=404, detail="Images API operation not found")

    authorization = request.headers.get("authorization", "")
    if not authorization.lower().startswith("bearer "):
        raise HTTPException(status_code=401, detail="XToken API key is required")

    content_length = request.headers.get("content-length")
    if content_length:
        try:
            if int(content_length) > XTOKEN_IMAGES_MAX_BODY_BYTES:
                raise HTTPException(status_code=413, detail="Image request exceeds 32 MiB")
        except ValueError as exc:
            raise HTTPException(status_code=400, detail="Invalid Content-Length header") from exc

    body = await request.body()
    if len(body) > XTOKEN_IMAGES_MAX_BODY_BYTES:
        raise HTTPException(status_code=413, detail="Image request exceeds 32 MiB")

    request_headers = {"authorization": authorization}
    for name in ("content-type", "accept", "user-agent"):
        value = request.headers.get(name)
        if value:
            request_headers[name] = value

    client = httpx.AsyncClient(
        timeout=httpx.Timeout(connect=20.0, read=None, write=180.0, pool=20.0),
        follow_redirects=False,
    )
    upstream_request = client.build_request(
        "POST",
        f"{XTOKEN_IMAGES_BASE_URL}/{operation}",
        headers=request_headers,
        content=body,
    )
    try:
        upstream_response = await client.send(upstream_request, stream=True)
    except httpx.RequestError as exc:
        await client.aclose()
        logger.warning("XToken Images upstream request failed: %s", exc)
        raise HTTPException(
            status_code=502,
            detail=f"XToken Images API unavailable: {exc}",
        ) from exc

    response_headers = {
        name: value
        for name, value in upstream_response.headers.items()
        if name.lower() not in HOP_BY_HOP_HEADERS
        and not name.lower().startswith("access-control-")
    }
    return StreamingResponse(
        upstream_response.aiter_raw(),
        status_code=upstream_response.status_code,
        headers=response_headers,
        background=BackgroundTask(_close_proxy_response, upstream_response, client),
    )


@app.post("/xtoken-images/{operation}")
async def proxy_xtoken_images(request: Request, operation: str) -> StreamingResponse:
    return await _proxy_xtoken_images(request, operation)


@app.on_event("startup")
def startup() -> None:
    expected_sources = {
        "huggingface": ("huggingface", "url"),
        "modelscope": ("modelscope", "modelscope"),
        "local": (inpainting_service.LAMA_MODEL_SOURCE, inpainting_service.MIGAN_MODEL_SOURCE),
    }
    if DEPLOYMENT_TARGET not in expected_sources:
        raise RuntimeError(f"Unsupported DEPLOYMENT_TARGET: {DEPLOYMENT_TARGET}")
    expected_lama, expected_migan = expected_sources[DEPLOYMENT_TARGET]
    if (inpainting_service.LAMA_MODEL_SOURCE, inpainting_service.MIGAN_MODEL_SOURCE) != (expected_lama, expected_migan):
        raise RuntimeError(
            "Deployment/model source mismatch: "
            f"target={DEPLOYMENT_TARGET}, lama={inpainting_service.LAMA_MODEL_SOURCE}, "
            f"migan={inpainting_service.MIGAN_MODEL_SOURCE}"
        )

    if inpainting_service.LOAD_MODEL_ON_STARTUP:
        for mode in inpainting_service.PRELOAD_WATERMARK_MODES:
            if inpainting_service._normalize_mode(mode) == "fast":
                inpainting_service._load_migan_model()
            else:
                inpainting_service._load_lama_session()

    if OCR_WARMUP_ON_STARTUP:
        ocr_service.registry.get("rapidocr").warmup()


@app.get("/")
def root() -> dict[str, object]:
    return health()


@app.get("/health")
def health() -> dict[str, object]:
    return {
        "status": "ok",
        "mode": "dual-selected-backends",
        "deployment": {
            "target": DEPLOYMENT_TARGET,
            "model_sources": {
                "lama": inpainting_service.LAMA_MODEL_SOURCE,
                "migan": inpainting_service.MIGAN_MODEL_SOURCE,
            },
        },
        "default_mode": inpainting_service.DEFAULT_MODE,
        "supported_modes": inpainting_service.MODES,
        "loaded": {
            "quality": inpainting_service.LAMA_SESSION is not None,
            "fast": inpainting_service.MIGAN_MODEL is not None,
        },
    }


@app.get("/douyin/health")
def douyin_health() -> dict[str, object]:
    return {"ok": True, "service": "douyin-parser"}


@app.post("/douyin/parse")
async def parse_douyin(request: DouyinParseRequest) -> JSONResponse:
    try:
        result = await resolve_douyin_share(request.share_text)
    except DouyinParseError as exc:
        return JSONResponse(
            status_code=400,
            content={
                "ok": False,
                "result": None,
                "error": {"code": exc.code, "message": exc.message},
            },
        )
    return JSONResponse(
        content={"ok": True, "result": result.to_dict(), "error": None}
    )


@app.get("/youtube/health")
def youtube_health() -> dict[str, object]:
    return {"ok": True, "service": "youtube-transcript"}


@app.post("/youtube/transcript")
async def youtube_transcript(request: YouTubeTranscriptRequest) -> JSONResponse:
    try:
        result = await run_in_threadpool(
            fetch_youtube_transcript,
            request.url,
            request.languages,
        )
    except YouTubeTranscriptError as exc:
        error: dict[str, object] = {"code": exc.code, "message": exc.message}
        if exc.details:
            error["details"] = exc.details
        return JSONResponse(
            status_code=400,
            content={"ok": False, "result": None, "error": error},
        )
    return JSONResponse(
        content={"ok": True, "result": result.to_dict(), "error": None}
    )


@app.post("/youtube/info")
async def youtube_info(request: YouTubeInfoRequest) -> JSONResponse:
    try:
        result = await run_in_threadpool(
            fetch_youtube_video,
            request.url,
            request.languages,
            include_transcript=request.include_transcript,
        )
    except YouTubeTranscriptError as exc:
        error: dict[str, object] = {"code": exc.code, "message": exc.message}
        if exc.details:
            error["details"] = exc.details
        return JSONResponse(
            status_code=400,
            content={"ok": False, "result": None, "error": error},
        )
    return JSONResponse(content={"ok": True, "result": result, "error": None})


@app.post("/youtube/caption")
async def youtube_caption(request: YouTubeCaptionRequest) -> Response:
    try:
        content = await fetch_signed_caption(request.url)
    except YouTubeTranscriptError as exc:
        return JSONResponse(
            status_code=400,
            content={
                "ok": False,
                "result": None,
                "error": {"code": exc.code, "message": exc.message},
            },
        )
    return Response(content=content, media_type="application/json; charset=utf-8")


@app.post("/inpaint")
async def inpaint(
    image: UploadFile = File(...),
    mask: UploadFile = File(...),
    mode: str = Form(inpainting_service.DEFAULT_MODE),
) -> Response:
    try:
        image_bytes = await image.read()
        mask_bytes = await mask.read()
        pil_image = Image.open(io.BytesIO(image_bytes))
        pil_mask = Image.open(io.BytesIO(mask_bytes))
        result, normalized_mode, elapsed = inpainting_service.run_inpaint(pil_image, pil_mask, mode=mode)
        output = io.BytesIO()
        result.save(output, format="PNG")
        return Response(
            content=output.getvalue(),
            media_type="image/png",
            headers={
                "X-Inpaint-Mode": normalized_mode,
                "X-Inpaint-Backend": inpainting_service.MODES[normalized_mode],
                "X-Inpaint-Elapsed": f"{elapsed:.3f}",
            },
        )
    except Exception as exc:
        raise HTTPException(status_code=500, detail=str(exc)) from exc


@app.post("/remove-doubao")
async def remove_doubao(
    image: UploadFile = File(...),
    mode: str = Form(inpainting_service.DEFAULT_MODE),
) -> Response:
    try:
        image_bytes = await image.read()
        pil_image = Image.open(io.BytesIO(image_bytes))
        result, normalized_mode, elapsed, metadata = run_remove_doubao(pil_image, mode=mode)
        output = io.BytesIO()
        result.save(output, format="PNG")
        return Response(
            content=output.getvalue(),
            media_type="image/png",
            headers={
                "X-Inpaint-Mode": normalized_mode,
                "X-Inpaint-Backend": inpainting_service.MODES[normalized_mode],
                "X-Inpaint-Elapsed": f"{elapsed:.3f}",
                "X-Doubao-Confidence": str(metadata["confidence"]),
                "X-Doubao-Coverage": str(metadata["coverage"]),
                "X-Doubao-Detection": str(metadata["detection"]),
                "X-Doubao-BBox": ",".join(str(v) for v in metadata["bbox"]),
            },
        )
    except Exception as exc:
        raise HTTPException(status_code=500, detail=str(exc)) from exc


@app.api_route(
    "/sub2api",
    methods=["GET", "POST", "PUT", "PATCH", "DELETE", "OPTIONS", "HEAD"],
)
async def sub2api_console() -> RedirectResponse:
    return RedirectResponse(url="/login", status_code=302)


@app.api_route(
    "/{path:path}",
    methods=["GET", "POST", "PUT", "PATCH", "DELETE", "OPTIONS", "HEAD"],
)
async def sub2api_console_assets(request: Request, path: str) -> StreamingResponse:
    return await _proxy_sub2api(request, f"/{path}")