main.py

import logging
import shutil
import tempfile
from pathlib import Path

import pillow_heif
import uvicorn
from fastapi import FastAPI, File, HTTPException, UploadFile
from fastapi.responses import Response

from config import load_config
from metadata import write_metadata
from vision import get_tags_and_caption

pillow_heif.register_heif_opener()

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger(__name__)

app = FastAPI(title="Image Tagger")
config = load_config()

SUPPORTED_SUFFIXES = {".heic", ".heif", ".jpg", ".jpeg"}
CONTENT_TYPES = {
    ".heic": "image/heic",
    ".heif": "image/heif",
    ".jpg": "image/jpeg",
    ".jpeg": "image/jpeg",
}


@app.post("/tag")
async def tag_image(file: UploadFile = File(...)) -> Response:
    filename = file.filename or "image"
    suffix = Path(filename).suffix.lower()

    if suffix not in SUPPORTED_SUFFIXES:
        raise HTTPException(
            status_code=415,
            detail=f"Unsupported format '{suffix}'. Accepted: {', '.join(SUPPORTED_SUFFIXES)}",
        )

    with tempfile.TemporaryDirectory() as tmpdir:
        input_path = str(Path(tmpdir) / f"input{suffix}")
        work_path = str(Path(tmpdir) / f"work{suffix}")

        # Save upload
        with open(input_path, "wb") as f:
            shutil.copyfileobj(file.file, f)

        # Copy to work file — exiftool will modify this in-place
        shutil.copy2(input_path, work_path)

        # Get tags and caption from vision API
        log.info("Calling vision API for %s", filename)
        try:
            result = await get_tags_and_caption(input_path, config)
        except Exception as exc:
            log.exception("Vision API call failed")
            raise HTTPException(status_code=502, detail=f"Vision API error: {exc}") from exc

        log.info("EN tags: %s | DE tags: %s", result.tags_en, result.tags_de)
        log.info("EN caption: %s", result.caption_en)
        log.info("DE caption: %s", result.caption_de)

        # Write metadata back to work file
        try:
            write_metadata(work_path, result)
        except Exception as exc:
            log.exception("exiftool write failed")
            raise HTTPException(status_code=500, detail=f"Metadata write error: {exc}") from exc

        # Read result into memory before temp dir is cleaned up
        with open(work_path, "rb") as f:
            data = f.read()

    media_type = CONTENT_TYPES[suffix]
    return Response(
        content=data,
        media_type=media_type,
        headers={"Content-Disposition": f'attachment; filename="{filename}"'},
    )


@app.get("/health")
def health() -> dict:
    return {"status": "ok", "model": config.vision.model, "base_url": config.vision.base_url}


if __name__ == "__main__":
    uvicorn.run("main:app", host=config.server.host, port=config.server.port, reload=False)