character_forge_image/api_server.py

"""
Character Forge API Server
===========================

REST API for automated character generation pipeline.

Workflow:
1. External tool → POST /api/v1/character/generate with description
2. Generate initial portrait with Nano Banana (Gemini)
3. Run Character Forge pipeline (6 stages)
4. Return all outputs (intermediates + composites)

Usage:
    python api_server.py

Endpoints:
    POST /api/v1/character/generate - Generate character from description
    GET /api/v1/health - Health check
    GET /api/v1/backends - Backend status

License: Apache 2.0
"""

import os
import sys
import json
import base64
import asyncio
import time
from pathlib import Path
from typing import Optional, Dict, Any, List
from datetime import datetime
from io import BytesIO
from threading import Lock

# Add parent directory to path for imports
sys.path.insert(0, str(Path(__file__).parent))

from fastapi import FastAPI, HTTPException, BackgroundTasks
from fastapi.responses import JSONResponse
from pydantic import BaseModel, Field
from PIL import Image
import uvicorn

# Import Character Forge components
from services.character_forge_service import CharacterForgeService
from core import BackendRouter
from models.generation_request import GenerationRequest
from utils.logging_utils import get_logger
from config.settings import Settings

logger = get_logger(__name__)

# =============================================================================
# RATE LIMITING & SEQUENTIAL PROCESSING
# =============================================================================

# Global lock to ensure ONLY ONE character generates at a time
generation_lock = Lock()

# Rate limiting configuration
RATE_LIMIT_CONFIG = {
    "gemini": {
        "delay_between_requests": 3.0,  # Minimum 3 seconds between API calls
        "delay_after_stage": 2.0,        # Wait 2 seconds after each stage completes
        "delay_after_safety_block": 30.0, # Wait 30 seconds after safety filter trigger
        "max_requests_per_minute": 15     # Conservative limit
    },
    "comfyui": {
        "delay_between_requests": 1.0,
        "delay_after_stage": 0.5,
        "delay_after_safety_block": 5.0,
        "max_requests_per_minute": 60
    }
}

# Track last request time for rate limiting
last_request_time = {"gemini": 0, "comfyui": 0}

def enforce_rate_limit(backend: str, delay_type: str = "delay_between_requests"):
    """
    Enforce rate limiting to avoid API bans/blacklisting.

    CRITICAL: This prevents parallel processing and enforces delays
    between requests to avoid hitting Google's rate limits.

    Args:
        backend: Backend name ("gemini" or "comfyui")
        delay_type: Type of delay to enforce
    """
    global last_request_time

    config = RATE_LIMIT_CONFIG.get(backend, RATE_LIMIT_CONFIG["gemini"])
    required_delay = config.get(delay_type, 3.0)

    # Calculate time since last request
    time_since_last = time.time() - last_request_time.get(backend, 0)

    # If not enough time has passed, wait
    if time_since_last < required_delay:
        wait_time = required_delay - time_since_last
        logger.info(f"[RATE LIMIT] Waiting {wait_time:.1f}s before next {backend} API call...")
        time.sleep(wait_time)

    # Update last request time
    last_request_time[backend] = time.time()

# =============================================================================
# API MODELS
# =============================================================================

class CharacterGenerationRequest(BaseModel):
    """Request model for character generation."""

    character_id: str = Field(..., description="Unique identifier for the character")
    description: str = Field(..., description="Text description of the character")
    character_name: Optional[str] = Field(None, description="Character name (defaults to character_id)")
    gender_term: Optional[str] = Field("character", description="Gender term: 'character', 'man', or 'woman'")
    costume_description: Optional[str] = Field(None, description="Costume/clothing description")
    backend: Optional[str] = Field(Settings.BACKEND_GEMINI, description="Backend to use for generation")
    return_intermediates: bool = Field(True, description="Return intermediate stage images")
    output_format: str = Field("base64", description="Output format: 'base64' or 'paths'")


class StageOutput(BaseModel):
    """Output model for a single generation stage."""

    stage_name: str
    status: str
    image: Optional[str] = None  # base64 encoded or path
    prompt: Optional[str] = None
    aspect_ratio: Optional[str] = None
    temperature: Optional[float] = None


class CharacterGenerationResponse(BaseModel):
    """Response model for character generation."""

    character_id: str
    character_name: str
    status: str  # "completed", "failed", "processing"
    message: str
    timestamp: str
    backend: str

    # Generated files
    initial_portrait: Optional[StageOutput] = None
    stages: Optional[Dict[str, StageOutput]] = None
    character_sheet: Optional[StageOutput] = None

    # File paths (if output_format == "paths")
    saved_to: Optional[str] = None

    # Error info
    error: Optional[str] = None


