"""
The app contains:
- a new UI for the OmniParser AI Agent.
- 
python app_new.py --windows_host_url localhost:8006 --omniparser_server_url localhost:8000
"""

import os
import io
import shutil
import mimetypes
from datetime import datetime
from enum import StrEnum
from functools import partial
from pathlib import Path
from typing import cast, List, Optional
import argparse
import gradio as gr
from anthropic import APIResponse
from anthropic.types import TextBlock
from anthropic.types.beta import BetaMessage, BetaTextBlock, BetaToolUseBlock
from anthropic.types.tool_use_block import ToolUseBlock
from loop import (
    APIProvider,
    sampling_loop_sync,
)
from tools import ToolResult
import requests
from requests.exceptions import RequestException
import base64

CONFIG_DIR = Path("~/.anthropic").expanduser()
API_KEY_FILE = CONFIG_DIR / "api_key"

INTRO_TEXT = '''
<div style="text-align: center; margin-bottom: 10px;">
    <h2>OmniParser AI Agent</h2>
    <p>Turn any vision-language model into an AI agent. We currently support <b>OpenAI (4o/o1/o3-mini), DeepSeek (R1), Qwen (2.5VL) or Anthropic Computer Use (Sonnet)</b>.</p>
    <p>Type a message and press send to start OmniTool. Press stop to pause, and press the trash icon in the chat to clear the message history.</p>
    <p>You can also upload files for analysis using the file upload section.</p>
</div>
'''

def parse_arguments():
    parser = argparse.ArgumentParser(description="Gradio App")
    parser.add_argument("--windows_host_url", type=str, default='localhost:8006')
    parser.add_argument("--omniparser_server_url", type=str, default="localhost:8000")
    parser.add_argument("--run_folder", type=str, default="./tmp/outputs")
    return parser.parse_args()
args = parse_arguments()

# Update upload folder from args if provided
RUN_FOLDER = Path(os.path.join(args.run_folder, datetime.now().strftime('%Y%m%d_%H%M')))
RUN_FOLDER.mkdir(parents=True, exist_ok=True)

class Sender(StrEnum):
    USER = "user"
    BOT = "assistant"
    TOOL = "tool"


def load_existing_files():
    """Load all existing files from the uploads folder"""
    files = []
    if RUN_FOLDER.exists():
        for file_path in RUN_FOLDER.iterdir():
            if file_path.is_file():
                files.append(str(file_path))
    return files

def setup_state(state):
    if "messages" not in state:
        state["messages"] = []
    if "model" not in state:
        state["model"] = "omniparser + gpt-4o-orchestrated"
    if "provider" not in state:
        state["provider"] = "openai"
    if "openai_api_key" not in state:  # Fetch API keys from environment variables
        state["openai_api_key"] = os.getenv("OPENAI_API_KEY", "")
    if "anthropic_api_key" not in state:
        state["anthropic_api_key"] = os.getenv("ANTHROPIC_API_KEY", "")
    if "api_key" not in state:
        state["api_key"] = ""
    if "auth_validated" not in state:
        state["auth_validated"] = False
    if "responses" not in state:
        state["responses"] = {}
    if "tools" not in state:
        state["tools"] = {}
    if "only_n_most_recent_images" not in state:
        state["only_n_most_recent_images"] = 2
    if 'chatbot_messages' not in state:
        state['chatbot_messages'] = []
    if 'stop' not in state:
        state['stop'] = False
    if 'uploaded_files' not in state:
        state['uploaded_files'] = []  # Start with an empty list instead of loading existing files

async def main(state):
    """Render loop for Gradio"""
    setup_state(state)
    return "Setup completed"

