import gradio as gr
import time
from apscheduler.schedulers.background import BackgroundScheduler
import threading
import globals
from utils.io import initialize_models_providers_file, save_results, load_results, load_models_providers, get_results_table
from utils.jobs import run_single_job, launch_jobs, update_job_statuses
from typing import List, Optional


def status_monitor() -> None:
    """Background thread to monitor job statuses."""
    while True:
        update_job_statuses()
        time.sleep(240)  # Check every 30 seconds


def daily_checkpoint() -> None:
    """Daily checkpoint - save current state."""
    print("Daily checkpoint - saving current state")
    save_results()


# Create Gradio interface
def create_app() -> gr.Blocks:
    with gr.Blocks(title="Inference Provider Testing Dashboard") as demo:
        gr.Markdown("# Inference Provider Testing Dashboard")
        gr.Markdown("Launch and monitor evaluation jobs for multiple models and providers.")

        output = gr.Textbox(label="Logs and status", interactive=False)

        with gr.Row():
            with gr.Column():
                gr.Markdown("## Initialize Config File")
                init_btn = gr.Button("Fetch and Initialize Models/Providers", variant="secondary")

        with gr.Row():
            with gr.Column():
                gr.Markdown("## Launch Jobs")
                launch_btn = gr.Button("Launch All Jobs", variant="primary")

        with gr.Row():
            with gr.Column():
                gr.Markdown("## Job Results")
                results_table = gr.Dataframe(
                    value=get_results_table(),
                    interactive=True,
                    show_search="search",
                    show_copy_button=True,
                    show_fullscreen_button=True,
                    wrap=True,
                    static_columns=list(range(7)),
                    datatype=["str", "str", "str", "str", "str", "str", "html", "str"],
                    elem_id="results_table"
                )
                refresh_btn = gr.Button("Refresh Results")

        def update_model_choices() -> gr.update:
            models_providers = load_models_providers(globals.LOCAL_CONFIG_FILE)
            model_choices = sorted(list(set([mp[0] for mp in models_providers])))
            return gr.update(choices=model_choices, value=model_choices[0] if model_choices else None)


        def update_provider_choices(model: Optional[str]) -> gr.update:
            """Update provider dropdown based on selected model."""
            if not model:
                return gr.update(choices=[])

            # Get providers for the selected model from the config file
            models_providers = load_models_providers(globals.LOCAL_CONFIG_FILE)
            providers = [mp[1] for mp in models_providers if mp[0] == model]

            return gr.update(choices=providers, value=providers[0] if providers else None)

        # Event handlers
        init_btn.click(
            fn=initialize_models_providers_file,
            outputs=output
        )

        launch_btn.click(
            fn=launch_jobs,
            outputs=output
        )

        refresh_btn.click(
            fn=get_results_table,
            outputs=results_table
        )

        # Handle dataframe cell selection for relaunch
        def handle_table_select(evt: gr.SelectData):
            """Handle when a cell in the results table is clicked."""
            print(f"[Relaunch] Cell selected - Row: {evt.index[0]}, Col: {evt.index[1]}, Value: {evt.value}")

            # Check if the Actions column (index 7) was clicked
            if evt.index[1] == 7:
                # Get the full row data from the dataframe
                df = get_results_table()
                row_data = df.data.iloc[evt.index[0]]

                model = row_data['Model']
                provider = row_data['Provider']
                print(f"[Relaunch] Relaunching job - Model: {model}, Provider: {provider}")

                return run_single_job(model, provider, globals.TASKS)

        results_table.select(
            fn=handle_table_select,
            inputs=[],
            outputs=output
        )

    return demo


if __name__ == "__main__":
    # Load previous results
    load_results()
    print("Starting Inference Provider Testing Dashboard")

    # Start status monitor thread
    monitor_thread = threading.Thread(target=status_monitor, daemon=True)
    monitor_thread.start()
    print("Job status monitor started")

    # Start APScheduler for daily checkpoint
    scheduler = BackgroundScheduler()
    scheduler.add_job(daily_checkpoint, 'cron', hour=0, minute=0)  # Run at midnight
    scheduler.start()
    print("Daily checkpoint scheduler started (saves at 00:00)")

    # Create and launch the Gradio interface
    demo = create_app()
    demo.launch(server_name="0.0.0.0", server_port=7860)