import gradio as gr
import os, json, re, time, base64, requests, urllib.parse
from PIL import Image, ImageDraw, ImageFont
from io import BytesIO

GEMINI_API_KEY = os.environ.get('GEMINI_API_KEY', '')
HF_TOKEN       = os.environ.get('HF_TOKEN', '')

GEMINI_TEXT_URL  = 'https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent'
GEMINI_IMAGE_URL  = 'https://generativelanguage.googleapis.com/v1beta/models/gemini-3-pro-image-preview:generateContent'
GEMINI_IMAGE_URL2 = 'https://generativelanguage.googleapis.com/v1beta/models/gemini-3.1-flash-image-preview:generateContent'
GEMINI_IMAGE_URL3 = 'https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-image:generateContent'

# ── FONTS ─────────────────────────────────────────────────────────────────────
_B = ['/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf',
      '/usr/share/fonts/truetype/liberation/LiberationSans-Bold.ttf',
      '/usr/share/fonts/truetype/freefont/FreeSansBold.ttf']
_R = ['/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf',
      '/usr/share/fonts/truetype/liberation/LiberationSans-Regular.ttf',
      '/usr/share/fonts/truetype/freefont/FreeSans.ttf']

def fnt(paths, size):
    for p in paths:
        try: return ImageFont.truetype(p, size)
        except: pass
    try: return ImageFont.load_default(size=size)
    except: return ImageFont.load_default()

# ── GEMINI TEXT (script generation) ──────────────────────────────────────────
def gemini_text(prompt, api_key):
    url = GEMINI_TEXT_URL + '?key=' + api_key
    body = {
        'contents': [{'parts': [{'text': prompt}]}],
        'generationConfig': {'temperature': 0.8, 'maxOutputTokens': 8192}
    }
    r = requests.post(url, json=body, timeout=60)
    r.raise_for_status()
    data = r.json()
    return data['candidates'][0]['content']['parts'][0]['text']

# ── GEMINI IMAGEN (page generation) ──────────────────────────────────────────
def gemini_imagen_call(url, prompt, api_key, width, height):
    """Single Gemini image model call with retry on 429."""
    body = {
        'contents': [{'parts': [{'text': prompt}]}],
        'generationConfig': {
            'responseModalities': ['IMAGE', 'TEXT'],
            'temperature': 1.0
        }
    }
    full_url = url + '?key=' + api_key
    for attempt in range(2):
        try:
            r = requests.post(full_url, json=body, timeout=120)
            if r.status_code == 200:
                data = r.json()
                candidates = data.get('candidates', [])
                if candidates:
                    parts = candidates[0].get('content', {}).get('parts', [])
                    for part in parts:
                        if 'inlineData' in part:
                            img_bytes = base64.b64decode(part['inlineData']['data'])
                            img = Image.open(BytesIO(img_bytes)).convert('RGB')
                            return img.resize((width, height), Image.LANCZOS)
                print('Gemini image: no image in response from ' + url.split('/models/')[1].split(':')[0])
                return None
            elif r.status_code == 429:
                print('Gemini image 429 (' + url.split('/models/')[1].split(':')[0] + '), waiting ' + str(20*(attempt+1)) + 's...')
                time.sleep(20 * (attempt + 1))
            else:
                print('Gemini image error ' + str(r.status_code) + ' (' + url.split('/models/')[1].split(':')[0] + '): ' + r.text[:200])
                return None
        except Exception as e:
            print('Gemini image exception: ' + str(e)[:100])
            return None
    return None

