import os
import io
from fastapi import FastAPI, UploadFile, File, HTTPException, Security
from fastapi.security import APIKeyHeader
from fastapi.middleware.cors import CORSMiddleware  # Required for website access
from PIL import Image
from starlette.status import HTTP_403_FORBIDDEN
from fastapi.responses import Response

# ---------------------------------------------------------
# 1. APP CONFIGURATION
# ---------------------------------------------------------
app = FastAPI(
    title="Secure Image Compressor API",
    description="Compresses images via API using a secure key. CORS enabled for external websites.",
    version="1.0"
)

# ---------------------------------------------------------
# 2. CORS MIDDLEWARE (Crucial for Website Integration)
# ---------------------------------------------------------
# This allows your API to be called from ANY domain (your portfolio, localhost, etc.)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Allows requests from ANY website
    allow_credentials=True,
    allow_methods=["*"],  # Allows GET, POST, etc.
    allow_headers=["*"],  # Allows Custom Headers like x-api-key
)

# ---------------------------------------------------------
# 3. SECURITY CONFIGURATION
# ---------------------------------------------------------
# Define the header name clients must use
API_KEY_NAME = "x-api-key"
api_key_header = APIKeyHeader(name=API_KEY_NAME, auto_error=False)

# Retrieve the secret from Hugging Face Environment Secrets
REAL_API_KEY = os.getenv("API_SECRET_KEY")

async def get_api_key(api_key_header: str = Security(api_key_header)):
    """Validates the API Key from the header."""
    if not REAL_API_KEY:
        # Fallback if secret isn't set in Settings
        raise HTTPException(
            status_code=500, 
            detail="Server Security Config Error: API_SECRET_KEY not set in env."
        )
        
    if api_key_header == REAL_API_KEY:
        return api_key_header
    
    raise HTTPException(
        status_code=HTTP_403_FORBIDDEN, 
        detail="Could not validate credentials. Invalid API Key."
    )

# ---------------------------------------------------------
# 4. ROUTE LOGIC
# ---------------------------------------------------------

@app.get("/")
def home():
    """Health check endpoint to ensure the server is running."""
    return {"message": "Image Compressor API is active. POST to /compress with your API Key."}

@app.post("/compress")
async def compress_image(
    file: UploadFile = File(...), 
    quality: int = 60, 
    api_key: str = Security(get_api_key)
):
    """
    Uploads an image, compresses it, and returns the bytes.
    
    - **file**: The image file (JPEG, PNG, WebP).
    - **quality**: Compression quality (1-100). Default is 60.
    - **x-api-key**: Header required for auth.
    """
    
    # 1. Validate file type
    if not file.content_type.startswith("image/"):
        raise HTTPException(400, detail="File must be an image.")

    try:
        # 2. Read image into Pillow
        image_data = await file.read()
        image = Image.open(io.BytesIO(image_data))
        
        # 3. Determine output format
        # Use original format if supported, otherwise fallback to JPEG
        output_format = image.format
        if output_format not in ["JPEG", "PNG", "WEBP"]:
            output_format = "JPEG" 
        
        # Handle transparency (RGBA) if converting to a format that doesn't support it (like JPEG)
        if output_format == "JPEG" and image.mode in ("RGBA", "P"):
            image = image.convert("RGB")

        # 4. Compress Image into Memory Buffer
        buffer = io.BytesIO()
        
        # optimize=True is CPU intensive but provides better compression
        image.save(
            buffer, 
            format=output_format, 
            quality=quality, 
            optimize=True
        )
        
        # Move cursor to start of buffer to read it
        buffer.seek(0)
        compressed_bytes = buffer.getvalue()
        
        # 5. Calculate stats for headers
        original_size = len(image_data)
        compressed_size = len(compressed_bytes)
        savings = original_size - compressed_size
        
        # 6. Return the raw image bytes with correct media type and stats headers
        return Response(
            content=compressed_bytes, 
            media_type=f"image/{output_format.lower()}",
            headers={
                "X-Original-Size": str(original_size),
                "X-Compressed-Size": str(compressed_size),
                "X-Savings-Bytes": str(savings)
            }
        )

    except Exception as e:
        # Catch unexpected errors (corrupt files, memory issues)
        raise HTTPException(500, detail=f"Image processing failed: {str(e)}")