import os from fastapi import APIRouter, HTTPException, UploadFile, File, Depends from fastapi.responses import FileResponse from pydantic import BaseModel from typing import List, Optional, Dict, Any import base64 # Import our LinkedIn image generation services from services.linkedin.image_generation import LinkedInImageGenerator, LinkedInImageStorage from services.linkedin.image_prompts import LinkedInPromptGenerator from services.onboarding.api_key_manager import APIKeyManager from middleware.auth_middleware import get_current_user # Set up logging from loguru import logger # Initialize router router = APIRouter(prefix="/api/linkedin", tags=["linkedin-image-generation"]) # Initialize services api_key_manager = APIKeyManager() image_generator = LinkedInImageGenerator(api_key_manager) prompt_generator = LinkedInPromptGenerator(api_key_manager) image_storage = LinkedInImageStorage(api_key_manager=api_key_manager) # Request/Response models class ImagePromptRequest(BaseModel): content_type: str topic: str industry: str content: str class ImageGenerationRequest(BaseModel): prompt: str content_context: Dict[str, Any] aspect_ratio: Optional[str] = "1:1" class ImagePromptResponse(BaseModel): style: str prompt: str description: str prompt_index: int enhanced_at: Optional[str] = None linkedin_optimized: Optional[bool] = None fallback: Optional[bool] = None content_context: Optional[Dict[str, Any]] = None class ImageGenerationResponse(BaseModel): success: bool image_url: Optional[str] = None image_id: Optional[str] = None style: Optional[str] = None aspect_ratio: Optional[str] = None error: Optional[str] = None class ImageEditRequest(BaseModel): image_base64: Optional[str] = None image_id: Optional[str] = None prompt: str content_context: Dict[str, Any] class ImageEditResponse(BaseModel): success: bool image_data: Optional[str] = None image_id: Optional[str] = None image_url: Optional[str] = None width: Optional[int] = None height: Optional[int] = None provider: Optional[str] = None model: Optional[str] = None error: Optional[str] = None @router.post("/generate-image-prompts", response_model=List[ImagePromptResponse]) async def generate_image_prompts(request: ImagePromptRequest): """ Generate three AI-optimized image prompts for LinkedIn content """ try: logger.info(f"Generating image prompts for {request.content_type} about {request.topic}") # Use our LinkedIn prompt generator service prompts = await prompt_generator.generate_three_prompts({ 'content_type': request.content_type, 'topic': request.topic, 'industry': request.industry, 'content': request.content }) logger.info(f"Generated {len(prompts)} image prompts successfully") return prompts except Exception as e: logger.error(f"Error generating image prompts: {str(e)}") raise HTTPException(status_code=500, detail=f"Failed to generate image prompts: {str(e)}") @router.post("/generate-image", response_model=ImageGenerationResponse) async def generate_linkedin_image( request: ImageGenerationRequest, current_user: Dict[str, Any] = Depends(get_current_user) ): """ Generate LinkedIn-optimized image from selected prompt """ try: user_id = current_user.get("id") logger.info(f"Generating LinkedIn image with prompt: {request.prompt[:100]}... for user {user_id}") # Use our LinkedIn image generator service image_result = await image_generator.generate_image( prompt=request.prompt, content_context=request.content_context, user_id=user_id ) if image_result and image_result.get('success'): # Store the generated image image_id = await image_storage.store_image( image_data=image_result['image_data'], metadata={ 'prompt': request.prompt, 'style': request.content_context.get('style', 'Generated'), 'aspect_ratio': request.aspect_ratio, 'content_type': request.content_context.get('content_type'), 'topic': request.content_context.get('topic'), 'industry': request.content_context.get('industry') }, user_id=user_id ) logger.info(f"Image generated and stored successfully with ID: {image_id}") return ImageGenerationResponse( success=True, image_url=image_result.get('image_url'), image_id=image_id, style=request.content_context.get('style', 'Generated'), aspect_ratio=request.aspect_ratio ) else: error_msg = image_result.get('error', 'Unknown error during image generation') logger.error(f"Image generation failed: {error_msg}") return ImageGenerationResponse( success=False, error=error_msg ) except Exception as e: logger.error(f"Error generating LinkedIn image: {str(e)}") return ImageGenerationResponse( success=False, error=f"Failed to generate image: {str(e)}" ) @router.post("/edit-image", response_model=ImageEditResponse) async def edit_linkedin_image( request: ImageEditRequest, current_user: Dict[str, Any] = Depends(get_current_user) ): """ Edit a LinkedIn-optimized image using natural language. Provide the image as base64 and describe the desired edits. """ try: user_id = current_user.get("id") if not user_id: raise HTTPException(status_code=401, detail="Authentication required") if not request.prompt or not request.prompt.strip(): raise HTTPException(status_code=400, detail="Prompt is required for image editing") logger.info(f"Editing LinkedIn image with prompt: {request.prompt[:100]}... for user {user_id}") # Get input image bytes — from image_id (fetch from storage) or image_base64 (direct decode) input_image_bytes = None if request.image_id: stored = await image_storage.retrieve_image(request.image_id, user_id) if not stored or not stored.get('success'): raise HTTPException(status_code=404, detail=f"Image not found: {request.image_id}") input_image_bytes = stored['image_data'] logger.info(f"Fetched image {request.image_id} from storage ({len(input_image_bytes)} bytes)") elif request.image_base64: input_image_bytes = base64.b64decode(request.image_base64) else: raise HTTPException(status_code=400, detail="Either image_id or image_base64 is required") # Use LinkedIn image generator with common editing infrastructure image_result = await image_generator.edit_image( input_image_bytes=input_image_bytes, edit_prompt=request.prompt, content_context=request.content_context, user_id=user_id, ) if image_result and image_result.get('success'): image_b64 = base64.b64encode(image_result['image_data']).decode("utf-8") # Store the edited image — log but don't fail if storage has issues new_image_id = None stored_result = await image_storage.store_image( image_data=image_result['image_data'], metadata={ 'prompt': request.prompt, 'style': request.content_context.get('style', 'Edited'), 'content_type': request.content_context.get('content_type'), 'topic': request.content_context.get('topic'), 'industry': request.content_context.get('industry'), 'is_edit': True, 'original_prompt': request.prompt, 'source_image_id': request.image_id, }, user_id=user_id ) if stored_result and stored_result.get('success'): new_image_id = stored_result.get('image_id') logger.info(f"Edited image stored with ID: {new_image_id}") else: logger.warning(f"Edited image not stored: {stored_result.get('error', 'unknown reason')}") return ImageEditResponse( success=True, image_data=image_b64, image_id=new_image_id, image_url=image_result.get('image_url'), width=image_result.get('width'), height=image_result.get('height'), provider=image_result.get('provider'), model=image_result.get('model'), ) else: error_msg = image_result.get('error', 'Unknown error during image editing') logger.error(f"Image editing failed: {error_msg}") return ImageEditResponse( success=False, error=error_msg ) except HTTPException: raise except Exception as e: logger.error(f"Error editing LinkedIn image: {str(e)}", exc_info=True) return ImageEditResponse( success=False, error=f"Failed to edit image: {str(e)}" ) @router.get("/image-status/{image_id}") async def get_image_status( image_id: str, current_user: Dict[str, Any] = Depends(get_current_user) ): """ Check the status of an image generation request """ try: user_id = current_user.get("id") # Get image metadata from storage metadata = await image_storage.get_image_metadata(image_id, user_id) if metadata: return { "success": True, "status": "completed", "metadata": metadata } else: return { "success": False, "status": "not_found", "error": "Image not found" } except Exception as e: logger.error(f"Error checking image status: {str(e)}") return { "success": False, "status": "error", "error": str(e) } @router.get("/images/{image_id}") async def get_generated_image( image_id: str, current_user: Dict[str, Any] = Depends(get_current_user) ): """ Retrieve a generated image by ID. Returns the image file directly as a PNG response. """ try: user_id = current_user.get("id") image_result = await image_storage.retrieve_image(image_id, user_id) if image_result.get('success') and image_result.get('image_path'): return FileResponse( path=image_result['image_path'], media_type="image/png", filename=f"{image_id}.png" ) else: raise HTTPException(status_code=404, detail="Image not found") except HTTPException: raise except Exception as e: logger.error(f"Error retrieving image: {str(e)}") raise HTTPException(status_code=500, detail=f"Failed to retrieve image: {str(e)}") @router.delete("/images/{image_id}") async def delete_generated_image( image_id: str, current_user: Dict[str, Any] = Depends(get_current_user) ): """ Delete a generated image by ID """ try: user_id = current_user.get("id") result = await image_storage.delete_image(image_id, user_id) if result.get('success'): return {"success": True, "message": "Image deleted successfully"} else: return {"success": False, "message": "Failed to delete image"} except Exception as e: logger.error(f"Error deleting image: {str(e)}") return {"success": False, "error": str(e)} # Health check endpoint @router.get("/image-generation-health") async def health_check(): """ Lightweight health check for image generation services. Verifies configuration and service availability without making API calls. """ try: services = {} all_healthy = True # Check API key configuration (no actual API call) image_api_key = api_key_manager.get_api_key("image_generation") or os.getenv("WAVESPEED_API_KEY") or os.getenv("HF_TOKEN") services["image_api_key_configured"] = bool(image_api_key) # Check storage accessibility stats = await image_storage.get_storage_stats() storage_ok = stats.get('success', False) services["image_storage"] = "operational" if storage_ok else "unavailable" if storage_ok: services["storage_stats"] = { "total_images": stats.get('total_files', 0), "total_size_gb": stats.get('total_size_gb', 0), "limit_gb": stats.get('storage_limit_gb', 0), } # Check prompt generator initialization prompt_ok = prompt_generator is not None and hasattr(prompt_generator, 'generate_three_prompts') services["prompt_generator"] = "operational" if prompt_ok else "unavailable" # Check image generator initialization gen_ok = image_generator is not None and hasattr(image_generator, 'generate_image') services["image_generator"] = "operational" if gen_ok else "unavailable" if not all(v == "operational" or v is True for v in services.values()): all_healthy = False return { "status": "healthy" if all_healthy else "degraded", "services": services } except Exception as e: logger.error(f"Health check failed: {str(e)}") return { "status": "unhealthy", "error": str(e) }