# =============================================================================
# API SERVER
# =============================================================================

app = FastAPI(
    title="Character Forge API",
    description="Automated character turnaround sheet generation pipeline",
    version="1.0.0"
)

# Initialize services
character_service = CharacterForgeService(api_key=os.environ.get("GEMINI_API_KEY"))
backend_router = BackendRouter(api_key=os.environ.get("GEMINI_API_KEY"))


# =============================================================================
# UTILITY FUNCTIONS
# =============================================================================

def image_to_base64(image: Image.Image) -> str:
    """Convert PIL Image to base64 string."""
    buffered = BytesIO()
    image.save(buffered, format="PNG")
    img_str = base64.b64encode(buffered.getvalue()).decode()
    return f"data:image/png;base64,{img_str}"


def generate_initial_portrait(description: str, backend: str, max_retries: int = 3) -> tuple[Optional[Image.Image], str]:
    """
    Generate initial frontal portrait using Nano Banana (Gemini).

    CRITICAL: Includes rate limiting and retry logic for safety filters.

    Args:
        description: Character description
        backend: Backend to use
        max_retries: Maximum retry attempts

    Returns:
        Tuple of (image, status_message)
    """
    logger.info(f"Generating initial portrait with {backend}...")
    logger.info(f"Description: {description}")

    # Create portrait-focused prompt
    base_prompt = f"Generate a high-quality frontal portrait photograph focusing on the upper shoulders and face. {description}. Professional studio lighting, neutral grey background. The face should fill the vertical space. Photorealistic, detailed facial features."

    prompt = base_prompt

    for attempt in range(max_retries):
        try:
            # CRITICAL: Enforce rate limiting BEFORE making request
            enforce_rate_limit(backend, "delay_between_requests")

            logger.info(f"Initial portrait attempt {attempt + 1}/{max_retries}")

            # Generate using backend router
            request = GenerationRequest(
                prompt=prompt,
                backend=backend,
                aspect_ratio="3:4",  # Portrait format
                temperature=0.4,
                input_images=[]
            )

            result = backend_router.generate(request)

            if result.success:
                logger.info(f"Initial portrait generated successfully: {result.image.size}")

                # CRITICAL: Wait after successful generation
                enforce_rate_limit(backend, "delay_after_stage")

                return result.image, "Success"

            # Check for safety filter blocks
            error_msg_upper = result.message.upper()
            if any(keyword in error_msg_upper for keyword in [
                'SAFETY', 'BLOCKED', 'PROHIBITED', 'CENSORED',
                'POLICY', 'NSFW', 'INAPPROPRIATE', 'IMAGE_OTHER'
            ]):
                logger.warning(f"⚠️ Safety filter triggered on attempt {attempt + 1}: {result.message}")

                # CRITICAL: Long delay after safety block
                enforce_rate_limit(backend, "delay_after_safety_block")

                # Modify prompt to add clothing if not already present
                if "wearing" not in prompt.lower() and "clothed" not in prompt.lower():
                    prompt = base_prompt + ", wearing appropriate casual clothing (shirt and pants)"
                    logger.info(f"Modified prompt to avoid safety filters: added clothing description")

                # Continue to next retry
                continue

            # Other error - retry with delay
            logger.warning(f"Attempt {attempt + 1} failed: {result.message}")
            if attempt < max_retries - 1:
                enforce_rate_limit(backend, "delay_after_safety_block")

        except Exception as e:
            logger.error(f"Attempt {attempt + 1} exception: {e}")
            if attempt < max_retries - 1:
                enforce_rate_limit(backend, "delay_after_safety_block")

    return None, f"All {max_retries} attempts failed"


