reactbin/api/app/routers/images.py

import asyncio
import logging
import re
import struct
from typing import Any

from fastapi import APIRouter, Depends, File, Form, HTTPException, Response, UploadFile
from sqlalchemy.ext.asyncio import AsyncSession

from app.auth.provider import Identity
from app.config import get_settings
from app.dependencies import get_db, get_storage, require_auth
from app.models import Image
from app.repositories.image_repo import ImageRepository
from app.repositories.tag_repo import TagRepository
from app.storage.backend import StorageBackend
from app.thumbnail import generate_thumbnail
from app.utils import compute_sha256, generate_short_id
from app.validation import FileSizeError, MimeTypeError, validate_file_size, validate_mime_type

logger = logging.getLogger(__name__)

router = APIRouter(tags=["images"])


_SHORT_ID_RE = re.compile(r"^[a-zA-Z0-9]{8}$")


def _error(detail: str, code: str, status: int):
    raise HTTPException(status_code=status, detail={"detail": detail, "code": code})


def _validate_short_id(short_id: str) -> str:
    if not _SHORT_ID_RE.match(short_id):
        raise HTTPException(
            status_code=422,
            detail={"detail": "Invalid image ID", "code": "invalid_short_id"},
        )
    return short_id


def _image_to_dict(
    image: Image, *, cdn_base: str | None = None, duplicate: bool | None = None
) -> dict[str, Any]:
    _base = cdn_base.strip().rstrip("/") if cdn_base else None
    file_url = f"{_base}/{image.storage_key}" if _base else f"/api/v1/i/{image.short_id}/file"
    thumbnail_url = (
        (f"{_base}/{image.thumbnail_key}" if _base else f"/api/v1/i/{image.short_id}/thumbnail")
        if image.thumbnail_key
        else None
    )
    data: dict[str, Any] = {
        "id": str(image.id),
        "short_id": image.short_id,
        "hash": image.hash,
        "filename": image.filename,
        "mime_type": image.mime_type,
        "size_bytes": image.size_bytes,
        "width": image.width,
        "height": image.height,
        "storage_key": image.storage_key,
        "thumbnail_key": image.thumbnail_key,
        "file_url": file_url,
        "thumbnail_url": thumbnail_url,
        "created_at": image.created_at.isoformat(),
        "tags": image.tags,
    }
    if duplicate is not None:
        data["duplicate"] = duplicate
    return data


def _read_image_dimensions(data: bytes, mime_type: str) -> tuple[int, int]:
    """Return (width, height) from raw image bytes. Falls back to (0, 0)."""
    try:
        if mime_type == "image/jpeg":
            return _jpeg_dimensions(data)
        elif mime_type == "image/png":
            return _png_dimensions(data)
        elif mime_type == "image/gif":
            return _gif_dimensions(data)
        elif mime_type == "image/webp":
            return _webp_dimensions(data)
    except Exception:
        pass
    return 0, 0


def _jpeg_dimensions(data: bytes) -> tuple[int, int]:
    i = 0
    while i < len(data):
        if data[i] != 0xFF:
            break
        i += 1
        marker = data[i]
        i += 1
        if marker in (0xD8, 0xD9):
            continue
        length = struct.unpack(">H", data[i : i + 2])[0]
        if marker in (0xC0, 0xC1, 0xC2):
            h, w = struct.unpack(">HH", data[i + 3 : i + 7])
            return w, h
        i += length
    return 0, 0


def _png_dimensions(data: bytes) -> tuple[int, int]:
    w, h = struct.unpack(">II", data[16:24])
    return w, h


def _gif_dimensions(data: bytes) -> tuple[int, int]:
    w, h = struct.unpack("<HH", data[6:10])
    return w, h


def _webp_dimensions(data: bytes) -> tuple[int, int]:
    if data[8:12] == b"VP8 ":
        w = struct.unpack("<H", data[26:28])[0] & 0x3FFF
        h = struct.unpack("<H", data[28:30])[0] & 0x3FFF
        return w, h
    elif data[8:12] == b"VP8L":
        bits = struct.unpack("<I", data[21:25])[0]
        w = (bits & 0x3FFF) + 1
        h = ((bits >> 14) & 0x3FFF) + 1
        return w, h
    return 0, 0