def gemini_imagen(prompt, api_key, width=1024, height=1408):
    """Try Gemini image models from best to fastest: 3-pro → 3.1-flash → 2.5-flash."""
    # Model 1: Gemini 3 Pro Image (best quality)
    print('  Trying Gemini 3 Pro Image...')
    img = gemini_imagen_call(GEMINI_IMAGE_URL, prompt, api_key, width, height)
    if img:
        print('  Gemini 3 Pro Image SUCCESS')
        return img
    # Model 2: Gemini 3.1 Flash Image (fast + good quality)
    print('  Trying Gemini 3.1 Flash Image...')
    img = gemini_imagen_call(GEMINI_IMAGE_URL2, prompt, api_key, width, height)
    if img:
        print('  Gemini 3.1 Flash Image SUCCESS')
        return img
    # Model 3: Gemini 2.5 Flash Image (fastest)
    print('  Trying Gemini 2.5 Flash Image...')
    img = gemini_imagen_call(GEMINI_IMAGE_URL3, prompt, api_key, width, height)
    if img:
        print('  Gemini 2.5 Flash Image SUCCESS')
        return img
    return None

# ── POLLINATIONS FALLBACK ─────────────────────────────────────────────────────
def pollinations(prompt, seed, retries=3):
    enc = urllib.parse.quote(prompt[:600])
    url = ('https://image.pollinations.ai/prompt/' + enc +
           '?width=1024&height=1408&nologo=true&model=flux&seed=' + str(seed) + '&enhance=true')
    for attempt in range(retries):
        try:
            if attempt > 0: time.sleep(5 * attempt)
            r = requests.get(url, timeout=120)
            if r.status_code == 200 and len(r.content) > 8000:
                img = Image.open(BytesIO(r.content)).convert('RGB')
                return img.resize((1024, 1408), Image.LANCZOS)
            print('Pollinations status ' + str(r.status_code) + ' attempt ' + str(attempt))
        except Exception as e:
            print('Pollinations err: ' + str(e)[:60])
    return None

def placeholder(page_num, title):
    img = Image.new('RGB', (1024, 1408), (30, 30, 42))
    d = ImageDraw.Draw(img)
    d.rectangle([6,6,1017,1401], outline=(90,90,110), width=4)
    d.text((512, 704), title + ' — page ' + str(page_num),
           font=fnt(_B, 30), fill=(130,130,150), anchor='mm')
    d.text((512, 754), 'Image generation unavailable', font=fnt(_R, 22), fill=(90,90,100), anchor='mm')
    return img