def validate_auth(provider: APIProvider, api_key: str | None):
    if provider == APIProvider.ANTHROPIC:
        if not api_key:
            return "Enter your Anthropic API key to continue."
    if provider == APIProvider.BEDROCK:
        import boto3

        if not boto3.Session().get_credentials():
            return "You must have AWS credentials set up to use the Bedrock API."
    if provider == APIProvider.VERTEX:
        import google.auth
        from google.auth.exceptions import DefaultCredentialsError

        if not os.environ.get("CLOUD_ML_REGION"):
            return "Set the CLOUD_ML_REGION environment variable to use the Vertex API."
        try:
            google.auth.default(scopes=["https://www.googleapis.com/auth/cloud-platform"])
        except DefaultCredentialsError:
            return "Your google cloud credentials are not set up correctly."

def load_from_storage(filename: str) -> str | None:
    """Load data from a file in the storage directory."""
    try:
        file_path = CONFIG_DIR / filename
        if file_path.exists():
            data = file_path.read_text().strip()
            if data:
                return data
    except Exception as e:
        print(f"Debug: Error loading {filename}: {e}")
    return None

def save_to_storage(filename: str, data: str) -> None:
    """Save data to a file in the storage directory."""
    try:
        CONFIG_DIR.mkdir(parents=True, exist_ok=True)
        file_path = CONFIG_DIR / filename
        file_path.write_text(data)
        # Ensure only user can read/write the file
        file_path.chmod(0o600)
    except Exception as e:
        print(f"Debug: Error saving {filename}: {e}")

def _api_response_callback(response: APIResponse[BetaMessage], response_state: dict):
    response_id = datetime.now().isoformat()
    response_state[response_id] = response

def _tool_output_callback(tool_output: ToolResult, tool_id: str, tool_state: dict):
    tool_state[tool_id] = tool_output

def chatbot_output_callback(message, chatbot_state, hide_images=False, sender="bot"):
    def _render_message(message: str | BetaTextBlock | BetaToolUseBlock | ToolResult, hide_images=False):
    
        print(f"_render_message: {str(message)[:100]}")
        
        if isinstance(message, str):
            return message
        
        is_tool_result = not isinstance(message, str) and (
            isinstance(message, ToolResult)
            or message.__class__.__name__ == "ToolResult"
        )
        if not message or (
            is_tool_result
            and hide_images
            and not hasattr(message, "error")
            and not hasattr(message, "output")
        ):  # return None if hide_images is True
            return
        # render tool result
        if is_tool_result:
            message = cast(ToolResult, message)
            if message.output:
                return message.output
            if message.error:
                return f"Error: {message.error}"
            if message.base64_image and not hide_images:
                # somehow can't display via gr.Image
                # image_data = base64.b64decode(message.base64_image)
                # return gr.Image(value=Image.open(io.BytesIO(image_data)))
                return f'<img src="data:image/png;base64,{message.base64_image}">'

        elif isinstance(message, BetaTextBlock) or isinstance(message, TextBlock):
            # Format reasoning text in a collapsible dropdown
            return f"Next step Reasoning: {message.text}"
            # reasoning_text = message.text
            # return f'''
            # <details>
            #     <summary><Current Step Reasoning (click to expand):</summary>
            #     <div style="padding: 10px; background-color: #f8f9fa; border-radius: 5px; margin-top: 5px;">
            #         <pre>{reasoning_text}</pre>
            #     </div>
            # </details>
            # '''
        elif isinstance(message, BetaToolUseBlock) or isinstance(message, ToolUseBlock):
            # return f"Next I will perform the following action: {message.input}"
            return None
        else:  
            return message

    def _truncate_string(s, max_length=500):
        """Truncate long strings for concise printing."""
        if isinstance(s, str) and len(s) > max_length:
            return s[:max_length] + "..."
        return s
    # processing Anthropic messages
    message = _render_message(message, hide_images)
    
    if sender == "bot":
        chatbot_state.append((None, message))
    else:
        chatbot_state.append((message, None))
    
    # Create a concise version of the chatbot state for printing
    concise_state = [(_truncate_string(user_msg), _truncate_string(bot_msg))
                        for user_msg, bot_msg in chatbot_state]
    # print(f"chatbot_output_callback chatbot_state: {concise_state} (truncated)")

