# --- app.py ---
import json
import math
import gradio as gr
from PIL import Image, ImageDraw
import os
import shutil
from utils.keyframe_utils import generate_keyframe_prompt, generate_all_keyframe_images

# Load segments JSON
def load_segments():
    with open("segments_full.json", "r", encoding="utf-8") as f:
        segments = json.load(f)
    return segments[:5]  # return only the first 5 segments for initial testing

# Load pre-generated keyframe images (3 per segment)
def get_keyframe_images(segment_id):
    images = []
    for i in range(1, 4):
        img_path = f"keyframes/segment_{segment_id}_v{i}.png"
        if os.path.exists(img_path):
            images.append(Image.open(img_path).resize((128, 72)))
        else:
            img = Image.new("RGB", (128, 72), color=(200, 200, 200))
            draw = ImageDraw.Draw(img)
            draw.text((10, 30), "No Image", fill=(0, 0, 0))
            images.append(img)
    return images

def segment_display(segment):
    seg_id = segment.get("segment_id")
    description = segment.get("description", "")
    short_desc = description if len(description) <= 120 else description[:117] + "..."

    def wrap_text(text, width=40):
        return "\n".join([text[i:i+width] for i in range(0, len(text), width)])

    display_desc = wrap_text(short_desc)
    row = [f"Segment {seg_id}", display_desc]
    row.extend(get_keyframe_images(seg_id))
    return row

def paginate_segments(page=1, page_size=15):
    segments = load_segments()
    total = len(segments)
    max_page = math.ceil(total / page_size)
    start = (page - 1) * page_size
    end = start + page_size
    subset = segments[start:end]

    headers = ["Segment ID", "Description", "Candidate 1", "Candidate 2", "Candidate 3"]
    data = [segment_display(seg) for seg in subset]
    return headers, data, max_page

def zip_keyframe_folder():
    zip_name = "keyframes.zip"
    if os.path.exists(zip_name):
        os.remove(zip_name)
    shutil.make_archive("keyframes", 'zip', "keyframes")
    return zip_name

def build_interface():
    with gr.Blocks() as demo:
        gr.Markdown("## 🎬 Keyframe Candidate Viewer")
        page_state = gr.State(1)
        table = gr.Dataframe(
            headers=["Segment ID", "Description", "Candidate 1", "Candidate 2", "Candidate 3"],
            datatype=["str", "str", "image", "image", "image"],
            row_count=15
        )
        total_pages_text = gr.Textbox(label="Page Info", interactive=False)

        def update(page):
            headers, rows, max_page = paginate_segments(page)
            return gr.update(headers=headers, value=rows), f"Page {page} of {max_page}", page

        def generate_and_refresh(_):
            segments = load_segments()
            generate_all_keyframe_images(segments)
            return update(1)

        with gr.Row():
            prev_btn = gr.Button("⬅", size="sm")
            next_btn = gr.Button("➡", size="sm")

        prev_btn.click(fn=lambda p: max(1, p - 1), inputs=page_state, outputs=page_state).then(update, inputs=page_state, outputs=[table, total_pages_text, page_state])
        next_btn.click(fn=lambda p: p + 1, inputs=page_state, outputs=page_state).then(update, inputs=page_state, outputs=[table, total_pages_text, page_state])

        demo.load(fn=update, inputs=page_state, outputs=[table, total_pages_text, page_state])

        gr.Markdown("### 🖼️ 生成所有分镜的关键帧图像")
        generate_btn = gr.Button("🛠️ 一键生成图像")
        generate_btn.click(fn=generate_and_refresh, inputs=[page_state], outputs=[table, total_pages_text, page_state])

        gr.Markdown("### 📦 打包下载 keyframes 文件夹（包含图片和 prompts.json）")
        zip_btn = gr.Button("📁 生成 ZIP")
        zip_file_output = gr.File(label="下载 ZIP")
        zip_btn.click(fn=zip_keyframe_folder, outputs=zip_file_output)

    return demo

if __name__ == "__main__":
    demo = build_interface()
    demo.launch()