@router.post("/images", status_code=201)
async def upload_image(
    file: UploadFile = File(...),
    tags: str | None = Form(None),
    db: AsyncSession = Depends(get_db),
    storage: StorageBackend = Depends(get_storage),
    _: Identity = Depends(require_auth),
    settings=Depends(get_settings),
):
    data = await file.read()
    mime_type = file.content_type or "application/octet-stream"

    try:
        validate_mime_type(mime_type)
    except MimeTypeError:
        raise HTTPException(
            status_code=422,
            detail={"detail": f"Unsupported file type: {mime_type}", "code": "invalid_mime_type"},
        )

    try:
        validate_file_size(len(data), max_bytes=settings.max_upload_bytes)
    except FileSizeError as exc:
        raise HTTPException(
            status_code=422,
            detail={"detail": str(exc), "code": "file_too_large"},
        )

    hash_hex = compute_sha256(data)
    image_repo = ImageRepository(db)
    _cdn_base = settings.s3_public_base_url
    existing = await image_repo.get_by_hash(hash_hex)
    if existing:
        return Response(
            content=__import__("json").dumps(
                _image_to_dict(existing, cdn_base=_cdn_base, duplicate=True)
            ),
            status_code=200,
            media_type="application/json",
        )

    # Parse tag names
    tag_names: list[str] = []
    if tags:
        tag_repo = TagRepository(db)
        raw = [t.strip() for t in tags.replace(",", " ").split() if t.strip()]
        try:
            tag_names = [tag_repo.normalise_and_validate(t) for t in raw]
        except ValueError as exc:
            raise HTTPException(
                status_code=422,
                detail={"detail": str(exc), "code": "invalid_tag"},
            )

    width, height = _read_image_dimensions(data, mime_type)

    from sqlalchemy.exc import IntegrityError

    for _ in range(10):
        short_id = generate_short_id()
        await storage.put(short_id, data, mime_type)

        thumbnail_key: str | None = None
        try:
            thumb_bytes = await asyncio.to_thread(generate_thumbnail, data, mime_type)
            await storage.put(f"{short_id}-thumb", thumb_bytes, "image/webp")
            thumbnail_key = f"{short_id}-thumb"
        except Exception:
            logger.warning(
                "Thumbnail generation failed for %s; proceeding without thumbnail", short_id
            )

        try:
            image = await image_repo.create(
                hash_hex=hash_hex,
                filename=file.filename or "upload",
                mime_type=mime_type,
                size_bytes=len(data),
                width=width,
                height=height,
                storage_key=short_id,
                short_id=short_id,
                thumbnail_key=thumbnail_key,
            )
            break
        except IntegrityError:
            await db.rollback()
            await storage.delete(short_id)
            if thumbnail_key:
                await storage.delete(thumbnail_key)
            thumbnail_key = None
            continue
    else:
        raise HTTPException(
            status_code=500,
            detail={"detail": "Failed to assign unique ID", "code": "id_collision"},
        )

    if tag_names:
        tag_repo = TagRepository(db)
        await tag_repo.attach_tags(image, tag_names)
        image = await image_repo.reload_with_tags(image.id)

    return _image_to_dict(image, cdn_base=_cdn_base, duplicate=False)


@router.get("/images")
async def list_images(
    tags: str | None = None,
    limit: int = 50,
    offset: int = 0,
    db: AsyncSession = Depends(get_db),
    settings=Depends(get_settings),
):
    limit = min(limit, 100)
    _cdn_base = settings.s3_public_base_url
    tag_names = [t.strip() for t in tags.split(",") if t.strip()] if tags else None
    image_repo = ImageRepository(db)
    images, total = await image_repo.list_images(tag_names=tag_names, limit=limit, offset=offset)
    return {
        "items": [_image_to_dict(img, cdn_base=_cdn_base) for img in images],
        "total": total,
        "limit": limit,
        "offset": offset,
    }


