import gradio as gr
import time
from utils.story_management import (
    generate_direct_comic,
    extract_comic_scenes,
    load_narration_from_file
)
from config import IMAGE_STYLES, IMAGE_STYLE_INFO, AGE_GROUPS
from datetime import datetime

def log_execution(func):
    def wrapper(*args, **kwargs):
        start_time = time.time()
        start_str = datetime.fromtimestamp(start_time).strftime('%Y-%m-%d %H:%M:%S')
        
        result = func(*args, **kwargs)
        
        end_time = time.time()
        end_str = datetime.fromtimestamp(end_time).strftime('%Y-%m-%d %H:%M:%S')
        duration = end_time - start_time
        
        # Write to file (works in Colab)
        with open('content/logs.txt', 'a') as f:
            f.write(f"{func.__name__}, start: {start_str}, end: {end_str}, duration: {duration:.4f}s\n")
        
        # Also print to see output immediately
        print(f"{func.__name__}, start: {start_str}, end: {end_str}, duration: {duration:.4f}s")
        
        return result
    return wrapper
@log_execution
def create_story_interface(demo: gr.Blocks) -> gr.Blocks:
    """Create the main story interface with comic generation functionality.

    This function initializes the primary UI interface for the comic generation system,
    setting up the main tab structure and components.

    Args:
        demo (gr.Blocks): The Gradio Blocks instance to build the interface on

    Returns:
        gr.Blocks: The configured Gradio interface with all components initialized
    """

    create_quick_comic_tab()

    return demo