def save_character_outputs(
    character_id: str,
    character_name: str,
    initial_portrait: Image.Image,
    metadata: Dict[str, Any]
) -> Path:
    """
    Save all character outputs to organized directory structure.

    Args:
        character_id: Character ID
        character_name: Character name
        initial_portrait: Initial portrait image
        metadata: Generation metadata with all stages

    Returns:
        Path to output directory
    """
    # Create output directory
    output_dir = Settings.CHARACTER_SHEETS_DIR / character_id
    output_dir.mkdir(parents=True, exist_ok=True)

    # Save initial portrait
    initial_path = output_dir / f"{character_id}_00_initial_portrait.png"
    initial_portrait.save(initial_path, format="PNG")
    logger.info(f"Saved initial portrait: {initial_path}")

    # Save all stage outputs
    stages = metadata.get("stages", {})
    stage_num = 1

    for stage_name, stage_data in stages.items():
        if isinstance(stage_data, dict) and "image" in stage_data:
            image = stage_data["image"]
            if isinstance(image, Image.Image):
                stage_path = output_dir / f"{character_id}_{stage_num:02d}_{stage_name}.png"
                image.save(stage_path, format="PNG")
                logger.info(f"Saved stage {stage_num}: {stage_path}")
                stage_num += 1

    # Save metadata
    metadata_clean = {
        "character_id": character_id,
        "character_name": character_name,
        "timestamp": metadata.get("timestamp"),
        "backend": metadata.get("backend"),
        "initial_image_type": metadata.get("initial_image_type"),
        "costume_description": metadata.get("costume_description"),
        "stages": {
            name: {
                "status": data.get("status"),
                "prompt": data.get("prompt"),
                "aspect_ratio": data.get("aspect_ratio"),
                "temperature": data.get("temperature")
            }
            for name, data in stages.items()
            if isinstance(data, dict)
        }
    }

    metadata_path = output_dir / f"{character_id}_metadata.json"
    with open(metadata_path, 'w') as f:
        json.dump(metadata_clean, f, indent=2)
    logger.info(f"Saved metadata: {metadata_path}")

    return output_dir


# =============================================================================
# API ENDPOINTS
# =============================================================================

@app.get("/")
async def root():
    """API root endpoint."""
    return {
        "name": "Character Forge API",
        "version": "1.0.0",
        "status": "operational",
        "endpoints": {
            "generate": "/api/v1/character/generate",
            "health": "/api/v1/health",
            "backends": "/api/v1/backends"
        }
    }


@app.get("/api/v1/health")
async def health_check():
    """Health check endpoint."""
    return {
        "status": "healthy",
        "timestamp": datetime.now().isoformat(),
        "service": "character-forge-api"
    }


@app.get("/api/v1/backends")
async def get_backends():
    """Get status of all available backends."""
    status = character_service.get_all_backend_status()
    return status