@router.get("/i/{short_id}")
async def get_image(
    short_id: str,
    db: AsyncSession = Depends(get_db),
    settings=Depends(get_settings),
):
    _validate_short_id(short_id)
    _cdn_base = settings.s3_public_base_url
    image_repo = ImageRepository(db)
    image = await image_repo.get_by_short_id(short_id)
    if not image:
        raise HTTPException(
            status_code=404,
            detail={"detail": "Image not found", "code": "image_not_found"},
        )
    return _image_to_dict(image, cdn_base=_cdn_base)


@router.get("/i/{short_id}/file")
async def serve_image_file(
    short_id: str,
    db: AsyncSession = Depends(get_db),
    storage: StorageBackend = Depends(get_storage),
):
    _validate_short_id(short_id)
    image_repo = ImageRepository(db)
    image = await image_repo.get_by_short_id(short_id)
    if not image:
        raise HTTPException(
            status_code=404,
            detail={"detail": "Image not found", "code": "image_not_found"},
        )
    try:
        data = await storage.get(image.storage_key)
    except Exception:
        raise HTTPException(
            status_code=500,
            detail={"detail": "Failed to retrieve image content", "code": "storage_error"},
        ) from None
    return Response(
        content=data,
        media_type=image.mime_type,
        headers={
            "ETag": f'"{image.hash}"',
            "Cache-Control": "public, max-age=31536000, immutable",
        },
    )


@router.get("/i/{short_id}/thumbnail")
async def serve_image_thumbnail(
    short_id: str,
    db: AsyncSession = Depends(get_db),
    storage: StorageBackend = Depends(get_storage),
):
    _validate_short_id(short_id)
    image_repo = ImageRepository(db)
    image = await image_repo.get_by_short_id(short_id)
    if not image:
        raise HTTPException(
            status_code=404,
            detail={"detail": "Image not found", "code": "image_not_found"},
        )
    key = image.thumbnail_key or image.storage_key
    media_type = "image/webp" if image.thumbnail_key else image.mime_type
    try:
        data = await storage.get(key)
    except Exception:
        raise HTTPException(
            status_code=500,
            detail={"detail": "Failed to retrieve image content", "code": "storage_error"},
        ) from None
    return Response(
        content=data,
        media_type=media_type,
        headers={
            "ETag": f'"{image.hash}"',
            "Cache-Control": "public, max-age=31536000, immutable",
        },
    )


@router.patch("/i/{short_id}/tags")
async def update_image_tags(
    short_id: str,
    body: dict,
    db: AsyncSession = Depends(get_db),
    _: Identity = Depends(require_auth),
    settings=Depends(get_settings),
):
    _validate_short_id(short_id)
    _cdn_base = settings.s3_public_base_url
    image_repo = ImageRepository(db)
    image = await image_repo.get_by_short_id(short_id)
    if not image:
        raise HTTPException(
            status_code=404,
            detail={"detail": "Image not found", "code": "image_not_found"},
        )

    raw_tags: list[str] = body.get("tags", [])
    tag_repo = TagRepository(db)
    try:
        tag_names = [tag_repo.normalise_and_validate(t) for t in raw_tags]
    except ValueError as exc:
        raise HTTPException(
            status_code=422,
            detail={"detail": str(exc), "code": "invalid_tag"},
        )

    await tag_repo.replace_tags_on_image(image, tag_names)
    image = await image_repo.reload_with_tags(image.id)
    return _image_to_dict(image, cdn_base=_cdn_base)


@router.delete("/i/{short_id}", status_code=204)
async def delete_image(
    short_id: str,
    db: AsyncSession = Depends(get_db),
    storage: StorageBackend = Depends(get_storage),
    _: Identity = Depends(require_auth),
):
    _validate_short_id(short_id)
    image_repo = ImageRepository(db)
    image = await image_repo.get_by_short_id(short_id)
    if not image:
        raise HTTPException(
            status_code=404,
            detail={"detail": "Image not found", "code": "image_not_found"},
        )
    storage_key = image.storage_key
    thumbnail_key = image.thumbnail_key
    await image_repo.delete(image)
    await storage.delete(storage_key)
    if thumbnail_key:
        await storage.delete(thumbnail_key)
    return Response(status_code=204)