backend/api/image_service.py

"""
Image Service API endpoints
Handles image compression, storage, and serving
"""

import base64
import re
from datetime import datetime
from pathlib import Path

import httpx
from fastapi import APIRouter, HTTPException, Response, UploadFile, File, Query
from fastapi.responses import JSONResponse
from pydantic import BaseModel, Field

from utils.public_url import get_public_base_url
from utils.storage import temp_images
from utils.image_processor import compress_and_store_image

router = APIRouter()
IMAGE_STORAGE_DIR = Path("storage/images")

# High-quality settings for reference/continuity frames (last frame of previous segment)
REFERENCE_FRAME_QUALITY = 92
REFERENCE_FRAME_MAX_WIDTH = 1920
REFERENCE_FRAME_MAX_HEIGHT = 1080

def _is_reference_frame_filename(filename: str) -> bool:
    if not filename:
        return False
    name = filename.lower()
    return bool(re.match(r"^(frame-|last-frame\.|whisper-frame-)", name) or "frame" in name and name.endswith((".jpg", ".jpeg", ".png")))


@router.post("/upload-image")
async def upload_image(
    file: UploadFile = File(...),
    reference: bool = Query(False, description="High quality for last-frame/reference uploads"),
):
    """
    Upload and host an image, returns public URL.
    Use ?reference=true when uploading a continuity/reference frame (last frame of previous segment)
    for higher quality and less downscaling.
    """
    try:
        image_bytes = await file.read()
        encoded = base64.b64encode(image_bytes).decode('utf-8')
        data_url = f"data:{file.content_type or 'image/jpeg'};base64,{encoded}"
        public_url = get_public_base_url()

        use_high_quality = reference or _is_reference_frame_filename(file.filename or "")
        if use_high_quality:
            hosted_url = await compress_and_store_image(
                data_url, public_url,
                max_width=REFERENCE_FRAME_MAX_WIDTH,
                max_height=REFERENCE_FRAME_MAX_HEIGHT,
                quality=REFERENCE_FRAME_QUALITY,
            )
        else:
            hosted_url = await compress_and_store_image(data_url, public_url)

        return JSONResponse(content={
            "url": hosted_url,
            "filename": file.filename,
        })
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Image upload failed: {str(e)}")


class HostImageUrlBody(BaseModel):
    url: str = Field(..., min_length=8, description="Public HTTPS image URL (e.g. CDN)")


@router.post("/host-image-url")
async def host_image_from_url(body: HostImageUrlBody):
    """
    Download an image server-side and host it like /upload-image (reference quality).
    Used after scraping so KIE and the planner can use a stable public URL.
    """
    url = body.url.strip()
    try:
        async with httpx.AsyncClient(timeout=30.0, follow_redirects=True) as client:
            r = await client.get(url, headers={"User-Agent": "ProductShowcase/1.0"})
            r.raise_for_status()
            content = r.content
            ct = (r.headers.get("content-type") or "image/jpeg").split(";")[0].strip() or "image/jpeg"
            if not ct.startswith("image/"):
                raise HTTPException(
                    status_code=400,
                    detail=f"URL did not return an image (content-type: {ct})",
                )
    except httpx.HTTPError as e:
        raise HTTPException(status_code=502, detail=f"Could not download image: {e}")

    try:
        b64 = base64.b64encode(content).decode("utf-8")
        data_url = f"data:{ct};base64,{b64}"
        public_url = get_public_base_url()
        hosted_url = await compress_and_store_image(
            data_url,
            public_url,
            max_width=REFERENCE_FRAME_MAX_WIDTH,
            max_height=REFERENCE_FRAME_MAX_HEIGHT,
            quality=REFERENCE_FRAME_QUALITY,
        )
        return JSONResponse(content={"url": hosted_url, "source_url": url})
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Host image failed: {str(e)}")


@router.get("/images/{image_id}")
async def serve_image(image_id: str):
    """
    Serve temporarily stored images.
    Prefer in-memory cache; fallback to disk for process restarts.
    """
    image_data = temp_images.get(image_id)
    if image_data is None:
        disk_path = IMAGE_STORAGE_DIR / f"{image_id}.jpg"
        if disk_path.exists():
            raw = disk_path.read_bytes()
            image_data = {
                "buffer": raw,
                "content_type": "image/jpeg",
            }
            # Warm memory cache for subsequent requests in this process.
            temp_images[image_id] = {
                "buffer": raw,
                "timestamp": datetime.now().timestamp(),
                "content_type": "image/jpeg",
            }
        else:
            raise HTTPException(status_code=404, detail="Image not found")
    
    return Response(
        content=image_data['buffer'],
        media_type=image_data['content_type'],
        headers={
            'Cache-Control': 'public, max-age=3600'
        }
    )