src/character_service.py

"""
Character Sheet Service
=======================

9-stage pipeline for generating 7-view character turnaround sheets.

Layout:
+------------------+------------------+------------------+
| Left Face Profile| Front Face       | Right Face Profile|  (3:4)
+------------------+------------------+------------------+
| Left Side Body   | Front Body | Right Side Body | Back Body |  (9:16)
+------------------+------------------+------------------+
"""

import time
import random
import logging
from pathlib import Path
from typing import Optional, Tuple, Dict, Any, Callable, List
from datetime import datetime
from PIL import Image

from .models import (
    GenerationRequest,
    GenerationResult,
    CharacterSheetConfig,
    CharacterSheetMetadata
)
from .gemini_client import GeminiClient
from .backend_router import BackendRouter, BackendType
from .utils import ensure_pil_image, save_image, sanitize_filename, preprocess_images_for_backend


logger = logging.getLogger(__name__)


class CharacterSheetService:
    """
    Service for generating 7-view character turnaround sheets.

    Pipeline (9 stages):
    0. Input normalization (face→body or body→face+body)
    1. Front face portrait
    2. Left face profile (90 degrees)
    3. Right face profile (90 degrees)
    4. Front full body (from normalized)
    5. Back full body
    6. Left side full body
    7. Right side full body
    8. Composite character sheet
    """

    def __init__(
        self,
        api_key: Optional[str] = None,
        use_pro_model: bool = False,
        config: Optional[CharacterSheetConfig] = None,
        backend: Optional[BackendType] = None,
        backend_router: Optional[BackendRouter] = None
    ):
        """
        Initialize character sheet service.

        Args:
            api_key: Gemini API key (for cloud backends)
            use_pro_model: Use Gemini Pro model (legacy, use backend param instead)
            config: Optional configuration
            backend: Specific backend to use
            backend_router: Pre-configured backend router
        """
        self.config = config or CharacterSheetConfig()

        # Determine backend
        if backend_router is not None:
            self.router = backend_router
            self.backend = backend or backend_router.default_backend
        else:
            # Determine default backend based on params
            if backend is not None:
                self.backend = backend
            elif use_pro_model:
                self.backend = BackendType.GEMINI_PRO
            else:
                self.backend = BackendType.GEMINI_FLASH

            self.router = BackendRouter(
                gemini_api_key=api_key,
                default_backend=self.backend
            )

        # For backward compatibility
        self.use_pro_model = use_pro_model
        self.client = self.router.get_client(self.backend)

        logger.info(f"CharacterSheetService initialized (backend: {self.backend.value})")

    def generate_character_sheet(
        self,
        initial_image: Optional[Image.Image],
        input_type: str = "Face Only",
        character_name: str = "Character",
        gender_term: str = "character",
        costume_description: str = "",
        costume_image: Optional[Image.Image] = None,
        face_image: Optional[Image.Image] = None,
        body_image: Optional[Image.Image] = None,
        include_costume_in_faces: bool = True,
        progress_callback: Optional[Callable[[int, int, str], None]] = None,
        stage_callback: Optional[Callable[[str, Image.Image, Dict[str, Any]], None]] = None,
        output_dir: Optional[Path] = None
    ) -> Tuple[Optional[Image.Image], str, Dict[str, Any]]:
        """
        Generate complete 7-view character turnaround sheet.

        Args:
            initial_image: Starting image (face or body)
            input_type: "Face Only", "Full Body", or "Face + Body (Separate)"
            character_name: Character name
            gender_term: "character", "man", or "woman"
            costume_description: Text costume description
            costume_image: Optional costume reference
            face_image: Face image (for Face + Body mode)
            body_image: Body image (for Face + Body mode)
            include_costume_in_faces: If True, include costume reference in face views.
                Set False for models like FLUX that confuse costume with framing.
            progress_callback: Optional callback(stage, total_stages, message)
            stage_callback: Optional callback(stage_name, image, stages_dict) called after each
                stage completes with the generated image. Enables streaming preview.
            output_dir: Optional output directory

        Returns:
            Tuple of (character_sheet, status_message, metadata)
        """
        try:
            total_stages = 9
            stages = {}

            logger.info("=" * 60)
            logger.info(f"STARTING CHARACTER SHEET: {character_name}")
            logger.info(f"Input type: {input_type}")
            logger.info(f"Costume: {costume_description or '(none)'}")
            logger.info("=" * 60)

            # Build costume instructions - separate for face and body views
            # For models like FLUX, costume refs confuse face generation
            costume_instruction_body = ""
            if costume_description:
                costume_instruction_body = f" wearing {costume_description}"
            elif costume_image:
                costume_instruction_body = " wearing the costume shown in the reference"

            # Face views only get costume instruction if flag is set
            if include_costume_in_faces:
                costume_instruction_face = costume_instruction_body
            else:
                costume_instruction_face = ""
                logger.info("Costume excluded from face views (include_costume_in_faces=False)")

            def update_progress(stage: int, message: str):
                if progress_callback:
                    progress_callback(stage, total_stages, message)
                logger.info(f"[Stage {stage}/{total_stages}] {message}")

            def notify_stage_complete(stage_name: str, image: Image.Image):
                """Notify callback when a stage completes for streaming preview."""
                if stage_callback and image is not None:
                    stage_callback(stage_name, image, stages)

            # =================================================================
            # Stage 0: Normalize input
            # =================================================================
            update_progress(0, "Normalizing input images...")

            reference_body, reference_face = self._normalize_input(
                initial_image=initial_image,
                input_type=input_type,
                face_image=face_image,
                body_image=body_image,
                costume_instruction=costume_instruction_body,  # Body normalization uses full costume
                costume_image=costume_image,
                gender_term=gender_term,
                stages=stages,
                progress_callback=lambda msg: update_progress(0, msg)
            )

            if reference_body is None or reference_face is None:
                return None, "Failed to normalize input images", {}

            time.sleep(1)

            # =================================================================
            # FACE VIEWS (3 portraits)
            # =================================================================

            # Stage 1: Front face portrait
            update_progress(1, "Generating front face portrait...")

            if input_type == "Face + Body (Separate)":
                prompt = f"Generate a close-up frontal facial portrait showing the {gender_term} from the first image (body/costume reference), extrapolate and extract exact facial details from the second image (face reference). Do NOT transfer clothing or hair style from the second image. The face should fill the entire vertical space, neutral grey background with professional studio lighting."
                input_images = [reference_body, reference_face]
            else:
                prompt = f"Generate a formal portrait view of this {gender_term}{costume_instruction_face} as depicted in the reference images, in front of a neutral grey background with professional studio lighting. The face should fill the entire vertical space. Maintain exact facial features from the reference."
                input_images = [reference_face, reference_body]
                # Only include costume in face views if flag is set (smarter models)
                if costume_image and include_costume_in_faces:
                    input_images.append(costume_image)

            front_face, status = self._generate_stage(
                prompt=prompt,
                input_images=input_images,
                aspect_ratio=self.config.face_aspect_ratio,
                temperature=self.config.face_temperature
            )

            if front_face is None:
                return None, f"Stage 1 failed: {status}", {}

            stages['front_face'] = front_face
            notify_stage_complete('front_face', front_face)
            time.sleep(1)

            # Stage 2: Left face profile
            update_progress(2, "Generating left face profile...")

            prompt = f"Create a left side profile view (90 degrees) of this {gender_term}'s face{costume_instruction_face}, showing the left side of the face filling the frame. Professional studio lighting against a neutral grey background. Maintain exact facial features from the reference."

            input_images = [front_face, reference_body]
            if input_type == "Face + Body (Separate)":
                input_images.append(reference_face)
            elif costume_image and include_costume_in_faces:
                # Only include costume in face views if flag is set (smarter models)
                input_images.append(costume_image)

            left_face, status = self._generate_stage(
                prompt=prompt,
                input_images=input_images,
                aspect_ratio=self.config.face_aspect_ratio,
                temperature=self.config.face_temperature
            )

            if left_face is None:
                return None, f"Stage 2 failed: {status}", {}

            stages['left_face'] = left_face
            notify_stage_complete('left_face', left_face)
            time.sleep(1)

            # Stage 3: Right face profile
            update_progress(3, "Generating right face profile...")

            prompt = f"Create a right side profile view (90 degrees) of this {gender_term}'s face{costume_instruction_face}, showing the right side of the face filling the frame. Professional studio lighting against a neutral grey background. Maintain exact facial features from the reference."

            input_images = [front_face, reference_body]
            if input_type == "Face + Body (Separate)":
                input_images.append(reference_face)
            elif costume_image and include_costume_in_faces:
                # Only include costume in face views if flag is set (smarter models)
                input_images.append(costume_image)

            right_face, status = self._generate_stage(
                prompt=prompt,
                input_images=input_images,
                aspect_ratio=self.config.face_aspect_ratio,
                temperature=self.config.face_temperature
            )

            if right_face is None:
                return None, f"Stage 3 failed: {status}", {}

            stages['right_face'] = right_face
            notify_stage_complete('right_face', right_face)
            time.sleep(1)

            # =================================================================
            # BODY VIEWS (4 views)
            # =================================================================

            # Stage 4: Front body (use normalized reference)
            update_progress(4, "Using front body from normalized reference...")
            front_body = reference_body
            stages['front_body'] = front_body
            notify_stage_complete('front_body', front_body)
            time.sleep(1)

            # Stage 5: Back body
            update_progress(5, "Generating back full body...")

            prompt = f"Generate a rear view image of this {gender_term}{costume_instruction_body} showing the back in a neutral standing pose against a neutral grey background with professional studio lighting. The full body should fill the vertical space. Maintain consistent appearance from the reference images."

            input_images = [reference_body, front_face]
            if costume_image:
                input_images.append(costume_image)

            back_body, status = self._generate_stage(
                prompt=prompt,
                input_images=input_images,
                aspect_ratio=self.config.body_aspect_ratio,
                temperature=self.config.body_temperature
            )

            if back_body is None:
                return None, f"Stage 5 failed: {status}", {}

            stages['back_body'] = back_body
            notify_stage_complete('back_body', back_body)
            time.sleep(1)

            # Stage 6: Left side body
            update_progress(6, "Generating left side full body...")

            prompt = f"Generate a left side view body of this {gender_term}{costume_instruction_body} from the left side in front of a neutral grey background. The {gender_term} should be shown from the left side (90 degree angle) in a neutral standing pose. Left side view. Full body fills vertical space. Professional studio lighting."

            input_images = [left_face, front_body, reference_body]

            left_body, status = self._generate_stage(
                prompt=prompt,
                input_images=input_images,
                aspect_ratio=self.config.body_aspect_ratio,
                temperature=self.config.body_temperature
            )

            if left_body is None:
                return None, f"Stage 6 failed: {status}", {}

            stages['left_body'] = left_body
            notify_stage_complete('left_body', left_body)
            time.sleep(1)

            # Stage 7: Right side body
            update_progress(7, "Generating right side full body...")

            prompt = f"Generate a right side view body of this {gender_term}{costume_instruction_body} from the right side in front of a neutral grey background. The {gender_term} should be shown from the right side (90 degree angle) in a neutral standing pose. Right side view. Full body fills vertical space. Professional studio lighting."

            input_images = [right_face, front_body, reference_body]

            right_body, status = self._generate_stage(
                prompt=prompt,
                input_images=input_images,
                aspect_ratio=self.config.body_aspect_ratio,
                temperature=self.config.body_temperature
            )

            if right_body is None:
                return None, f"Stage 7 failed: {status}", {}

            stages['right_body'] = right_body
            notify_stage_complete('right_body', right_body)
            time.sleep(1)

            # =================================================================
            # Stage 8: Composite character sheet
            # =================================================================
            update_progress(8, "Compositing character sheet...")

            character_sheet = self.composite_character_sheet(
                left_face=left_face,
                front_face=front_face,
                right_face=right_face,
                left_body=left_body,
                front_body=front_body,
                right_body=right_body,
                back_body=back_body
            )

            stages['character_sheet'] = character_sheet

            # Build metadata
            metadata = CharacterSheetMetadata(
                character_name=character_name,
                input_type=input_type,
                costume_description=costume_description,
                backend=self.router.get_active_backend_name(),
                stages={
                    "left_face": {"size": left_face.size},
                    "front_face": {"size": front_face.size},
                    "right_face": {"size": right_face.size},
                    "left_body": {"size": left_body.size},
                    "front_body": {"size": front_body.size},
                    "right_body": {"size": right_body.size},
                    "back_body": {"size": back_body.size},
                }
            )

            success_msg = f"Character sheet generated! 7 views of {character_name}"

            # Save to disk if requested
            if output_dir:
                save_dir = self._save_outputs(
                    character_name=character_name,
                    stages=stages,
                    output_dir=output_dir
                )
                success_msg += f"\nSaved to: {save_dir}"

            update_progress(9, "Complete!")
            return character_sheet, success_msg, {"metadata": metadata, "stages": stages}

        except Exception as e:
            logger.exception(f"Character sheet generation failed: {e}")
            return None, f"Error: {str(e)}", {}

    def _normalize_input(
        self,
        initial_image: Optional[Image.Image],
        input_type: str,
        face_image: Optional[Image.Image],
        body_image: Optional[Image.Image],
        costume_instruction: str,
        costume_image: Optional[Image.Image],
        gender_term: str,
        stages: dict,
        progress_callback: Optional[Callable]
    ) -> Tuple[Optional[Image.Image], Optional[Image.Image]]:
        """Normalize input images to create reference body and face."""

        if input_type == "Face + Body (Separate)":
            if face_image is None or body_image is None:
                return None, None

            if progress_callback:
                progress_callback("Normalizing body image...")

            prompt = f"Front view full body portrait of this person{costume_instruction}, standing, neutral background"
            input_images = [body_image, face_image]
            if costume_image:
                input_images.append(costume_image)

            normalized_body, _ = self._generate_stage(
                prompt=prompt,
                input_images=input_images,
                aspect_ratio=self.config.body_aspect_ratio,
                temperature=self.config.normalize_temperature
            )

            if normalized_body is None:
                return None, None

            stages['normalized_body'] = normalized_body
            return normalized_body, face_image

        elif input_type == "Face Only":
            if initial_image is None:
                return None, None

            if progress_callback:
                progress_callback("Generating full body from face...")

            prompt = f"Create a full body image of the {gender_term}{costume_instruction} standing in a neutral pose in front of a grey background with professional studio lighting. The {gender_term}'s face and features should match the reference image exactly."

            input_images = [initial_image]
            if costume_image:
                input_images.append(costume_image)

            full_body, _ = self._generate_stage(
                prompt=prompt,
                input_images=input_images,
                aspect_ratio=self.config.body_aspect_ratio,
                temperature=self.config.normalize_temperature
            )

            if full_body is None:
                return None, None

            stages['generated_body'] = full_body
            return full_body, initial_image

        else:  # Full Body
            if initial_image is None:
                return None, None

            # Normalize body
            if progress_callback:
                progress_callback("Normalizing full body...")

            prompt = f"Front view full body portrait of this person{costume_instruction}, standing, neutral background"

            input_images = [initial_image]
            if costume_image:
                input_images.append(costume_image)

            normalized_body, _ = self._generate_stage(
                prompt=prompt,
                input_images=input_images,
                aspect_ratio=self.config.body_aspect_ratio,
                temperature=self.config.normalize_temperature
            )

            if normalized_body is None:
                return None, None

            stages['normalized_body'] = normalized_body
            time.sleep(1)

            # Extract face
            if progress_callback:
                progress_callback("Generating face closeup...")

            prompt = f"Create a frontal closeup portrait of this {gender_term}'s face{costume_instruction}, focusing only on the face and head. Professional studio lighting against a neutral grey background. The face should fill the entire vertical space. Maintain exact facial features from the reference."

            input_images = [normalized_body, initial_image]
            if costume_image:
                input_images.append(costume_image)

            face_closeup, _ = self._generate_stage(
                prompt=prompt,
                input_images=input_images,
                aspect_ratio=self.config.face_aspect_ratio,
                temperature=self.config.face_temperature
            )

            if face_closeup is None:
                return None, None

            stages['extracted_face'] = face_closeup
            return normalized_body, face_closeup

    def _generate_stage(
        self,
        prompt: str,
        input_images: List[Image.Image],
        aspect_ratio: str,
        temperature: float,
        max_retries: int = 3
    ) -> Tuple[Optional[Image.Image], str]:
        """Generate single stage with retry logic."""

        modified_prompt = prompt
        cfg = self.config

        # Preprocess images for the current backend
        backend_type = self.backend.value if self.backend else "unknown"
        processed_images = preprocess_images_for_backend(
            input_images, backend_type, aspect_ratio
        )
        logger.info(f"Preprocessed {len(processed_images)} images for {backend_type}")

        for attempt in range(max_retries):
            try:
                if attempt > 0:
                    wait_time = cfg.retry_delay
                    logger.info(f"Retry {attempt + 1}/{max_retries}, waiting {wait_time}s...")
                    time.sleep(wait_time)

                request = GenerationRequest(
                    prompt=modified_prompt,
                    input_images=processed_images,
                    aspect_ratio=aspect_ratio,
                    temperature=temperature
                )

                result = self.client.generate(request)

                if result.success:
                    delay = random.uniform(cfg.rate_limit_delay_min, cfg.rate_limit_delay_max)
                    time.sleep(delay)
                    return result.image, result.message

                # Check for safety block
                error_upper = result.message.upper()
                if any(kw in error_upper for kw in ['SAFETY', 'BLOCKED', 'PROHIBITED', 'IMAGE_OTHER']):
                    if 'wearing' not in modified_prompt.lower():
                        if 'body' in modified_prompt.lower():
                            modified_prompt = prompt + ", fully clothed in casual wear"
                        else:
                            modified_prompt = prompt + ", wearing appropriate clothing"
                        logger.info("Modified prompt to avoid safety filters")

                logger.warning(f"Attempt {attempt + 1} failed: {result.message}")

            except Exception as e:
                logger.error(f"Attempt {attempt + 1} exception: {e}")
                if attempt == max_retries - 1:
                    return None, str(e)

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

    def composite_character_sheet(
        self,
        left_face: Image.Image,
        front_face: Image.Image,
        right_face: Image.Image,
        left_body: Image.Image,
        front_body: Image.Image,
        right_body: Image.Image,
        back_body: Image.Image
    ) -> Image.Image:
        """
        Composite all 7 views into character sheet.

        Layout:
        +------------------+------------------+------------------+
        | Left Face Profile| Front Face       | Right Face Profile|
        +------------------+------------------+------------------+
        | Left Side Body   | Front Body | Right Side Body | Back Body |
        +------------------+------------------+------------------+
        """
        # Normalize all inputs
        left_face = ensure_pil_image(left_face, "left_face")
        front_face = ensure_pil_image(front_face, "front_face")
        right_face = ensure_pil_image(right_face, "right_face")
        left_body = ensure_pil_image(left_body, "left_body")
        front_body = ensure_pil_image(front_body, "front_body")
        right_body = ensure_pil_image(right_body, "right_body")
        back_body = ensure_pil_image(back_body, "back_body")

        spacing = self.config.spacing

        # Calculate dimensions
        face_row_width = left_face.width + front_face.width + right_face.width
        body_row_width = left_body.width + front_body.width + right_body.width + back_body.width
        canvas_width = max(face_row_width, body_row_width)
        canvas_height = front_face.height + spacing + front_body.height

        # Create canvas
        canvas = Image.new('RGB', (canvas_width, canvas_height), color=self.config.background_color)

        # Upper row: 3 face portraits
        x = 0
        canvas.paste(left_face, (x, 0))
        x += left_face.width
        canvas.paste(front_face, (x, 0))
        x += front_face.width
        canvas.paste(right_face, (x, 0))

        # Lower row: 4 body views
        x = 0
        y = front_face.height + spacing
        canvas.paste(left_body, (x, y))
        x += left_body.width
        canvas.paste(front_body, (x, y))
        x += front_body.width
        canvas.paste(right_body, (x, y))
        x += right_body.width
        canvas.paste(back_body, (x, y))

        return canvas

    def extract_views_from_sheet(
        self,
        character_sheet: Image.Image
    ) -> Dict[str, Image.Image]:
        """
        Extract individual views from character sheet.

        Returns:
            Dictionary with 7 extracted views
        """
        sheet_width, sheet_height = character_sheet.size
        spacing = self.config.spacing

        # Find separator by scanning for dark bar
        scan_start = sheet_height // 3
        scan_end = (2 * sheet_height) // 3

        min_brightness = 255
        separator_y = scan_start

        for y in range(scan_start, scan_end):
            line = character_sheet.crop((0, y, min(200, sheet_width), y + 1))
            pixels = list(line.getdata())
            avg_brightness = sum(
                sum(p[:3]) / 3 if isinstance(p, tuple) else p
                for p in pixels
            ) / len(pixels)

            if avg_brightness < min_brightness:
                min_brightness = avg_brightness
                separator_y = y

        face_height = separator_y
        body_start_y = separator_y + spacing
        body_height = sheet_height - body_start_y

        # Calculate widths from aspect ratios
        face_width = (face_height * 3) // 4
        body_width = (body_height * 9) // 16

        # Extract views
        views = {
            'left_face': character_sheet.crop((0, 0, face_width, face_height)),
            'front_face': character_sheet.crop((face_width, 0, 2 * face_width, face_height)),
            'right_face': character_sheet.crop((2 * face_width, 0, 3 * face_width, face_height)),
            'left_body': character_sheet.crop((0, body_start_y, body_width, body_start_y + body_height)),
            'front_body': character_sheet.crop((body_width, body_start_y, 2 * body_width, body_start_y + body_height)),
            'right_body': character_sheet.crop((2 * body_width, body_start_y, 3 * body_width, body_start_y + body_height)),
            'back_body': character_sheet.crop((3 * body_width, body_start_y, 4 * body_width, body_start_y + body_height)),
        }

        return views

    def _save_outputs(
        self,
        character_name: str,
        stages: dict,
        output_dir: Path
    ) -> Path:
        """Save all outputs to directory."""
        output_dir = Path(output_dir)
        safe_name = sanitize_filename(character_name)
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        char_dir = output_dir / f"{safe_name}_{timestamp}"
        char_dir.mkdir(parents=True, exist_ok=True)

        for name, image in stages.items():
            if isinstance(image, Image.Image):
                save_image(image, char_dir, f"{safe_name}_{name}")

        logger.info(f"Saved outputs to: {char_dir}")
        return char_dir