def valid_params(user_input, state):
    """Validate all requirements and return a list of error messages."""
    errors = []
    
    for server_name, url in [('Windows Host', 'localhost:5000'), ('OmniParser Server', args.omniparser_server_url)]:
        try:
            url = f'http://{url}/probe'
            response = requests.get(url, timeout=3)
            if response.status_code != 200:
                errors.append(f"{server_name} is not responding")
        except RequestException as e:
            errors.append(f"{server_name} is not responding")
    
    if not state["api_key"].strip():
        errors.append("LLM API Key is not set")

    if not user_input:
        errors.append("no computer use request provided")
    
    return errors

def process_input(user_input, state):
    # Reset the stop flag
    if state["stop"]:
        state["stop"] = False

    errors = valid_params(user_input, state)
    if errors:
        raise gr.Error("Validation errors: " + ", ".join(errors))
    
    # Append the user message to state["messages"]
    state["messages"].append(
        {
            "role": Sender.USER,
            "content": [TextBlock(type="text", text=user_input)],
        }
    )

    # Append the user's message to chatbot_messages with None for the assistant's reply
    state['chatbot_messages'].append((user_input, None))
    yield state['chatbot_messages'], gr.update()  # Yield to update the chatbot UI with the user's message

    print("state")
    print(state)

    # Run sampling_loop_sync with the chatbot_output_callback
    for loop_msg in sampling_loop_sync(
        model=state["model"],
        provider=state["provider"],
        messages=state["messages"],
        output_callback=partial(chatbot_output_callback, chatbot_state=state['chatbot_messages'], hide_images=False),
        tool_output_callback=partial(_tool_output_callback, tool_state=state["tools"]),
        api_response_callback=partial(_api_response_callback, response_state=state["responses"]),
        api_key=state["api_key"],
        only_n_most_recent_images=state["only_n_most_recent_images"],
        max_tokens=16384,
        omniparser_url=args.omniparser_server_url,
        save_folder=str(RUN_FOLDER)
    ):  
        if loop_msg is None or state.get("stop"):
            # Detect and add new files to the state
            file_choices_update = detect_new_files(state)
            yield state['chatbot_messages'], file_choices_update
            print("End of task. Close the loop.")
            break
            
        yield state['chatbot_messages'], gr.update()  # Yield the updated chatbot_messages to update the chatbot UI
    
    # Final detection of new files
    file_choices_update = detect_new_files(state)
    yield state['chatbot_messages'], file_choices_update

def stop_app(state):
    state["stop"] = True
    return "App stopped"

def get_header_image_base64():
    try:
        # Get the absolute path to the image relative to this script
        script_dir = Path(__file__).parent
        image_path = script_dir.parent.parent / "imgs" / "header_bar_thin.png"
        
        with open(image_path, "rb") as image_file:
            encoded_string = base64.b64encode(image_file.read()).decode()
            return f'data:image/png;base64,{encoded_string}'
    except Exception as e:
        print(f"Failed to load header image: {e}")
        return None

