import os
import time
import logging
from typing import List, Optional
from sqlmodel import Session
from app.models.models import ModelConfig
from google import genai
from google.genai import types
from PIL import Image
import io

logger = logging.getLogger(__name__)

class AIService:
    def __init__(self, session: Session):
        self.session = session
        
    def _get_client(self, model_type: str):
        from app.cruds.crud_config import get_active_config
        config = get_active_config(self.session, model_type)
        
        if not config:
            raise ValueError(f"No active configuration found for {model_type} model.")
            
        if config.provider.lower() == "google":
            # Initialize Google Client
            return genai.Client(api_key=config.api_key), config.model_name
            
        raise NotImplementedError(f"Provider {config.provider} not supported yet.")

    def generate_storyboard(self, system_prompt: str, user_input: str) -> str:
        client, model_name = self._get_client("text")
        
        full_prompt = f"{system_prompt}\n\nUser Input: {user_input}\n\nPlease generate the full storyboard in JSON format as requested."
        
        max_retries = 3
        for attempt in range(max_retries):
            try:
                response = client.models.generate_content(
                    model=model_name,
                    contents=full_prompt
                )
                return response.text
            except Exception as e:
                logger.error(f"Error generating storyboard (Attempt {attempt + 1}/{max_retries}): {e}")
                if attempt < max_retries - 1:
                    time.sleep(2 ** attempt)
                    continue
                raise e

    def generate_image(self, prompt: str, context_images: List[str] = None, aspect_ratio: str = "16:9", resolution: str = "2K") -> bytes:
        client, model_name = self._get_client("image")
        
        contents = [prompt]
        if context_images:
            for img_path in context_images:
                if os.path.exists(img_path):
                    try:
                        prev_img = Image.open(img_path)
                        contents.append(prev_img)
                    except Exception as e:
                        logger.warning(f"Failed to load context image {img_path}: {e}")
                else:
                    # Log missing context image but don't fail, just skip it
                    logger.warning(f"Warning: Context image not found at {img_path}, skipping.")

        # Retry loop
        max_retries = 3
        for attempt in range(max_retries):
            try:
                logger.info(f"DEBUG: Starting image generation attempt {attempt + 1}/{max_retries} with model {model_name}...")
                logger.info(f"DEBUG: Prompt length: {len(prompt)}")
                if context_images:
                    logger.info(f"DEBUG: Context images count: {len(context_images)}")

                response = client.models.generate_content(
                    model=model_name,
                    contents=contents,
                    config=types.GenerateContentConfig(
                        image_config=types.ImageConfig(
                            aspect_ratio=aspect_ratio,
                            image_size=resolution
                        ),
                    )
                )
                logger.info(f"DEBUG: Generation API call completed for attempt {attempt + 1}")
                
                if response.parts:
                    for part in response.parts:
                        if part.inline_data is not None:
                            image_data = part.inline_data.data
                            if len(image_data) > 0:
                                logger.info(f"DEBUG: Successfully received image data ({len(image_data)} bytes)")
                                return image_data
                            else:
                                logger.warning(f"Warning: Received empty image data on attempt {attempt + 1}")
                
                # Check for text refusal/error
                if response.text:
                    logger.warning(f"Model response text (no image): {response.text}")
                    
                logger.warning(f"Attempt {attempt + 1} failed: No valid image data found in response.")
                if attempt == max_retries - 1:
                    raise ValueError(f"No image found in response after {max_retries} retries. Last response: {response.text if response.text else 'Empty'}")
                
                time.sleep(2 ** attempt)
                
            except Exception as e:
                logger.error(f"Error generating image (Attempt {attempt + 1}/{max_retries}): {e}")
                if attempt < max_retries - 1:
                    time.sleep(2 ** attempt)
                    continue
                raise e
        return b""