import gradio as gr
import json
from typing import Dict, List
import os

class DatasetViewer:
    def __init__(self, json_path: str):
        with open(json_path, 'r', encoding='utf-8') as f:
            self.data = json.load(f)
        self.current_index = 0
        self.total_items = len(self.data)
    
    def get_current_item(self) -> tuple:
        item = self.data[self.current_index]
        return (
            item.get('corrected_description', ''),  # Left Text
            item.get('pddl_domain_processed', ''), # Right Text
            f"{self.current_index + 1}. {item.get('file_name', '')}"  # Current Location
        )
    
    def next_item(self) -> tuple:
        self.current_index = (self.current_index + 1) % self.total_items
        return self.get_current_item()
    
    def prev_item(self) -> tuple:
        self.current_index = (self.current_index - 1) % self.total_items
        return self.get_current_item()
    
    def jump_to_index(self, selected_key: str) -> tuple:
        try:
            # Extract index from the selected text
            index = int(selected_key.split('.')[0]) - 1
            if 0 <= index < self.total_items:
                self.current_index = index
        except:
            pass
        return self.get_current_item()
    
    def get_choices(self, search_term: str = '') -> List[str]:
        choices = []
        for idx, item in enumerate(self.data):
            preview = f"{item.get('file_name', '')}"
            choice = f"{idx + 1}. {preview}"
            if search_term.lower() in choice.lower():
                choices.append(choice)
        return choices

def create_ui():
    # Initialize dataset viewer
    viewer = DatasetViewer('our_benchmark_modified.json')
    
    with gr.Blocks(css="""
        .container { margin: 15px; }
        .text-display { min-height: 200px; }
        .navigation { text-align: center; margin: 10px; }
        .hf-yellow-btn {
            background: #FFA726 !important; /* Hugging Face 黄色 */
            border: none !important;
        }
        .hf-yellow-btn:hover {
            background: #FF9800 !important; /* 鼠标悬停时的颜色 */
            box-shadow: 0 0 0 0.2rem rgba(255, 167, 38, 0.5) !important;
        }
    """) as demo:
        gr.HTML("<h1 style='text-align: center; margin-bottom: 1em;'>Data Viewer</h1>")
        
        with gr.Row():
            # Navigation control
            prev_btn = gr.Button(
                "Previous\nPage", 
                scale=1,
                min_width=50,
                size="lg",
                variant="primary",
                elem_classes="hf-yellow-btn"
            )
            current_position = gr.Dropdown(
                choices=viewer.get_choices(),
                label="File Name",
                scale=2,
                allow_custom_value=True
            )
            next_btn = gr.Button(
                "Next\nPage", 
                scale=1,
                min_width=50,
                size="lg",
                variant="primary",
                elem_classes="hf-yellow-btn"
            )
        
        with gr.Row():
            # Text display area
            # left_text = gr.TextArea(label="Description", interactive=False, elem_classes=["text-display"])
            left_text = gr.Code(
                label="Description", 
                language="markdown",
                interactive=False, 
                elem_classes=["text-display"],
                wrap_lines=True
            )
            right_text = gr.Code(
                label="PDDL Domain", 
                language="typescript",
                interactive=False, 
                elem_classes=["text-display"],
                # wrap_lines=True
            )
        
        # Initialize display
        initial_data = viewer.get_current_item()
        left_text.value = initial_data[0]
        right_text.value = initial_data[1]
        current_position.value = initial_data[2]
        
        # Event Handling
        prev_btn.click(
            fn=viewer.prev_item,
            outputs=[left_text, right_text, current_position]
        )
        
        next_btn.click(
            fn=viewer.next_item,
            outputs=[left_text, right_text, current_position]
        )
        
        current_position.change(
            fn=viewer.jump_to_index,
            inputs=[current_position],
            outputs=[left_text, right_text, current_position]
        )
        
        # Search function
        current_position.input(
            fn=viewer.get_choices,
            inputs=[current_position],
            outputs=[current_position]
        )
    
    return demo

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