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

# Load segments JSON
def load_segments():
    with open("segments_full.json", "r", encoding="utf-8") as f:
        segments = json.load(f)
    return segments

# 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))
        else:
            images.append(Image.new("RGB", (256, 144), color=(200, 200, 200)))
    return images


def segment_display(segment):
    """
    Format one row for Gradio display:
    - Truncate and wrap description
    - Load and resize keyframe images to thumbnail size
    """
    seg_id = segment.get("segment_id")
    description = segment.get("description", "")
    short_desc = description if len(description) <= 120 else description[:117] + "..."

    # Wrap long text with newlines every ~40 chars
    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]

    for i in range(1, 4):
        img_path = f"keyframes/segment_{seg_id}_v{i}.png"
        if os.path.exists(img_path):
            img = Image.open(img_path).resize((128, 72))
        else:
            # Placeholder if not found
            img = Image.new("RGB", (128, 72), color=(200, 200, 200))
            draw = ImageDraw.Draw(img)
            draw.text((10, 30), "No Image", fill=(0, 0, 0))
        row.append(img)

    return row

# Pagination logic
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

# Function to return downloadable JSON file
def download_prompts():
    if os.path.exists("all_prompts_output.json"):
        return "all_prompts_output.json"
    else:
        return None

# Gradio interface
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

        prev_btn = gr.Button("⬅ Prev")
        next_btn = gr.Button("Next ➡")

        table.change(fn=lambda: None, inputs=[], outputs=[])
        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("### 📥 下载所有生成的图像提示词")
        with gr.Row():
            download_btn = gr.Button("📥 导出 prompts JSON")
            download_output = gr.File(label="下载文件")
            download_btn.click(fn=download_prompts, outputs=download_output)

    return demo

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