import os
import json
from fastapi import FastAPI, UploadFile, File, HTTPException, Security, Depends, Request
from fastapi.security.api_key import APIKeyHeader
from starlette.status import HTTP_403_FORBIDDEN, HTTP_429_TOO_MANY_REQUESTS
from fastapi.responses import StreamingResponse
from io import BytesIO
from PIL import Image

# Import the background removal function
from rembg import remove

# --- Rate Limiting ---
from slowapi import Limiter, _rate_limit_exceeded_handler
from slowapi.util import get_remote_address
from slowapi.errors import RateLimitExceeded

# Initialize the rate limiter
# This allows 5 requests per minute per IP address. Adjust as needed.
limiter = Limiter(key_func=get_remote_address, default_limits=["5/minute"])

# --- API Initialization ---
app = FastAPI(
    title="Background Remover API",
    description="A simple API to remove the background from images, hosted on Hugging Face Spaces.",
    version="1.0.0"
)
# Add the rate limiter to the app
app.state.limiter = limiter
app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)


# --- Security: API Key Authentication ---

# Define the API Key header
api_key_header = APIKeyHeader(name="X-API-Key", auto_error=False)

# Load API keys from Hugging Face Secrets (as a JSON string)
# For local testing, you can create a .env file and use python-dotenv
# Example JSON string: '{"your-internal-key-for-website": "internal", "a-paid-user-key": "paid"}'
API_KEYS_JSON = os.getenv("API_KEYS_JSON")
if not API_KEYS_JSON:
    # Fallback for local development if the secret is not set
    print("WARNING: API_KEYS_JSON secret not found. Using dummy keys for local development.")
    API_KEYS = {
        "my-secret-website-key": "internal",
        "paid-user-123-key": "paid"
    }
else:
    API_KEYS = json.loads(API_KEYS_JSON)


async def get_api_key(api_key: str = Security(api_key_header)):
    """
    Dependency function to validate the API key from the X-API-Key header.
    """
    if not api_key or api_key not in API_KEYS:
        raise HTTPException(
            status_code=HTTP_403_FORBIDDEN, detail="Could not validate credentials"
        )
    return {"key": api_key, "type": API_KEYS[api_key]}

# --- API Endpoints ---

@app.get("/", tags=["General"])
def read_root():
    """ A simple health check endpoint. """
    return {"status": "ok", "message": "Welcome to the Background Remover API!"}


@app.post("/remove-background/",
    tags=["Image Processing"],
    summary="Remove background from an image",
    response_description="A PNG image with a transparent background.",
    responses={
        200: {
            "content": {"image/png": {}},
            "description": "Successful Response",
        },
        400: {"description": "Invalid file type or file too large"},
        401: {"description": "Missing or invalid API Key"},
        429: {"description": "Rate limit exceeded"},
    }
)
@limiter.limit("5/minute") # Apply the rate limit to this specific endpoint
async def remove_background(
    request: Request, # Required for the rate limiter
    file: UploadFile = File(..., description="Image file to process (JPG, PNG)"),
    api_key: dict = Depends(get_api_key) # This enforces API key authentication
):
    """
    This endpoint receives an image, validates it, removes the background,
    and returns the resulting image with a transparent background.

    - **Authentication**: Requires a valid `X-API-Key` in the header.
    - **Rate Limiting**: Limited to 5 requests per minute per IP.
    - **File Validation**: Accepts only `image/jpeg` and `image/png` up to 10MB.
    """
    # --- Input Validation ---
    # 1. Check file type
    if file.content_type not in ["image/jpeg", "image/png"]:
        raise HTTPException(status_code=400, detail="Invalid file type. Please upload a JPG or PNG.")

    # 2. Check file size (e.g., 10 MB limit)
    # The `await file.read()` will load the whole file into memory.
    # We read it once and then use the bytes.
    contents = await file.read()
    if len(contents) > 10 * 1024 * 1024: # 10 MB
        raise HTTPException(status_code=400, detail="File too large. Maximum size is 10 MB.")

    # Log the request (you can expand this for analytics)
    print(f"Processing request for API key type: {api_key['type']}")

    # --- Core Logic: Background Removal ---
    try:
        # Use rembg to remove the background
        output_bytes = remove(contents)

        # Return the result as a streaming response (efficient for binary data)
        return StreamingResponse(BytesIO(output_bytes), media_type="image/png")

    except Exception as e:
        # If anything goes wrong during processing
        print(f"Error processing image: {e}")
        raise HTTPException(status_code=500, detail="Internal server error while processing the image.")