import os
import gradio as gr
import traceback
from PIL import Image, ImageEnhance, ImageFilter
import io
import base64
import google.generativeai as genai

# --- Environment Configuration ---
GEMINI_KEY = os.environ.get("GEMINI_KEY", "")

# --- Style Template Optimization ---
BASE_TEMPLATE = """Describe this design as a concise Flux 1.1 Pro prompt focusing on:
- Key visual elements
- Technical specifications
- Style consistency
- Functional requirements"""

STYLE_INSTRUCTIONS = {
    "General": BASE_TEMPLATE,
    "Realistic": f"{BASE_TEMPLATE}\nPHOTOREALISTIC RULES: Use photography terms, texture details, accurate lighting",
    "Kawaii": f"{BASE_TEMPLATE}\nKAWAII RULES: Rounded shapes, pastel colors, cute expressions",
    "Vector": f"{BASE_TEMPLATE}\nVECTOR RULES: Clean lines, geometric shapes, B&W gradients",
    "Silhouette": f"{BASE_TEMPLATE}\nSILHOUETTE RULES: High contrast, minimal details, strong outlines"
}

# --- Flux Configuration ---
FLUX_SPECS = {
    "aspect_ratios": ["1:1", "16:9", "4:3", "9:16"],
    "color_modes": ["B&W", "CMYK", "RGB"],
    "dpi_options": [72, 150, 300, 600]
}

# --- Image Processing Pipeline ---
def preprocess_image(img):
    try:
        if isinstance(img, str):  # Handle file paths
            img = Image.open(img)
        img = img.convert("RGB")
        img = ImageEnhance.Contrast(img).enhance(1.2)
        img = img.filter(ImageFilter.SHARPEN)
        return img
    except Exception as e:
        raise ValueError(f"🔴 Image processing failed: {str(e)}")

# --- Core Generation Engine ---
def generate_prompt(image, api_key, style, creativity, neg_prompt, aspect, color_mode, dpi):
    try:
        # Step 1: Input Validation
        if not image:
            return "", "⚠️ Please upload an image."

        api_key = api_key or GEMINI_KEY
        if not api_key:
            return "", "🔑 API key required - set in env (GEMINI_KEY) or input field."

        # Step 2: Gemini Setup
        try:
            genai.configure(api_key=api_key)
            model = genai.GenerativeModel("gemini-1.5-pro")
        except ImportError:
            return "", "🚫 Failed to import google.generativeai. Install with: pip install google-generativeai"
        except Exception as e:
            if "authentication" in str(e).lower():
                return "", "🔐 Invalid API key or authentication error."
            else:
                return "", f"⚙️ API initialization error: {str(e)}"

        # Step 3: Preprocess Image
        try:
            img = preprocess_image(image)
            img_bytes = io.BytesIO()
            img.save(img_bytes, format="PNG")
            img_b64 = base64.b64encode(img_bytes.getvalue()).decode()
        except Exception as e:
            return "", f"🖼️ Image preparation failed: {str(e)}"

        # Step 4: Build Instruction Prompt
        try:
            instruction = f"{STYLE_INSTRUCTIONS[style]}\nAVOID: {neg_prompt}\n"
            instruction += f"ASPECT: {aspect}, COLORS: {color_mode}, DPI: {dpi}\n"
        except KeyError:
            return "", "🛠️ Invalid style selected. Please choose from available options."

        # Step 5: Call Gemini API
        try:
            response = model.generate_content(
                contents=[instruction, {"mime_type": "image/png", "data": img_b64}],
                generation_config={"temperature": creativity}
            )
            raw_prompt = response.text
        except Exception as e:
            return "", f"🤖 Prompt generation failed: {str(e)}"

        return raw_prompt, "✅ Prompt generated successfully!"

    except Exception as e:
        traceback.print_exc()
        return "", f"💥 Unexpected error: {str(e)}"


# --- Main Interface ---
def build_interface():
    global STYLE_INSTRUCTIONS  # Ensure access to global dict

    with gr.Blocks(title="Flux Pro Generator") as app:
        gr.Markdown("# 🎨 Flux Pro Prompt Generator")
        gr.Markdown("Generate optimized design prompts from images using Google's Gemini.")

        with gr.Row():
            # Left Panel: Image Upload & API Key
            with gr.Column(scale=1):
                api_key = gr.Textbox(
                    label="🔑 Gemini API Key",
                    value=GEMINI_KEY,
                    type="password",
                    info="Set GEMINI_KEY environment variable for production."
                )

                img_input = gr.Image(
                    label="🖼️ Upload Design",
                    type="pil",
                    sources=["upload"],
                    interactive=True
                )

            # Right Panel: Prompt Output + Controls Below
            with gr.Column(scale=2):
                status_msg = gr.Textbox(label="📢 Status", interactive=False)
                prompt_output = gr.Textbox(
                    label="📝 Optimized Prompt",
                    lines=8,
                    interactive=True
                )

                # Controls under prompt box
                style = gr.Dropdown(
                    list(STYLE_INSTRUCTIONS.keys()),
                    value="General",
                    label="🎨 Target Style"
                )

                with gr.Accordion("⚙️ Advanced Settings", open=False):
                    creativity = gr.Slider(0.0, 1.0, 0.7, label="🧠 Creativity Level")
                    neg_prompt = gr.Textbox(label="🚫 Negative Prompts", placeholder="What to avoid")
                    aspect = gr.Dropdown(FLUX_SPECS["aspect_ratios"], value="1:1", label="Aspect Ratio")
                    color_mode = gr.Dropdown(FLUX_SPECS["color_modes"], value="RGB", label="Color Mode")
                    dpi = gr.Dropdown([str(d) for d in FLUX_SPECS["dpi_options"]], value="300", label="Output DPI")

                gen_btn = gr.Button("✨ Generate Prompt", variant="primary")

        # Event Handling
        gen_btn.click(
            fn=generate_prompt,
            inputs=[
                img_input, api_key, style, creativity,
                neg_prompt, aspect, color_mode, dpi
            ],
            outputs=[prompt_output, status_msg],
            api_name="generate"
        )

    return app


# --- Production Launch ---
if __name__ == "__main__":
    app = build_interface()
    app.launch()