def get_file_viewer_html(file_path=None):
    """Generate HTML to view a file based on its type"""
    if not file_path:
        # Return the VNC viewer iframe
        return f'<iframe src="http://{args.windows_host_url}/vnc.html?view_only=1&autoconnect=1&resize=scale" width="100%" height="580" allow="fullscreen"></iframe>'
    
    file_path = Path(file_path)
    if not file_path.exists():
        return f'<div class="error-message">File not found: {file_path.name}</div>'
    
    # Determine the file type
    mime_type, _ = mimetypes.guess_type(file_path)
    file_type = mime_type.split('/')[0] if mime_type else 'unknown'
    file_extension = file_path.suffix.lower()
    
    # Handle different file types
    if file_type == 'image':
        # For images, display them directly
        with open(file_path, "rb") as image_file:
            encoded_string = base64.b64encode(image_file.read()).decode()
            return f'<div class="file-viewer"><h3>{file_path.name}</h3><img src="data:{mime_type};base64,{encoded_string}" style="max-width:100%; max-height:500px;"></div>'
    
    elif file_extension in ['.txt', '.py', '.js', '.html', '.css', '.json', '.md', '.csv'] or file_type == 'text':
        # For text files, display the content with syntax highlighting for code
        try:
            content = file_path.read_text(errors='replace')  # Use 'replace' to handle encoding issues
            # Escape HTML characters
            content = content.replace('&', '&amp;').replace('<', '&lt;').replace('>', '&gt;')
            
            # Add syntax highlighting class based on file extension
            highlight_class = ""
            if file_extension == '.py':
                highlight_class = "language-python"
            elif file_extension == '.js':
                highlight_class = "language-javascript"
            elif file_extension == '.html':
                highlight_class = "language-html"
            elif file_extension == '.css':
                highlight_class = "language-css"
            elif file_extension == '.json':
                highlight_class = "language-json"
            
            return f'''
            <div class="file-viewer">
                <h3>{file_path.name}</h3>
                <pre class="{highlight_class}" style="background-color: #f5f5f5; padding: 10px; border-radius: 5px; overflow: auto; max-height: 500px; white-space: pre-wrap;"><code>{content}</code></pre>
                <script>
                    // Add basic syntax highlighting with CSS
                    if (document.querySelector('.language-python')) {{
                        const keywords = ['def', 'class', 'import', 'from', 'return', 'if', 'else', 'elif', 'for', 'while', 'try', 'except', 'with', 'as', 'in', 'not', 'and', 'or', 'True', 'False', 'None'];
                        const code = document.querySelector('.language-python code');
                        let html = code.innerHTML;
                        keywords.forEach(keyword => {{
                            const regex = new RegExp('\\\\b' + keyword + '\\\\b', 'g');
                            html = html.replace(regex, `<span style="color: #0000FF; font-weight: bold;">$&</span>`);
                        }});
                        // Highlight strings
                        html = html.replace(/(["'])(?:(?=(\\\\?))\2.)*?\1/g, '<span style="color: #008000;">$&</span>');
                        // Highlight comments
                        html = html.replace(/(#.*)$/gm, '<span style="color: #808080;">$1</span>');
                        code.innerHTML = html;
                    }}
                </script>
            </div>
            '''
        except UnicodeDecodeError:
            return f'<div class="error-message">Cannot display binary file: {file_path.name}</div>'
    
    elif file_type == 'video':
        # For videos, use video tag
        with open(file_path, "rb") as video_file:
            encoded_string = base64.b64encode(video_file.read()).decode()
            return f'''
            <div class="file-viewer">
                <h3>{file_path.name}</h3>
                <video controls style="max-width:100%; max-height:500px;">
                    <source src="data:{mime_type};base64,{encoded_string}" type="{mime_type}">
                    Your browser does not support the video tag.
                </video>
            </div>
            '''
    
    elif file_type == 'audio':
        # For audio, use audio tag
        with open(file_path, "rb") as audio_file:
            encoded_string = base64.b64encode(audio_file.read()).decode()
            return f'''
            <div class="file-viewer">
                <h3>{file_path.name}</h3>
                <audio controls>
                    <source src="data:{mime_type};base64,{encoded_string}" type="{mime_type}">
                    Your browser does not support the audio tag.
                </audio>
            </div>
            '''
    
    elif file_extension == '.pdf':
        # For PDFs, embed them using an iframe with base64 data
        try:
            with open(file_path, "rb") as pdf_file:
                encoded_string = base64.b64encode(pdf_file.read()).decode()
                return f'''
                <div class="file-viewer">
                    <h3>{file_path.name}</h3>
                    <iframe src="data:application/pdf;base64,{encoded_string}" width="100%" height="500px" style="border: none;"></iframe>
                </div>
                '''
        except Exception as e:
            return f'<div class="error-message">Error displaying PDF: {str(e)}</div>'
    
    else:
        # For other file types, show info but can't display
        size_kb = file_path.stat().st_size / 1024
        return f'<div class="file-viewer"><h3>{file_path.name}</h3><p>File type: {mime_type or "Unknown"}</p><p>Size: {size_kb:.2f} KB</p><p>This file type cannot be displayed in the browser.</p></div>'

