import os
from typing import List
from PIL import Image
from google import genai
from google.genai import types
from dotenv import load_dotenv

load_dotenv()

def get_client():
    api_key = os.getenv("GEMINI_API_KEY")
    if not api_key:
        raise ValueError("GEMINI_API_KEY not found in environment variables")
    return genai.Client(api_key=api_key)

def generate_edited_slide(
    target_image: Image.Image,
    style_reference_images: List[Image.Image],
    full_text_context: str,
    user_prompt: str,
    resolution: str = "4K",
    enable_search: bool = False
) -> Image.Image:
    """
    Sends the target image, style refs, and text context to Gemini 3 Pro Image.
    Returns the generated PIL Image.
    """
    client = get_client()

    # Construct the prompt
    prompt_parts = []

    prompt_parts.append(user_prompt)
    prompt_parts.append(target_image)

    if style_reference_images:
        prompt_parts.append("Match the visual style (fonts, colors, layout) of these reference images:")
        for img in style_reference_images:
            prompt_parts.append(img)

    if full_text_context:
        prompt_parts.append(f"DOCUMENT CONTEXT:\n{full_text_context}\n")

    # Build config
    config = types.GenerateContentConfig(
        response_modalities=['IMAGE'],
        image_config=types.ImageConfig(
            image_size=resolution
        )
    )
    if enable_search:
        config.tools = [{"google_search": {}}]

    # Call the model
    try:
        response = client.models.generate_content(
            model='gemini-3-pro-image-preview',
            contents=prompt_parts,
            config=config
        )
    except Exception as e:
        error_msg = str(e).lower()
        if "quota" in error_msg or "billing" in error_msg or "payment" in error_msg:
            raise RuntimeError(
                "Gemini API Error: This tool requires a PAID API key with billing enabled.\n"
                "Free tier keys do not support image generation. Please:\n"
                "1. Visit https://aistudio.google.com/api-keys\n"
                "2. Enable billing on your Google Cloud project\n"
                f"Original error: {e}"
            )
        elif "api key" in error_msg or "authentication" in error_msg or "unauthorized" in error_msg:
            raise RuntimeError(
                "Gemini API Error: Invalid API key.\n"
                "Please check that your GEMINI_API_KEY environment variable is set correctly.\n"
                f"Original error: {e}"
            )
        else:
            raise RuntimeError(f"Gemini API Error: {e}")

    # Extract the image from the response
    generated_image = None
    if response.candidates and response.candidates[0].content.parts:
        for part in response.candidates[0].content.parts:
            if part.inline_data:
                # Convert bytes to PIL Image
                from io import BytesIO
                generated_image = Image.open(BytesIO(part.inline_data.data))
                break
    
    if not generated_image:
        raise RuntimeError("No image generated by the model.")

    return generated_image

def generate_new_slide(
    style_reference_images: List[Image.Image],
    user_prompt: str,
    full_text_context: str = "",
    resolution: str = "4K",
    enable_search: bool = False
) -> Image.Image:
    """
    Generates a completely new slide based on style references and a prompt.
    Returns the generated PIL Image.
    """
    client = get_client()

    # Construct the prompt
    prompt_parts = []

    prompt_parts.append(user_prompt)

    if style_reference_images:
        prompt_parts.append("Match the visual style (fonts, colors, layout) of these reference images:")
        for img in style_reference_images:
            prompt_parts.append(img)

    if full_text_context:
        prompt_parts.append(f"DOCUMENT CONTEXT:\n{full_text_context}\n")

    # Build config
    config = types.GenerateContentConfig(
        response_modalities=['IMAGE'],
        image_config=types.ImageConfig(
            image_size=resolution
        )
    )
    if enable_search:
        config.tools = [{"google_search": {}}]

    # Call the model
    try:
        response = client.models.generate_content(
            model='gemini-3-pro-image-preview',
            contents=prompt_parts,
            config=config
        )
    except Exception as e:
        error_msg = str(e).lower()
        if "quota" in error_msg or "billing" in error_msg or "payment" in error_msg:
            raise RuntimeError(
                "Gemini API Error: This tool requires a PAID API key with billing enabled.\n"
                "Free tier keys do not support image generation. Please:\n"
                "1. Visit https://aistudio.google.com/api-keys\n"
                "2. Enable billing on your Google Cloud project\n"
                f"Original error: {e}"
            )
        elif "api key" in error_msg or "authentication" in error_msg or "unauthorized" in error_msg:
            raise RuntimeError(
                "Gemini API Error: Invalid API key.\n"
                "Please check that your GEMINI_API_KEY environment variable is set correctly.\n"
                f"Original error: {e}"
            )
        else:
            raise RuntimeError(f"Gemini API Error: {e}")

    # Extract the image from the response
    generated_image = None
    if response.candidates and response.candidates[0].content.parts:
        for part in response.candidates[0].content.parts:
            if part.inline_data:
                # Convert bytes to PIL Image
                from io import BytesIO
                generated_image = Image.open(BytesIO(part.inline_data.data))
                break

    if not generated_image:
        raise RuntimeError("No image generated by the model.")

    return generated_image