from fastapi import APIRouter, Header
from typing import Optional
import base64

from services.flux_client import FluxClient
from api.models import ImageGenerationRequest

router = APIRouter()

def get_image_router(flux_client: FluxClient):
    @router.post("/generate-image")
    async def generate_image(
        request: ImageGenerationRequest,
        x_session_id: Optional[str] = Header(None)
    ):
        try:
            print(f"Generating image with dimensions: {request.width}x{request.height}")
            print(f"Using prompt: {request.prompt}")

            image_bytes = await flux_client.generate_image(
                prompt=request.prompt,
                width=request.width,
                height=request.height
            )
            
            if image_bytes:
                if isinstance(image_bytes, str):
                    image_bytes = image_bytes.encode('utf-8')
                base64_image = base64.b64encode(image_bytes).decode('utf-8').strip('"')
                return {"success": True, "image_base64": base64_image}
            else:
                return {"success": False, "error": "Failed to generate image"}

        except Exception as e:
            print(f"Error generating image: {str(e)}")
            return {"success": False, "error": str(e)}
    
    return router