def create_quick_comic_tab() -> None:
    """Create a simple tab for direct prompt-to-image comic generation.

    Sets up the main comic generation interface with the following components:
    - Story prompt input field
    - AI prompt enhancement option
    - Visual style selection
    - Number of scenes selector
    - Generation controls
    - Image display area
    - Scene navigation system

    The interface allows users to:
    1. Input their story description
    2. Configure generation parameters
    3. Generate a multi-panel comic
    4. View and navigate through individual scenes
    """

    with gr.Column():
    
        gr.Markdown("Welcome to Hekaya ")

        with gr.Row():
            with gr.Column(scale=3):
                user_prompt = gr.Textbox(
                    label="What Hekaya story would you like to visualize?",
                    placeholder="Describe your story with main characters and settings... (e.g., 'A young wizard learning magic in an ancient castle')",
                    lines=4
                )

                enrich_prompt = gr.Checkbox(
                    label="Enhance prompt with AI for coherence",
                    value=True,
                    info="Use AI to add just enough detail and coherence for consistent visual storytelling across all scenes"
                )

            with gr.Column(scale=1):
                comic_style = gr.Dropdown(
                    label="Visual Style",
                    choices=IMAGE_STYLES,
                    value="Comic Book Style"
                )

                style_description = gr.Markdown(
                    value=f"*{IMAGE_STYLE_INFO['Comic Book Style']}*",
                    label="Style Description"
                )

                age_group = gr.Dropdown(
                    label="Target Age Group",
                    choices=AGE_GROUPS,
                    value="9-12 (Pre-teen)",
                    info="Select the audience age group. Narration language, detail, and length will adapt automatically."
                )

                image_quality = gr.Dropdown(
                    label="Image Quality",
                    choices=["Low", "Medium", "High"],
                    value="Low",
                    info="Select the quality level for generated images. Higher quality may take longer to generate."
                )

        generate_btn = gr.Button("Generate Hekaya Story", variant="primary")
        status_display = gr.Markdown("")

        with gr.Row():
            with gr.Column(scale=2):
                comic_image = gr.Image(label="Generated Hekaya Story", type="filepath")

            with gr.Column(scale=1, elem_id="save_info_container"):
            
                gr.Markdown("Your generated story images are automatically saved locally.")
                save_path_display = gr.Markdown("", elem_id="save_path_info")

        narration_display = gr.Markdown(
        
            visible=True,
            elem_id="story_narration",
            elem_classes=["story-narration-box"]
        )

        with gr.Column(visible=False) as scene_viewer_container:
            
            gr.Markdown("Use the navigation buttons to view each upscaled scene individually.")

            with gr.Row(equal_height=True):
                prev_scene_btn = gr.Button("← Previous Scene", variant="secondary")
                scene_counter = gr.Markdown("Scene 1 of 1", elem_id="scene_counter")
                next_scene_btn = gr.Button("Next Scene →", variant="secondary")

            scene_image = gr.Image(label="Current Scene", type="filepath", height=768)
            scene_caption_display = gr.Markdown("", elem_id="scene_caption", elem_classes=["scene-caption-box"])
            scene_save_path = gr.Markdown("", elem_id="scene_save_path_info")

        scene_info = gr.State([])
        current_scene_index = gr.State(0)


        def update_style_description(style: str) -> str:
            """Update the style description text when a new style is selected."""
            return f"*{IMAGE_STYLE_INFO[style]}*"

        def show_generating_message() -> str:
            """Display a loading message while story scenes are being generated."""
            return "🔄 Generating your story scenes... Please wait..."

        def generate_comic_with_length(user_prompt, comic_style, enrich_prompt, age_group, image_quality):
            """Wrapper that handles the fixed num_scenes value while passing the age group and image quality."""
            comic_image, save_path_display, status_display, narration = generate_direct_comic(
                user_prompt,
                comic_style,
                12,
                enrich_prompt,
                3,  
                age_group,
            )

            if narration and narration.strip():
                narration_formatted = f""
                narration_update = gr.update(visible=True, value=narration_formatted)
            else:
                narration_update = gr.update(visible=True, value="")

            return comic_image, save_path_display, status_display, narration_update

        def init_scene_viewer(comic_path: str | None) -> tuple:
            """Initialize the scene viewer with extracted scenes from the comic image."""
            if not comic_path:
                return [], 0, gr.update(visible=False), None, "", "Scene 0 of 0", "No story image generated"

            scene_data, save_message = extract_comic_scenes(comic_path, 0)

            if not scene_data:
                return [], 0, gr.update(visible=False), None, "", "Scene 0 of 0", "Failed to extract scenes"

            first_scene = scene_data[0]

            return (
                scene_data,
                0,
                gr.update(visible=True),
                first_scene['path'],
                first_scene['caption'],
                f"Scene 1 of {len(scene_data)}",
                save_message
            )

        def update_scene_display(scene_data: list, current_index: int) -> tuple:
            """Update the scene viewer display with the current scene."""
            if not scene_data:
                return None, "", "Scene 0 of 0"

            index = max(0, min(current_index, len(scene_data) - 1))
            scene = scene_data[index]

            return scene['path'], scene['caption'], f"Scene {index + 1} of {len(scene_data)}"

        def navigate_to_previous_scene(idx: int) -> int:
            """Navigate to the previous scene in the sequence."""
            return max(0, idx - 1)

        def navigate_to_next_scene(paths: list, idx: int) -> int:
            """Navigate to the next scene in the sequence."""
            return min(len(paths) - 1, idx + 1) if paths else 0


        comic_style.change(
            fn=update_style_description,
            inputs=[comic_style],
            outputs=[style_description]
        )

        generate_btn.click(
            fn=show_generating_message,
            inputs=None,
            outputs=status_display
        ).then(
            fn=generate_comic_with_length,
            inputs=[user_prompt, comic_style, enrich_prompt, age_group, image_quality],
            outputs=[comic_image, save_path_display, status_display, narration_display]
        ).then(
            fn=init_scene_viewer,
            inputs=[comic_image],
            outputs=[
                scene_info,
                current_scene_index,
                scene_viewer_container,
                scene_image,
                scene_caption_display,
                scene_counter,
                scene_save_path
            ]
        )

        prev_scene_btn.click(
            fn=navigate_to_previous_scene,
            inputs=[current_scene_index],
            outputs=[current_scene_index]
        ).then(
            fn=update_scene_display,
            inputs=[scene_info, current_scene_index],
            outputs=[scene_image, scene_caption_display, scene_counter]
        )

        next_scene_btn.click(
            fn=navigate_to_next_scene,
            inputs=[scene_info, current_scene_index],
            outputs=[current_scene_index]
        ).then(
            fn=update_scene_display,
            inputs=[scene_info, current_scene_index],
            outputs=[scene_image, scene_caption_display, scene_counter]
        )