from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.responses import Response, JSONResponse
from PIL import Image
import io
import logging
from typing import Optional
from bg_remover import BackgroundRemover

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI(
    title="Background Remover API",
    description="API for removing backgrounds from images",
    version="1.0.0"
)

background_remover: Optional[BackgroundRemover] = None

def get_background_remover() -> BackgroundRemover:
    global background_remover
    if background_remover is None:
        logger.info("Initializing background remover...")
        background_remover = BackgroundRemover()
        logger.info("Background remover initialized")
    return background_remover

@app.get("/")
async def root():
    return {
        "message": "Background Remover API",
        "version": "1.0.0",
        "endpoints": {
            "/": "API information",
            "/health": "Health check",
            "/remove-background": "Remove background from image (POST)"
        }
    }

@app.get("/health")
async def health_check():
    try:
        get_background_remover()
        return {"status": "healthy", "service": "background_remover"}
    except Exception as e:
        logger.error(f"Health check failed: {str(e)}")
        raise HTTPException(status_code=503, detail=f"Service unhealthy: {str(e)}")

@app.post("/remove-background")
async def remove_background(
    file: UploadFile = File(..., description="Image file to process")
):
    try:
        if not file.content_type or not file.content_type.startswith('image/'):
            raise HTTPException(
                status_code=400,
                detail="File must be an image. Supported formats: JPEG, PNG, WEBP, etc."
            )
        
        image_data = await file.read()
        
        try:
            input_image = Image.open(io.BytesIO(image_data))
            if input_image.mode != 'RGBA':
                input_image = input_image.convert('RGB')
        except Exception as e:
            raise HTTPException(
                status_code=400,
                detail=f"Invalid image file: {str(e)}"
            )
        
        remover = get_background_remover()
        
        logger.info(f"Processing image: {file.filename}")
        output_image = remover.remove_background(input_image)
        
        output_buffer = io.BytesIO()
        output_image.save(output_buffer, format='PNG')
        output_buffer.seek(0)
        
        logger.info(f"Successfully processed image: {file.filename}")
        
        return Response(
            content=output_buffer.getvalue(),
            media_type="image/png",
            headers={
                "Content-Disposition": f"attachment; filename=removed_bg_{file.filename or 'image'}.png"
            }
        )
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Error processing image: {str(e)}", exc_info=True)
        raise HTTPException(
            status_code=500,
            detail=f"Error processing image: {str(e)}"
        )

@app.post("/remove-background/json")
async def remove_background_json(
    file: UploadFile = File(..., description="Image file to process")
):
    import base64
    
    try:
        if not file.content_type or not file.content_type.startswith('image/'):
            raise HTTPException(
                status_code=400,
                detail="File must be an image. Supported formats: JPEG, PNG, WEBP, etc."
            )
        
        image_data = await file.read()
        
        try:
            input_image = Image.open(io.BytesIO(image_data))
            if input_image.mode != 'RGBA':
                input_image = input_image.convert('RGB')
        except Exception as e:
            raise HTTPException(
                status_code=400,
                detail=f"Invalid image file: {str(e)}"
            )
        
        remover = get_background_remover()
        
        logger.info(f"Processing image: {file.filename}")
        output_image = remover.remove_background(input_image)
        
        output_buffer = io.BytesIO()
        output_image.save(output_buffer, format='PNG')
        output_buffer.seek(0)
        image_base64 = base64.b64encode(output_buffer.getvalue()).decode('utf-8')
        
        logger.info(f"Successfully processed image: {file.filename}")
        
        return JSONResponse({
            "success": True,
            "message": "Background removed successfully",
            "image": f"data:image/png;base64,{image_base64}",
            "filename": f"removed_bg_{file.filename or 'image'}.png"
        })
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Error processing image: {str(e)}", exc_info=True)
        return JSONResponse(
            status_code=500,
            content={
                "success": False,
                "message": f"Error processing image: {str(e)}"
            }
        )

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)