def handle_file_upload(files, state):
    """Handle file uploads and store them in the upload directory"""
    if not files:
        return gr.update(choices=[])
    
    file_choices = []
    
    for file in files:
        # Get the file name and create a path in the upload directory
        file_name = Path(file.name).name
        file_path = RUN_FOLDER / file_name
        
        # Save the file
        shutil.copy(file.name, file_path)
        
        # Add to the list of uploaded files
        file_path_str = str(file_path)
        file_choices.append((file_name, file_path_str))
        
        # Add to state
        if file_path_str not in state['uploaded_files']:
            state['uploaded_files'].append(file_path_str)
    
    # Update the view file dropdown with all uploaded files
    all_file_choices = [(Path(path).name, path) for path in state['uploaded_files']]
    
    return gr.update(choices=all_file_choices)

def toggle_view(view_mode, file_path=None, state=None):
    """Toggle between OmniTool Computer view and file viewer"""
    # If switching to File Viewer mode, detect and add new files to the state
    file_choices_update = gr.update()
    if view_mode == "File Viewer" and state is not None:
        file_choices_update = detect_new_files(state)
    
    # Return the appropriate view
    if view_mode == "OmniTool Computer":
        return get_file_viewer_html(), file_choices_update  # This returns the VNC iframe
    else:  # File Viewer mode
        if file_path:
            return get_file_viewer_html(file_path), file_choices_update
        else:
            return get_file_viewer_html(), file_choices_update  # Default to VNC if no file selected

def detect_new_files(state):
    """Detect new files in the uploads folder and add them to the state"""
    new_files_count = 0
    if RUN_FOLDER.exists():
        current_files = set(state['uploaded_files'])
        for file_path in RUN_FOLDER.iterdir():
            if file_path.is_file():
                file_path_str = str(file_path)
                if file_path_str not in current_files:
                    # This is a new file not yet in the state
                    state['uploaded_files'].append(file_path_str)
                    new_files_count += 1
                    print(f"Added new file to state: {file_path_str}")
    
    # Return updated file choices
    file_choices = [(Path(path).name, path) for path in state['uploaded_files']]
    print(f"Detected {new_files_count} new files. Total files in state: {len(state['uploaded_files'])}")
    return gr.update(choices=file_choices)

def refresh_files(state):
    """Refresh the list of files from the current session and detect new files"""
    return detect_new_files(state)

def auto_refresh_files(state):
    """Automatically refresh the list of files from the current session and detect new files"""
    return detect_new_files(state)

