import os
import io
import time
import base64
import json
import urllib.request
from pathlib import Path
from typing import Dict, Any, List, Optional
from PIL import Image, ImageDraw, ImageFont

from config import Config

class FlyerGenerator:
    """
    Experimental AI Flyer Generator.
    Supports:
    1. Direct Imagen 3 generation via Gemini API.
    2. Intelligent Graphic Flyer Rendering using Pillow with modern layouts.
    3. Preset sample flyers catalog from 'contoh flyer/'.
    """

    ASPECT_RATIO_DIMENSIONS = {
        "4:5": (1080, 1350),   # Instagram Portrait / Flyer Standard
        "1:1": (1080, 1080),   # Square Feed
        "9:16": (1080, 1920)   # Story / TikTok / Reels
    }

    THEMES = {
        "corporate_blue": {
            "name": "Corporate Blue (Professional)",
            "bg_top": (10, 25, 47),
            "bg_bottom": (15, 40, 75),
            "accent": (0, 210, 255),
            "card_bg": (255, 255, 255, 22),
            "text_main": (255, 255, 255),
            "text_sub": (200, 225, 255)
        },
        "emerald_growth": {
            "name": "Emerald Growth (Bisnis & Edukasi)",
            "bg_top": (6, 35, 25),
            "bg_bottom": (12, 60, 45),
            "accent": (52, 211, 153),
            "card_bg": (255, 255, 255, 24),
            "text_main": (255, 255, 255),
            "text_sub": (210, 250, 235)
        },
        "vibrant_coral": {
            "name": "Vibrant Coral (Promo & Diskon)",
            "bg_top": (45, 12, 30),
            "bg_bottom": (80, 20, 50),
            "accent": (255, 90, 95),
            "card_bg": (255, 255, 255, 25),
            "text_main": (255, 255, 255),
            "text_sub": (255, 220, 225)
        },
        "purple_tech": {
            "name": "Purple Tech (Modern Digital)",
            "bg_top": (20, 10, 40),
            "bg_bottom": (45, 20, 85),
            "accent": (192, 132, 252),
            "card_bg": (255, 255, 255, 22),
            "text_main": (255, 255, 255),
            "text_sub": (235, 215, 255)
        }
    }

    def __init__(self):
        self.output_dir = Config.AI_FLYERS_FOLDER
        self.samples_dir = Config.CONTOH_FLYER_FOLDER
        self.output_dir.mkdir(parents=True, exist_ok=True)

    def get_sample_flyers(self) -> List[Dict[str, Any]]:
        """List all reference flyers from 'contoh flyer/' folder."""
        samples = []
        if not self.samples_dir.exists():
            return samples

        for f in self.samples_dir.glob("*.*"):
            if f.suffix.lower() in [".jpg", ".jpeg", ".png", ".webp"]:
                samples.append({
                    "filename": f.name,
                    "title": f.stem,
                    "preview_url": f"/api/experimental/sample-image/{f.name}",
                    "filepath": str(f)
                })
        return samples

    def expand_prompt_with_gemini(self, simple_prompt: str, api_key: Optional[str] = None) -> Dict[str, Any]:
        """Use Gemini to expand a simple prompt into structured Indonesian marketing flyer copy & Imagen prompt."""
        key = api_key or Config.GEMINI_API_KEY
        if not key:
            # Fallback structure
            return {
                "headline": simple_prompt.upper()[:45],
                "subtitle": "Program Unggulan & Terpercaya 2026",
                "benefits": [
                    "Materi Praktis Langsung Aplikasikan",
                    "Didampingi Mentor & Konsultan Ahli",
                    "Fasilitas Lengkap & E-Sertifikat Resmi",
                    "Jejaring Komunitas Bisnis Eksklusif"
                ],
                "cta": "Daftar & Konsultasi Sekarang: Hubungi Kami",
                "imagen_prompt": f"A professional Indonesian business marketing flyer and workshop poster about {simple_prompt}. Modern clean commercial graphic design, 4:5 vertical portrait orientation, bold typography, vibrant lighting, clean background, premium corporate visual layout."
            }

        prompt_instruction = f"""
You are a creative advertising director. Expand this simple Indonesian marketing idea: "{simple_prompt}" into a structured marketing flyer layout.
Respond strictly in valid JSON with these exact keys:
{{
  "headline": "Punchy eye-catching headline in Indonesian (max 40 chars, UPPERCASE)",
  "subtitle": "Supportive value proposition subtitle in Indonesian (max 60 chars)",
  "benefits": [
    "Benefit 1 (short bullet point)",
    "Benefit 2 (short bullet point)",
    "Benefit 3 (short bullet point)",
    "Benefit 4 (short bullet point)"
  ],
  "cta": "Compelling Call to Action text in Indonesian (max 40 chars)",
  "imagen_prompt": "Detailed photorealistic commercial flyer description in English for Imagen 3. Specifying high resolution 4:5 vertical layout, crisp marketing graphics, clean professional typography space, commercial lighting, elegant color palette."
}}
"""
        try:
            import google.generativeai as gai
            gai.configure(api_key=key)
            model = None
            for m_name in ["gemini-3.6-flash", "gemini-flash-latest"]:
                try:
                    model = gai.GenerativeModel(m_name)
                    res = model.generate_content(prompt_instruction)
                    break
                except Exception:
                    continue
            if not model:
                raise RuntimeError("Tidak ada model Gemini yang merespon.")
            text = res.text.strip()
            if text.startswith("```json"):
                text = text[7:]
            if text.startswith("```"):
                text = text[3:]
            if text.endswith("```"):
                text = text[:-3]
            data = json.loads(text.strip())
            return data
        except Exception as e:
            print(f"[GEMINI EXPAND WARN] {e}. Using fallback layout.")
            return {
                "headline": simple_prompt.upper()[:45],
                "subtitle": "Solusi Terbaik & Terpercaya 2026",
                "benefits": [
                    "Solusi Terintegrasi & Efisien",
                    "Didukung Tim Profesional Berpengalaman",
                    "Biaya Transparan & Bergaransi",
                    "Layanan Konsultasi Penuh 24/7"
                ],
                "cta": "Hubungi Kami Untuk Info & Penawaran Spesial",
                "imagen_prompt": f"A high quality marketing flyer poster about {simple_prompt}, vertical portrait 4:5 aspect ratio, corporate aesthetic, clean graphic layout."
            }

    def generate_with_imagen(self, imagen_prompt: str, api_key: str, aspect_ratio: str = "4:5") -> Optional[bytes]:
        """Generate image bytes using Google Imagen 3 REST API."""
        if not api_key:
            return None

        url = f"https://generativelanguage.googleapis.com/v1beta/models/imagen-3.0-generate-002:predict?key={api_key}"
        # Imagen accepts 1:1, 3:4, 4:3, 9:16, 16:9
        api_ratio = "3:4" if aspect_ratio in ["4:5", "3:4"] else ("9:16" if aspect_ratio == "9:16" else "1:1")

        payload = {
            "instances": [{"prompt": imagen_prompt}],
            "parameters": {
                "sampleCount": 1,
                "aspectRatio": api_ratio
            }
        }

        try:
            req = urllib.request.Request(
                url,
                data=json.dumps(payload).encode("utf-8"),
                headers={"Content-Type": "application/json"}
            )
            with urllib.request.urlopen(req, timeout=40) as resp:
                data = json.loads(resp.read().decode())
                if "predictions" in data and len(data["predictions"]) > 0:
                    b64_str = data["predictions"][0].get("bytesBase64Encoded")
                    if b64_str:
                        return base64.b64decode(b64_str)
        except Exception as e:
            print(f"[IMAGEN 3 API ERROR] {e}")
        return None

    def render_graphic_flyer(
        self,
        headline: str,
        subtitle: str,
        benefits: List[str],
        cta: str,
        theme_key: str = "corporate_blue",
        aspect_ratio: str = "4:5"
    ) -> Image.Image:
        """
        Renders a high-resolution marketing flyer using Pillow.
        Adheres to the proven Indonesian marketing poster format (as seen in contoh flyer/):
        - Top Brand / Badge
        - Prominent Headline
        - Sub-headline
        - Benefit Cards with Checkmarks
        - Bottom CTA and Contact Bar
        """
        width, height = self.ASPECT_RATIO_DIMENSIONS.get(aspect_ratio, (1080, 1350))
        theme = self.THEMES.get(theme_key, self.THEMES["corporate_blue"])

        # Create vertical gradient background
        base = Image.new("RGBA", (width, height), (0, 0, 0, 255))
        draw = ImageDraw.Draw(base)

        r1, g1, b1 = theme["bg_top"]
        r2, g2, b2 = theme["bg_bottom"]

        for y in range(height):
            ratio = y / height
            r = int(r1 + (r2 - r1) * ratio)
            g = int(g1 + (g2 - g1) * ratio)
            b = int(b1 + (b2 - b1) * ratio)
            draw.line([(0, y), (width, y)], fill=(r, g, b, 255))

        # Ambient decorative glow circles
        accent_color = theme["accent"]
        draw.ellipse([width - 300, -100, width + 200, 400], fill=(accent_color[0], accent_color[1], accent_color[2], 30))
        draw.ellipse([-150, height - 400, 350, height + 100], fill=(r1, g1, b1, 60))

        # Fonts - load default or truetype
        try:
            font_title = ImageFont.truetype("arialbd.ttf", 52)
            font_sub = ImageFont.truetype("arial.ttf", 28)
            font_badge = ImageFont.truetype("arialbd.ttf", 22)
            font_bullet = ImageFont.truetype("arialbd.ttf", 28)
            font_cta = ImageFont.truetype("arialbd.ttf", 34)
            font_contact = ImageFont.truetype("arial.ttf", 24)
        except Exception:
            font_title = font_sub = font_badge = font_bullet = font_cta = font_contact = ImageFont.load_default()

        # 1. Top Badge Pill
        badge_text = "OFFICIAL MARKETING PROGRAM • 2026"
        badge_w = 480
        badge_x = (width - badge_w) // 2
        badge_y = 65
        draw.rounded_rectangle([badge_x, badge_y, badge_x + badge_w, badge_y + 46], radius=23, fill=(10, 25, 45), outline=accent_color, width=2)
        draw.text((badge_x + 35, badge_y + 11), badge_text, fill=accent_color, font=font_badge)

        # 2. Main Headline
        words = headline.split()
        lines = []
        cur_line = []
        for w in words:
            if len(" ".join(cur_line + [w])) <= 20:
                cur_line.append(w)
            else:
                lines.append(" ".join(cur_line))
                cur_line = [w]
        if cur_line:
            lines.append(" ".join(cur_line))

        text_y = 150
        for line in lines[:3]:
            bbox = draw.textbbox((0, 0), line, font=font_title)
            lw = bbox[2] - bbox[0]
            draw.text(((width - lw) // 2, text_y), line, fill=(255, 255, 255), font=font_title)
            text_y += 65

        # 3. Subtitle Pill
        sub_y = text_y + 15
        sub_bbox = draw.textbbox((0, 0), subtitle, font=font_sub)
        sub_w = min(sub_bbox[2] - sub_bbox[0] + 60, width - 120)
        sub_x = (width - sub_w) // 2
        draw.rounded_rectangle([sub_x, sub_y, sub_x + sub_w, sub_y + 55], radius=12, fill=accent_color)
        draw.text((sub_x + 30, sub_y + 12), subtitle, fill=(10, 20, 35), font=font_sub)

        # 4. Feature / Benefit Cards Box (Dark High-Contrast Card with Accent Border)
        cards_y = sub_y + 90
        box_margin = 60
        box_w = width - (box_margin * 2)
        box_h = 420
        # Dark contrasting card fill so text is 100% crystal clear
        draw.rounded_rectangle([box_margin, cards_y, box_margin + box_w, cards_y + box_h], radius=20, fill=(12, 22, 38), outline=accent_color, width=2)

        # Card Title Header
        header_text = "✨ KEUNGGULAN & BENEFIT PROGRAM:"
        draw.text((box_margin + 35, cards_y + 25), header_text, fill=accent_color, font=font_badge)

        # Bullets inside box
        bullet_y = cards_y + 80
        for b in benefits[:4]:
            circle_x = box_margin + 35
            # Draw Checkmark circle icon
            draw.ellipse([circle_x, bullet_y + 2, circle_x + 32, bullet_y + 34], fill=accent_color)
            draw.text((circle_x + 9, bullet_y + 4), "v", fill=(10, 25, 45), font=font_badge)

            # Bullet text in high-contrast crisp white
            draw.text((circle_x + 50, bullet_y + 3), b[:45], fill=(255, 255, 255), font=font_bullet)
            bullet_y += 75

        # 5. Bottom CTA Banner
        cta_y = height - 220
        cta_margin = 60
        cta_w = width - (cta_margin * 2)
        draw.rounded_rectangle([cta_margin, cta_y, cta_margin + cta_w, cta_y + 90], radius=16, fill=(255, 90, 95) if theme_key == "vibrant_coral" else accent_color)
        cta_bbox = draw.textbbox((0, 0), cta, font=font_cta)
        cta_tw = cta_bbox[2] - cta_bbox[0]
        draw.text(((width - cta_tw) // 2, cta_y + 24), cta, fill=(10, 15, 25), font=font_cta)

        # 6. Bottom Contact Info
        contact_text = "DSE Automation Marketing • Info & Pemesanan Terbuka Untuk Umum"
        c_bbox = draw.textbbox((0, 0), contact_text, font=font_contact)
        c_w = c_bbox[2] - c_bbox[0]
        draw.text(((width - c_w) // 2, height - 85), contact_text, fill=theme["text_sub"], font=font_contact)

        return base.convert("RGB")

    def create_flyer(
        self,
        prompt: str,
        theme_key: str = "corporate_blue",
        aspect_ratio: str = "4:5",
        use_imagen: bool = True,
        api_key: Optional[str] = None
    ) -> Dict[str, Any]:
        """
        Main entry point for generating a flyer.
        Returns a dict containing file metadata, generated caption, and preview URL.
        """
        key = api_key or Config.GEMINI_API_KEY
        expanded = self.expand_prompt_with_gemini(prompt, api_key=key)

        filename = f"flyer_ai_{int(time.time())}.png"
        filepath = self.output_dir / filename

        generated_via = "graphic_renderer"

        # Try Imagen 3 first if enabled and API key exists
        if use_imagen and key:
            print(f"[FLYER STUDIO] Mencoba generasi menggunakan Google Imagen 3 API...")
            img_bytes = self.generate_with_imagen(expanded.get("imagen_prompt", prompt), key, aspect_ratio)
            if img_bytes:
                with open(filepath, "wb") as f:
                    f.write(img_bytes)
                generated_via = "imagen_3"
                print(f"[FLYER STUDIO] Sukses membuat gambar via Imagen 3: {filename}")

        # Fallback to smart Graphic Banner Renderer if Imagen didn't run or failed
        if not filepath.exists() or filepath.stat().st_size == 0:
            print(f"[FLYER STUDIO] Merender grafik flyer presisi tinggi via Pillow...")
            image = self.render_graphic_flyer(
                headline=expanded.get("headline", prompt.upper()),
                subtitle=expanded.get("subtitle", "Solusi Unggulan 2026"),
                benefits=expanded.get("benefits", []),
                cta=expanded.get("cta", "Daftar Sekarang: Hubungi Kami"),
                theme_key=theme_key,
                aspect_ratio=aspect_ratio
            )
            image.save(str(filepath), format="PNG", quality=95)
            generated_via = "graphic_renderer"

        # Construct suggested marketing caption
        suggested_caption = f"🔥 {expanded.get('headline', prompt)} 🔥\n\n{expanded.get('subtitle', '')}\n\nKeunggulan:\n"
        for b in expanded.get("benefits", []):
            suggested_caption += f"• {b}\n"
        suggested_caption += f"\n👉 {expanded.get('cta', 'Hubungi kami sekarang!')}\n\n#marketing #promobisnis #event #viral #dseautomation"

        return {
            "success": True,
            "filename": filename,
            "filepath": str(filepath),
            "preview_url": f"/uploads/ai_flyers/{filename}",
            "headline": expanded.get("headline", prompt),
            "subtitle": expanded.get("subtitle", ""),
            "suggested_caption": suggested_caption,
            "mode": generated_via
        }