@app.post("/api/v1/character/generate")
async def generate_character(
    request: CharacterGenerationRequest,
    background_tasks: BackgroundTasks
) -> CharacterGenerationResponse:
    """
    Generate complete character turnaround sheet from description.

    CRITICAL: This endpoint is STRICTLY SEQUENTIAL. Only ONE character
    can be generated at a time to avoid Google API rate limits and bans.

    The entire pipeline runs to completion before responding to ensure:
    1. No parallel requests to Google API
    2. Proper delays between API calls
    3. All files saved before response
    4. Rate limits respected

    Pipeline:
    1. Generate initial frontal portrait (Nano Banana)
    2. Run Character Forge 6-stage pipeline (SEQUENTIAL)
    3. Save all outputs to disk
    4. Return response with file paths

    Args:
        request: Character generation request

    Returns:
        Character generation response with all outputs
    """

    # CRITICAL: Acquire lock to ensure ONLY ONE generation at a time
    # This prevents parallel processing which could trigger rate limits/bans
    acquired = generation_lock.acquire(blocking=True, timeout=3600)  # 1 hour max wait

    if not acquired:
        raise HTTPException(
            status_code=503,
            detail="Server busy - another character is being generated. Please retry in a few minutes."
        )

    try:
        character_id = request.character_id
        character_name = request.character_name or character_id

        logger.info("="*80)
        logger.info(f"API: SEQUENTIAL generation started for '{character_id}'")
        logger.info(f"Description: {request.description}")
        logger.info(f"Backend: {request.backend}")
        logger.info(f"Lock acquired - no other generations can run")
        logger.info("="*80)

        # Stage 1: Generate initial portrait with Nano Banana
        logger.info("[Stage 0/6] Generating initial portrait with Nano Banana...")

        initial_portrait, status = generate_initial_portrait(
            description=request.description,
            backend=request.backend
        )

        if initial_portrait is None:
            raise HTTPException(
                status_code=500,
                detail=f"Initial portrait generation failed: {status}"
            )

        # Stage 2: Run Character Forge pipeline
        # CRITICAL: This runs SEQUENTIALLY with built-in delays between stages
        logger.info("[Stages 1-6] Running Character Forge pipeline SEQUENTIALLY...")
        logger.info("Each stage waits for previous to complete + rate limit delay")

        character_sheet, message, metadata = character_service.generate_character_sheet(
            initial_image=initial_portrait,
            initial_image_type="Face Only",  # We generated a face portrait
            character_name=character_name,
            gender_term=request.gender_term,
            costume_description=request.costume_description or "",
            costume_image=None,
            face_image=None,
            body_image=None,
            backend=request.backend,
            progress_callback=None,
            output_dir=None  # We'll save manually
        )

        if character_sheet is None:
            raise HTTPException(
                status_code=500,
                detail=f"Character forge pipeline failed: {message}"
            )

        # CRITICAL: Wait before saving to ensure last API call is fully complete
        logger.info("Pipeline complete - waiting before file save to ensure API cooldown...")
        enforce_rate_limit(request.backend, "delay_after_stage")

        # Save all outputs to disk
        # CRITICAL: Files MUST be saved before returning response
        logger.info("Saving outputs to disk...")
        output_dir = save_character_outputs(
            character_id=character_id,
            character_name=character_name,
            initial_portrait=initial_portrait,
            metadata=metadata
        )

        logger.info(f"All files saved to: {output_dir}")

        # CRITICAL: Final delay before releasing lock
        # This ensures complete cooldown before next generation can start
        logger.info("Files saved - final cooldown before releasing lock...")
        enforce_rate_limit(request.backend, "delay_after_stage")

        # Build response
        response_data = {
            "character_id": character_id,
            "character_name": character_name,
            "status": "completed",
            "message": f"Character generated successfully! Saved to {output_dir}",
            "timestamp": datetime.now().isoformat(),
            "backend": request.backend,
            "saved_to": str(output_dir)
        }

        # Add stage outputs if requested
        if request.return_intermediates:
            stages_output = {}

            for stage_name, stage_data in metadata.get("stages", {}).items():
                if isinstance(stage_data, dict):
                    stage_output = StageOutput(
                        stage_name=stage_name,
                        status=stage_data.get("status", "unknown"),
                        prompt=stage_data.get("prompt"),
                        aspect_ratio=stage_data.get("aspect_ratio"),
                        temperature=stage_data.get("temperature")
                    )

                    # Add image if format is base64
                    if request.output_format == "base64" and "image" in stage_data:
                        image = stage_data["image"]
                        if isinstance(image, Image.Image):
                            stage_output.image = image_to_base64(image)

                    stages_output[stage_name] = stage_output

            response_data["stages"] = stages_output

            # Add initial portrait
            response_data["initial_portrait"] = StageOutput(
                stage_name="initial_portrait",
                status="generated",
                image=image_to_base64(initial_portrait) if request.output_format == "base64" else None,
                prompt=request.description,
                aspect_ratio="3:4",
                temperature=0.4
            )

            # Add character sheet
            response_data["character_sheet"] = StageOutput(
                stage_name="character_sheet",
                status="composited",
                image=image_to_base64(character_sheet) if request.output_format == "base64" else None,
                aspect_ratio="composite"
            )

        logger.info(f"API: Character generation completed successfully for '{character_id}'")
        logger.info("="*80)
        logger.info("SEQUENTIAL generation complete - releasing lock")
        logger.info("="*80)

        return CharacterGenerationResponse(**response_data)

    except HTTPException:
        raise
    except Exception as e:
        logger.exception(f"API: Character generation failed: {e}")
        return CharacterGenerationResponse(
            character_id=request.character_id,
            character_name=request.character_name or request.character_id,
            status="failed",
            message="Generation failed",
            timestamp=datetime.now().isoformat(),
            backend=request.backend,
            error=str(e)
        )

    finally:
        # CRITICAL: ALWAYS release the lock, even if generation fails
        # This ensures the server doesn't get stuck
        generation_lock.release()
        logger.info("Lock released - next generation can proceed")


# =============================================================================
# MAIN
# =============================================================================

def main():
    """Run API server."""
    logger.info("="*80)
    logger.info("CHARACTER FORGE API SERVER")
    logger.info("="*80)
    logger.info(f"Starting server on http://0.0.0.0:8000")
    logger.info(f"Swagger docs: http://localhost:8000/docs")
    logger.info(f"ReDoc: http://localhost:8000/redoc")
    logger.info("="*80)

    uvicorn.run(
        app,
        host="0.0.0.0",
        port=8000,
        log_level="info"
    )


if __name__ == "__main__":
    main()