with gr.Blocks(theme=gr.themes.Default()) as demo:
    gr.HTML("""
        <style>
        .no-padding {
            padding: 0 !important;
        }
        .no-padding > div {
            padding: 0 !important;
        }
        .markdown-text p {
            font-size: 18px;  /* Adjust the font size as needed */
        }
        </style>
    """)
    state = gr.State({})
    
    setup_state(state.value)
    
    header_image = get_header_image_base64()
    if header_image:
        gr.HTML(f'<img src="{header_image}" alt="OmniTool Header" width="100%">', elem_classes="no-padding")
        gr.HTML('<h1 style="text-align: center; font-weight: normal; margin-bottom: 20px;">Omni<span style="font-weight: bold;">Tool</span></h1>')
    else:
        gr.Markdown("# OmniTool", elem_classes="text-center")

    if not os.getenv("HIDE_WARNING", False):
        gr.HTML(INTRO_TEXT, elem_classes="markdown-text")

    with gr.Accordion("Settings", open=True, elem_classes="accordion-header"): 
        with gr.Row():
            with gr.Column():
                model = gr.Dropdown(
                    label="Model",
                    choices=["omniparser + gpt-4o", "omniparser + o1", "omniparser + o3-mini", "omniparser + R1", "omniparser + qwen2.5vl", "claude-3-5-sonnet-20241022", "omniparser + gpt-4o-orchestrated", "omniparser + o1-orchestrated", "omniparser + o3-mini-orchestrated", "omniparser + R1-orchestrated", "omniparser + qwen2.5vl-orchestrated"],
                    value="omniparser + gpt-4o-orchestrated",
                    interactive=True,
                    container=True
                )
            with gr.Column():
                only_n_images = gr.Slider(
                    label="N most recent screenshots",
                    minimum=0,
                    maximum=10,
                    step=1,
                    value=2,
                    interactive=True
                )
        with gr.Row():
            with gr.Column(1):
                provider = gr.Dropdown(
                    label="API Provider",
                    choices=[option.value for option in APIProvider],
                    value="openai",
                    interactive=False,
                    container=True
                )
            with gr.Column(2):
                api_key = gr.Textbox(
                    label="API Key",
                    type="password",
                    value=state.value.get("api_key", ""),
                    placeholder="Paste your API key here",
                    interactive=True,
                    container=True
                )

    # File Upload Section
    with gr.Accordion("File Upload & Management", open=True, elem_classes="accordion-header"):
        with gr.Row():
            with gr.Column():
                file_upload = gr.File(
                    label="Upload Files",
                    file_count="multiple",
                    type="filepath",
                    elem_classes="file-upload-area"
                )
            with gr.Column():
                with gr.Row():
                    upload_button = gr.Button("Upload Files", variant="primary", elem_classes="primary-button")
                    refresh_button = gr.Button("Refresh Files", variant="secondary", elem_classes="secondary-button")
        
        with gr.Row():
            # Initialize file choices as an empty list
            view_file_dropdown = gr.Dropdown(
                label="View File",
                choices=[],
                interactive=True,
                container=True
            )
            view_toggle = gr.Radio(
                label="Display Mode",
                choices=["OmniTool Computer", "File Viewer"],
                value="OmniTool Computer",
                interactive=True
            )

    with gr.Row():
        with gr.Column(scale=8):
            chat_input = gr.Textbox(
                show_label=False, 
                placeholder="Type a message to send to Omniparser + X ...", 
                container=False
            )
        with gr.Column(scale=1, min_width=50):
            submit_button = gr.Button(value="Send", variant="primary", elem_classes="primary-button")
        with gr.Column(scale=1, min_width=50):
            stop_button = gr.Button(value="Stop", variant="secondary", elem_classes="secondary-button")

    with gr.Row():
        with gr.Column(scale=2):
            chatbot = gr.Chatbot(
                label="Chatbot History", 
                autoscroll=True, 
                height=580,
                avatar_images=("👤", "🤖")
            )
        with gr.Column(scale=3):
            display_area = gr.HTML(
                get_file_viewer_html(),
                elem_classes="no-padding"
            )

    def update_model(model_selection, state):
        state["model"] = model_selection
        print(f"Model updated to: {state['model']}")
        
        if model_selection == "claude-3-5-sonnet-20241022":
            provider_choices = [option.value for option in APIProvider if option.value != "openai"]
        elif model_selection in set(["omniparser + gpt-4o", "omniparser + o1", "omniparser + o3-mini", "omniparser + gpt-4o-orchestrated", "omniparser + o1-orchestrated", "omniparser + o3-mini-orchestrated"]):
            provider_choices = ["openai"]
        elif model_selection == "omniparser + R1":
            provider_choices = ["groq"]
        elif model_selection == "omniparser + qwen2.5vl":
            provider_choices = ["dashscope"]
        else:
            provider_choices = [option.value for option in APIProvider]
        default_provider_value = provider_choices[0]

        provider_interactive = len(provider_choices) > 1
        api_key_placeholder = f"{default_provider_value.title()} API Key"

        # Update state
        state["provider"] = default_provider_value
        state["api_key"] = state.get(f"{default_provider_value}_api_key", "")

        # Calls to update other components UI
        provider_update = gr.update(
            choices=provider_choices,
            value=default_provider_value,
            interactive=provider_interactive
        )
        api_key_update = gr.update(
            placeholder=api_key_placeholder,
            value=state["api_key"]
        )

        return provider_update, api_key_update

    def update_only_n_images(only_n_images_value, state):
        state["only_n_most_recent_images"] = only_n_images_value
   
    def update_provider(provider_value, state):
        # Update state
        state["provider"] = provider_value
        state["api_key"] = state.get(f"{provider_value}_api_key", "")
        
        # Calls to update other components UI
        api_key_update = gr.update(
            placeholder=f"{provider_value.title()} API Key",
            value=state["api_key"]
        )
        return api_key_update
                
    def update_api_key(api_key_value, state):
        state["api_key"] = api_key_value
        state[f'{state["provider"]}_api_key'] = api_key_value

    def clear_chat(state):
        # Reset message-related state
        state["messages"] = []
        state["responses"] = {}
        state["tools"] = {}
        state['chatbot_messages'] = []
        return state['chatbot_messages']

    def view_file(file_path, view_mode):
        """Generate HTML to view the selected file if in File Viewer mode"""
        if view_mode == "File Viewer" and file_path:
            return get_file_viewer_html(file_path)
        elif view_mode == "OmniTool Computer":
            return get_file_viewer_html()  # Return VNC viewer
        else:
            return display_area.value  # Keep current display

    def update_view_file_dropdown(uploaded_files):
        """Update the view file dropdown when uploaded files change"""
        if not uploaded_files:
            return gr.update(choices=[])
        
        file_choices = [(Path(path).name, path) for path in uploaded_files]
        return gr.update(choices=file_choices)

    def reset_view():
        """Reset the view to the VNC viewer"""
        return get_file_viewer_html()

    model.change(fn=update_model, inputs=[model, state], outputs=[provider, api_key])
    only_n_images.change(fn=update_only_n_images, inputs=[only_n_images, state], outputs=None)
    provider.change(fn=update_provider, inputs=[provider, state], outputs=api_key)
    api_key.change(fn=update_api_key, inputs=[api_key, state], outputs=None)
    chatbot.clear(fn=clear_chat, inputs=[state], outputs=[chatbot])

    # File upload event handlers
    upload_button.click(
        fn=handle_file_upload,
        inputs=[file_upload, state],
        outputs=[view_file_dropdown]
    )
    
    # File viewing handlers
    view_file_dropdown.change(
        fn=view_file,
        inputs=[view_file_dropdown, view_toggle],
        outputs=[display_area]
    )
    
    submit_button.click(process_input, [chat_input, state], [chatbot, view_file_dropdown])
    stop_button.click(stop_app, [state], None)
    
    # Toggle view handler
    view_toggle.change(
        fn=toggle_view, 
        inputs=[view_toggle, view_file_dropdown, state], 
        outputs=[display_area, view_file_dropdown]
    )
    
    # Refresh files handler
    refresh_button.click(fn=refresh_files, inputs=[state], outputs=[view_file_dropdown])
    
    # Add JavaScript for auto-refresh instead of using demo.load()
    js_refresh = """
    function() {
        // Auto-refresh files every 5 seconds
        const refreshInterval = setInterval(function() {
            // Find and click the refresh button
            const refreshButtons = document.querySelectorAll('button');
            for (const button of refreshButtons) {
                if (button.textContent.includes('Refresh Files')) {
                    button.click();
                    break;
                }
            }
        }, 5000);
        
        // Return a cleanup function
        return () => clearInterval(refreshInterval);
    }
    """
    
    # Add the JavaScript to the page
    gr.HTML("<script>(" + js_refresh + ")();</script>")
    
if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7888)