from fastapi import FastAPI, File, UploadFile, Form, HTTPException
from fastapi.responses import Response
from fastapi.middleware.cors import CORSMiddleware
from typing import Optional
from utils import (
    adjust_brightness_contrast, adjust_saturation, sharpen,
    auto_white_balance, adjust_exposure, recover_shadows_highlights,
    advanced_denoise, deblur
)
from model import upscale_image

app = FastAPI(title="Advanced Image Enhancement Service")

app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])

@app.post("/api/v1/enhance")
async def enhance(
    file: UploadFile = File(...),
    # upscaling
    upscale_factor: Optional[int] = Form(1),          # 1,2,4
    restore_faces: Optional[bool] = Form(False),
    # adjustments
    brightness: Optional[float] = Form(1.0),
    contrast: Optional[float] = Form(1.0),
    saturation: Optional[float] = Form(1.0),
    sharpness: Optional[float] = Form(0.0),
    auto_white_balance: Optional[bool] = Form(False),
    exposure: Optional[float] = Form(1.0),
    shadow_boost: Optional[float] = Form(1.0),
    highlight_reduce: Optional[float] = Form(1.0),
    # denoise
    denoise: Optional[bool] = Form(False),
    denoise_strength: Optional[float] = Form(1.0),    # 0.5-2.0
    # deblur
    deblur_enabled: Optional[bool] = Form(False),
    deblur_kernel: Optional[int] = Form(5),
    # auto mode (ignores most manual settings)
    auto: Optional[bool] = Form(False)
):
    if file.content_type not in ["image/jpeg", "image/png", "image/webp"]:
        raise HTTPException(400, "Unsupported image type")
    
    contents = await file.read()
    
    # Auto mode: apply white balance, denoise, contrast stretch, and slight sharpening
    if auto:
        result = auto_white_balance(contents)
        result = advanced_denoise(result, h=0.8)
        result = adjust_brightness_contrast(result, brightness=1.1, contrast=1.05)
        result = sharpen(result, radius=0.5)
    else:
        result = contents
        if auto_white_balance:
            result = auto_white_balance(result)
        if exposure != 1.0:
            result = adjust_exposure(result, exposure)
        if brightness != 1.0 or contrast != 1.0:
            result = adjust_brightness_contrast(result, brightness, contrast)
        if saturation != 1.0:
            result = adjust_saturation(result, saturation)
        if sharpness > 0:
            result = sharpen(result, sharpness)
        if shadow_boost != 1.0 or highlight_reduce != 1.0:
            result = recover_shadows_highlights(result, shadow_boost, highlight_reduce)
        if denoise:
            # map strength 0.5-2.0 to h factor 0.5-1.5
            h_factor = max(0.5, min(1.5, denoise_strength))
            result = advanced_denoise(result, h=h_factor)
        if deblur_enabled:
            result = deblur(result, kernel_size=deblur_kernel)
    
    # Upscaling with optional face restoration (requires upscaling first)
    if upscale_factor > 1:
        result = upscale_image(result, scale=upscale_factor, restore_face=restore_faces)
    
    return Response(content=result, media_type="image/png")

@app.get("/health")
async def health():
    return {"status": "enhancer ready with 4x and advanced features"}