# ── PAGE HEADER / FOOTER ──────────────────────────────────────────────────────
def add_header_footer(art_img, page_num, total, title):
    PW, PH = 1240, 1754
    HH, FH = 52, 36
    page = Image.new('RGB', (PW, PH), (255, 255, 255))
    d = ImageDraw.Draw(page)
    d.rectangle([0, 0, PW, HH], fill=(10, 10, 10))
    d.text((PW//2, HH//2), title.upper()[:60], font=fnt(_B, 24), fill=(255,255,255), anchor='mm')
    fy = PH - FH
    d.rectangle([0, fy, PW, PH], fill=(215, 215, 215))
    d.text((PW//2, fy + FH//2), str(page_num) + ' / ' + str(total), font=fnt(_R, 18), fill=(55,55,55), anchor='mm')
    cx0, cy0, cx1, cy1 = 8, HH+6, PW-8, fy-6
    art_scaled = art_img.resize((cx1-cx0, cy1-cy0), Image.LANCZOS)
    page.paste(art_scaled, (cx0, cy0))
    return page

# ── SCRIPT GENERATION ─────────────────────────────────────────────────────────
def generate_script(topic, num_pages, style, language, api_key):
    lang = 'in Spanish' if language == 'es' else 'in English'
    prompt = (
        'You are a professional comic book writer. '
        'Create a complete ' + str(num_pages) + '-page ' + style + ' comic script ' + lang + ' about: ' + str(topic) + '.\n'
        'Each page has exactly 4 panels.\n'
        'Define 1-3 main characters with detailed visual descriptions.\n'
        'Keep each dialogue text under 12 words.\n'
        'Each image_prompt describes the FULL PAGE as a comic layout (not individual panels).\n'
        'The page_image_prompt must describe:\n'
        '  - A 4-panel comic book PAGE layout\n'
        '  - Each panel scene, action and characters\n'
        '  - The speech bubbles with exact text inside them\n'
        '  - The panel borders/gutters\n'
        'Return ONLY valid JSON (no markdown, no extra text):\n'
        '{"title":"string","genre":"string","story_summary":"string",'
        '"characters":[{"name":"string","description":"detailed visual appearance"}],'
        '"pages":[{"page_number":1,'
        '"page_image_prompt":"Full-page comic layout prompt in English for AI image generation. Describe 4 panels arranged as a comic page, each panel scene, all speech bubbles with text visible inside white oval bubbles, panel dividers, comic book art style",'
        '"panels":[{"panel_number":1,"scene":"string","action":"string",'
        '"dialogue":[{"character":"string","text":"short dialogue"}],'
        '"mood":"string"}]}]}'
    )
    try:
        raw = gemini_text(prompt, api_key)
        m = re.search(r'\{[\s\S]*\}', raw)
        if m:
            data = json.loads(m.group())
            if 'pages' in data and data['pages']:
                print('Script generated by Gemini: ' + str(len(data['pages'])) + ' pages')
                return data
    except Exception as e:
        print('Gemini script error: ' + str(e)[:120])
    return fallback_script(topic, num_pages, style, language)

def fallback_script(topic, num_pages, style, language):
    DLGS = [
        [{'character':'Hero','text':'This ends now! No more running!'}],
        [{'character':'Hero','text':'I will protect them all!'},
         {'character':'Villain','text':'Your courage is useless here!'}],
        [{'character':'Ally','text':'Watch out, they surround us!'},
         {'character':'Hero','text':'Stay close. We fight together.'}],
        [{'character':'Hero','text':'Is this the power they warned us about?'}],
    ]
    pages = []
    for p in range(1, num_pages+1):
        panels = []
        for i in range(1, 5):
            dlg = DLGS[(p*4+i) % len(DLGS)]
            panels.append({'panel_number':i,'scene':'Scene','action':'Action','dialogue':dlg,'mood':'intense'})
        # Build a descriptive page prompt
        pg_prompt = (
            style + ' comic book full page, 4 panels layout, page ' + str(p) + ' of the story about ' + str(topic) +
            '. Top-left panel: ' + str(panels[0]['action']) +
            '. Top-right panel: ' + str(panels[1]['action']) +
            '. Bottom-left panel: ' + str(panels[2]['action']) +
            '. Bottom-right panel: ' + str(panels[3]['action']) +
            '. Each panel has white speech bubbles with visible text. Strong black panel borders. '
            'Professional comic book art, dramatic lighting, vivid colors.'
        )
        pages.append({'page_number':p,'page_image_prompt':pg_prompt,'panels':panels})
    return {'title':str(topic),'genre':'action','story_summary':str(topic),
            'characters':[{'name':'Hero','description':'warrior in red gold armor, black hair, strong build'},
                          {'name':'Villain','description':'dark robe, silver mask, glowing purple eyes'}],
            'pages':pages}

# ── PAGE IMAGE PROMPT BUILDER ─────────────────────────────────────────────────
def build_page_prompt(page_data, script, style):
    """Build a rich, detailed prompt for Gemini 3 Pro image generation."""
    title     = script.get('title', 'Comic')
    genre     = script.get('genre', 'adventure')
    chars     = script.get('characters', [])
    panels    = page_data.get('panels', [])[:4]
    page_num  = page_data.get('page_number', 1)
    total_p   = script.get('total_pages', 1)

    # Character descriptions for consistency
    char_descs = []
    for c in chars[:3]:
        char_descs.append(c['name'] + ': ' + c.get('description','') + (', wears ' + c.get('costume','') if c.get('costume') else ''))
    char_ref = '. '.join(char_descs) if char_descs else ''

    # Build per-panel descriptions with full dialogue
    panel_parts = []
    num_panels = len(panels)
    for i, p in enumerate(panels):
        scene   = str(p.get('scene', ''))
        action  = str(p.get('action', ''))
        caption = str(p.get('caption', ''))
        dialogues = p.get('dialogue', [])
        dlg_list = []
        for d in dialogues:
            dlg_list.append(d['character'] + ' says: "' + d['text'] + '"')
        dlg_str = '; '.join(dlg_list) if dlg_list else 'no dialogue'
        cap_str = (' Narrator caption box: "' + caption + '"') if caption else ''
        panel_parts.append(
            'PANEL ' + str(i+1) + ': Scene=' + scene +
            '. Action=' + action + '. ' + dlg_str + cap_str
        )

    # Choose layout based on panel count
    if num_panels == 1:
        layout_desc = 'ONE large full-page panel'
    elif num_panels == 2:
        layout_desc = 'TWO panels: large top panel, large bottom panel, separated by thick black gutter'
    elif num_panels == 3:
        layout_desc = 'THREE panels: wide panoramic top panel, two equal panels side-by-side at bottom, separated by thick black gutters'
    else:
        layout_desc = 'FOUR panels: large wide panel at top, two medium square panels in middle row, one wide panel at bottom, all separated by thick black gutters (8px)'

    panels_text = ' | '.join(panel_parts)

    prompt = (
        'CREATE A PROFESSIONAL ' + style.upper() + ' COMIC BOOK PAGE. '
        'Comic: "' + title + '" (' + genre + '). Page ' + str(page_num) + ' of ' + str(total_p) + '. '
        'LAYOUT: ' + layout_desc + '. '
        'PANELS: ' + panels_text + '. '
        'SPEECH BUBBLES: White rounded oval bubbles with clear BLACK TEXT for every dialogue line. '
        'Bubbles sized to fit text (max 25% of panel area), tail pointing to speaker. '
        'Text must be legible. Do NOT overlap speech bubbles with main action. '
        'STYLE REQUIREMENTS: '
        + style + ' art style. Bold black panel borders (8-10px). White gutter between panels. '
        'Dynamic compositions. Dramatic lighting and shadows. Vivid expressive colors. '
        'Professional comic book quality. NO watermarks, NO logos, NO page numbers in image. '
        + (('CHARACTERS (keep consistent): ' + char_ref + '. ') if char_ref else '') +
        'Make every panel visually compelling and clearly readable as a comic book.'
    )
    return prompt[:1800]

# ── GENERATE ONE PAGE ─────────────────────────────────────────────────────────

def draw_speech_bubbles(img, panels):
    """Draw readable speech bubbles with Pillow over the comic page."""
    draw = ImageDraw.Draw(img)
    W, H = img.size
    
    # Calculate panel zones (4 panels: top-wide, mid-left, mid-right, bottom-wide)
    panel_zones = [
        (0, 0, W, H//2),           # Panel 1 top half
        (0, H//2, W//2, H),        # Panel 2 bottom-left
        (W//2, H//2, W, H),        # Panel 3 bottom-right
    ]
    if len(panels) >= 4:
        panel_zones = [
            (0, 0, W, H//3),           # Panel 1 top strip
            (0, H//3, W//2, 2*H//3),   # Panel 2 mid-left
            (W//2, H//3, W, 2*H//3),   # Panel 3 mid-right
            (0, 2*H//3, W, H),         # Panel 4 bottom strip
        ]
    
    font_b = fnt(_B, 22)
    font_r = fnt(_R, 20)
    
    for i, panel in enumerate(panels[:len(panel_zones)]):
        if i >= len(panel_zones):
            break
        px0, py0, px1, py1 = panel_zones[i]
        dialogues = panel.get('dialogue', [])
        caption   = panel.get('caption', '')
        
        bubble_y = py1 - 15  # Start from bottom of panel
        
        # Draw dialogues (from bottom up)
        for j, dlg in enumerate(reversed(dialogues[-2:])):  # max 2 bubbles per panel
            char = dlg.get('character', '')
            text = dlg.get('text', '')
            if not text:
                continue
            
            # Wrap text
            max_chars = 25
            words = text.split()
            wrapped_lines = []
            line = ''
            for w in words:
                if len(line) + len(w) + 1 <= max_chars:
                    line = (line + ' ' + w).strip()
                else:
                    if line: wrapped_lines.append(line)
                    line = w
            if line: wrapped_lines.append(line)
            wrapped_lines = wrapped_lines[:3]  # max 3 lines
            
            # Calculate bubble dimensions
            line_h = 26
            text_h = len(wrapped_lines) * line_h + 8
            char_w = max(len(l) for l in wrapped_lines) * 13 + 20
            char_w = min(char_w, (px1 - px0) - 20)
            bub_h  = text_h + 30  # space for char name
            
            # Position bubble
            bub_x0 = px0 + 15
            bub_y1 = bubble_y
            bub_y0 = bub_y1 - bub_h
            bub_x1 = bub_x0 + char_w
            
            if bub_y0 < py0 + 10:
                break  # No space
            
            # Draw white rounded bubble
            draw.rounded_rectangle([bub_x0, bub_y0, bub_x1, bub_y1], radius=12,
                                   fill=(255,255,255), outline=(20,20,20), width=2)
            
            # Character name in bold
            draw.text((bub_x0+10, bub_y0+6), char[:12] + ':', font=font_b, fill=(20,20,80))
            
            # Dialogue text
            for k, wl in enumerate(wrapped_lines):
                draw.text((bub_x0+10, bub_y0+28+k*line_h), wl, font=font_r, fill=(20,20,20))
            
            bubble_y = bub_y0 - 8  # Stack bubbles upward
        
        # Draw caption box at top of panel if exists
        if caption:
            cap_text = caption[:50]
            cap_w = len(cap_text) * 10 + 20
            cap_w = min(cap_w, px1 - px0 - 10)
            draw.rectangle([px0+5, py0+5, px0+cap_w, py0+30], fill=(255,255,220), outline=(100,100,0), width=1)
            draw.text((px0+10, py0+8), cap_text, font=fnt(_R, 16), fill=(60,40,0))
    
    return img

def generate_page(page_data, script, style, title, page_num, total, api_key):
    prompt = build_page_prompt(page_data, script, style)
    seed = page_num * 137

    print('Page ' + str(page_num) + '/' + str(total) + ' generating...')

    # Try Gemini image models (3-pro → 3.1-flash → 2.5-flash → Pollinations)
    art = None
    if api_key:
        art = gemini_imagen(prompt, api_key)

    # Fallback to Pollinations
    if art is None:
        print('  Trying Pollinations...')
        art = pollinations(prompt, seed)

    # Final fallback
    if art is None:
        print('  Using placeholder')
        art = placeholder(page_num, title)

    # Add readable speech bubbles with Pillow (works even with Pollinations)
    panels = page_data.get('panels', [])
    art = draw_speech_bubbles(art.copy(), panels)
    
    final = add_header_footer(art, page_num, total, title)
    print('Page ' + str(page_num) + ' done')
    return final

# ── MAIN ──────────────────────────────────────────────────────────────────────
def generate_comic(topic, num_pages, style, language, user_gemini_key):
    global GEMINI_API_KEY
    api_key = user_gemini_key.strip() if user_gemini_key and user_gemini_key.strip() else GEMINI_API_KEY
    num_pages = max(1, min(int(num_pages), 32))
    print('=== Generating ' + str(num_pages) + 'p comic: ' + str(topic) + ' ===')

    script = generate_script(topic, num_pages, style, language, api_key)
    title  = script.get('title', str(topic))
    pages_data = script.get('pages', [])

    while len(pages_data) < num_pages:
        fb = fallback_script(topic, 1, style, language)
        pg = fb['pages'][0]; pg['page_number'] = len(pages_data)+1
        pages_data.append(pg)

    page_images = []
    for pg in pages_data[:num_pages]:
        img = generate_page(pg, script, style, title, pg['page_number'], num_pages, api_key)
        page_images.append(img)
        time.sleep(2)  # Rate limit respect

    return script, page_images

def generate_comic_api(topic, num_pages=4, style='manga', language='en', gemini_key=''):
    if not str(topic).strip(): return {'error':'Topic required','success':False}
    script, pages = generate_comic(str(topic), num_pages, style, language, gemini_key)
    out = []
    for i, pg in enumerate(pages):
        buf = BytesIO(); pg.save(buf, format='JPEG', quality=92)
        out.append({'page_number':i+1,'image_base64':base64.b64encode(buf.getvalue()).decode(),'size':'1240x1754'})
    return {'success':True,'title':script.get('title',''),'story_summary':script.get('story_summary',''),
            'characters':script.get('characters',[]),'total_pages':len(pages),'pages':out}

def generate_comic_ui(topic, num_pages, style, language, gemini_key):
    if not topic.strip(): return '','',[],'{}' 
    script, pages = generate_comic(topic, num_pages, style, language, gemini_key)
    info = json.dumps({'title':script.get('title',''),'summary':script.get('story_summary',''),
                       'characters':script.get('characters',[]),'total_pages':len(pages),
                       'engine':'Gemini 3 Pro Image → 3.1 Flash → 2.5 Flash → Pollinations fallback'},
                      ensure_ascii=False, indent=2)
    return script.get('title',topic), script.get('story_summary',''), pages, info

STYLES=['manga','superhero comic','watercolor comic','noir comic','cartoon',
        'fantasy illustration','sci-fi comic','childrens book illustration']

with gr.Blocks(title='Comic AI Generator') as demo:
    gr.Markdown('# Comic AI Generator')
    gr.Markdown('Powered by **Gemini 3 Pro Image** — each page is one AI-generated full comic page (1240×1754 px) with panels, borders and speech bubbles drawn by AI.')
    with gr.Tabs():
        with gr.Tab('Generate Comic'):
            with gr.Row():
                with gr.Column(scale=1):
                    topic_in   = gr.Textbox(label='Story / Topic', lines=3,
                                    placeholder='A pirate captain who discovers a mythical treasure island...')
                    pages_in   = gr.Slider(1, 32, value=2, step=1, label='Number of Pages')
                    style_in   = gr.Dropdown(choices=STYLES, value='manga', label='Art Style')
                    lang_in    = gr.Radio(choices=[('English','en'),('Spanish','es')], value='en', label='Language')
                    key_in     = gr.Textbox(label='Gemini API Key (optional — uses server key if empty)',
                                    placeholder='AIza...', type='password')
                    gen_btn    = gr.Button('Generate Comic', variant='primary', size='lg')
                    gr.Markdown('_1 page ≈ 20s · 4 pages ≈ 1.5 min · 32 pages ≈ 12 min_')
                with gr.Column(scale=2):
                    title_out   = gr.Textbox(label='Title', interactive=False)
                    summary_out = gr.Textbox(label='Story Summary', interactive=False, lines=3)
                    gallery_out = gr.Gallery(label='Comic Pages', columns=1, rows=4, height='auto')
                    with gr.Accordion('Script / API info (JSON)', open=False):
                        json_out = gr.Code(language='json')
            gen_btn.click(fn=generate_comic_ui,
                          inputs=[topic_in, pages_in, style_in, lang_in, key_in],
                          outputs=[title_out, summary_out, gallery_out, json_out])
        with gr.Tab('API'):
            gr.Markdown(
                '## REST API\n\n'
                '**POST** `https://bukbuk-comic-ai-generator.hf.space/call/generate_comic_api`\n\n'
                '```json\n{"data":["topic", num_pages, "style", "language", "gemini_api_key"]}\n```\n\n'
                'Returns `pages[]` — each is a **1240×1754 JPEG** full comic page generated by Gemini Imagen 3.'
            )
    gr.Interface(fn=generate_comic_api,
        inputs=[gr.Textbox(label='topic'),gr.Number(label='num_pages',value=4),
                gr.Textbox(label='style',value='manga'),gr.Textbox(label='language',value='en'),
                gr.Textbox(label='gemini_key',value='')],
        outputs=gr.JSON(label='comic_data'), api_name='generate_comic_api',
        title='Comic API', description='Returns full 1240x1754 comic pages as base64 JPEG. Gemini 2.5 Flash Image powered.')

if __name__ == '__main__':
    demo.launch(share=False)