import gradio as gr

from config.settings import *
from data.contexts import CONTEXTS
from data.dataset_manager import load_hf_dataset
from services.annotation_service import submit_pair_annotation
from admin.admin_service import admin_login
from services.audio_service import generate_audio_pair


def update_sentence(context):
        return CONTEXTS[context]

# ============================
# SUBMIT BUTTON CONTROL
# ============================

def check_submit_ready(
        user_id,
        audio_played, 
        mos_A, 
        mos_B
):
    ready = (
        len(user_id.strip()) > 1 
        and audio_played['played'] == 1 
        and mos_A != "None" 
        and mos_B != "None"
    )
    return gr.update(interactive=ready)

def mark_audio_played():
    return {"played": 1}

# ============================
# GRADIO UI
# ============================

with gr.Blocks() as demo:

    gr.Markdown("# Audio MOS Annotation Tool")

    # --------------------------
    # USER SECTION
    # --------------------------

    audio_played = gr.State({"played": 0})

    hidden_models = gr.State()

    with gr.Row():
        user_id = gr.Textbox(label="User ID")
        gender = gr.Dropdown(["Male", "Female", "Other"], label="Gender", value="Other")

    with gr.Row():
        age_group = gr.Dropdown(
            choices=[
                "10-19",
                "20-29",
                "30-39",
                "40-49",
                "50-59",
                "60-69",
                "70+"
            ],
            label="Age Group"
        )

        native_language = gr.Textbox(
            label="Native Language"
        )

        tts_experience = gr.Dropdown(
            choices=[
                "Never",
                "Rarely",
                "Occasionally",
                "Frequently",
                "Expert"
            ],
            label="Experience with speech synthesis"
        )

        device_type = gr.Dropdown(
            choices=[
                "Headphones",
                "Earbuds",
                "Computer Speakers",
                "Smartphone Speaker"
            ],
            label="Listening Device"
        )

    with gr.Row():
        context_selector = gr.Dropdown(
            choices=list(CONTEXTS.keys()),
            value="Voice assistant",
            label="Context"
        )

        #sentence_display = gr.Textbox(
        #    value = CONTEXTS["Voice assistant"],
        #    label = "Sentence",
        #    interactive = False
        #)

        #model_selector = gr.Dropdown(
        #    choices = MODEL_OPTIONS,
        #    value = "OmniVoice",
        #    label = "Model"
        #)

        #context_selector.change(
        #    update_sentence,
        #    inputs=context_selector,
        #    outputs=sentence_display
        #)

        generate_btn = gr.Button("Generate Audio")

    #audio_player = gr.Audio(label="Generated audio")

    audio_A = gr.Audio(
        label = "Audio A"
    )

    audio_B = gr.Audio(
        label = "Audio B"
    )
    
    generate_btn.click(
        generate_audio_pair,
        inputs=context_selector,
        outputs = [
            audio_A,
            audio_B,
            audio_played,
            hidden_models
        ]
    )

    with gr.Row():
        naturalness_A = gr.Slider(
            minimum=1,
            maximum=5,
            step=1,
            value=3,
            label="Naturalness A"
        )

        intelligibility_A = gr.Slider(
            minimum=1,
            maximum=5,
            step=1,
            value=3,
            label="Intelligibility A"
        )

        context_A = gr.Slider(
            minimum=1,
            maximum=5,
            step=1,
            value=3, 
            label="Context Adequacy A"
        )

        mos_A = gr.Slider(
            minimum=1,
            maximum=5,
            step=1,
            value=3,
            label="MOS Audio A"
        )

    with gr.Row():
        naturalness_B = gr.Slider(
            minimum=1,
            maximum=5,
            step=1,
            value=3,
            label="Naturalness B"
        )

        intelligibility_B = gr.Slider(
            minimum=1,
            maximum=5,
            step=1,
            value=3,
            label="Intelligibility B"
        )

        context_B = gr.Slider(
            minimum=1,
            maximum=5,
            step=1,
            value=3, 
            label="Context Adequacy for audio B"
        )

        mos_B = gr.Slider(
            minimum=1,
            maximum=5,
            step=1,
            value=3,
            label="MOS Audio B"
        )

    with gr.Row():
        preferred_audio = gr.Radio(
            choices=[
                "Audio A",
                "Audio B",
                "No preference"
            ],
            label="Preferred Audio"
        )

        submit_btn = gr.Button("Submit Score", interactive=False)


    status = gr.Textbox(label="Status", interactive=False)


    # Mark audio as played
    audio_A.play(
        mark_audio_played,
        None,
        audio_played
    )

    audio_B.play(
        mark_audio_played,
        None,
        audio_played
    )

    # Enable submit button only when conditions are met
    user_id.change(
        check_submit_ready,
        inputs=[user_id, audio_played, mos_A, mos_B],
        outputs=submit_btn
    )
    audio_played.change(
        check_submit_ready,
        inputs=[user_id, audio_played, mos_A, mos_B],
        outputs=submit_btn
    )
    mos_A.change(
        check_submit_ready,
        inputs=[user_id, audio_played, mos_A, mos_B],
        outputs=submit_btn
    )
    mos_B.change(
        check_submit_ready,
        inputs=[user_id, audio_played, mos_A, mos_B],
        outputs=submit_btn
    )

    # Save annotation
    submit_btn.click(
        submit_pair_annotation,
        inputs=[
            user_id, 
            age_group, 
            gender, 
            native_language, 
            tts_experience, 
            device_type, 
            hidden_models,
            context_selector, 
            naturalness_A,
            intelligibility_A,
            context_A,
            mos_A,
            naturalness_B,
            intelligibility_B,
            context_B,
            mos_B,
            preferred_audio
        ],
        outputs=status
    )

    # ============================
    # ADMIN DASHBOARD
    # ============================

    gr.Markdown("## Admin Dashboard (Restricted Access)")

    with gr.Row():
        admin_password = gr.Textbox(label="Admin Password", type="password")
        admin_login_btn = gr.Button("Login")

    login_status = gr.Textbox(label="Login Status:", interactive=False)

    with gr.Column(visible=False) as admin_panel:
        gr.Markdown("### Annotation Results")
        pd_loaded=load_hf_dataset().to_pandas()
        results_table = gr.DataFrame(interactive=False, value=pd_loaded)
        pd_loaded.to_csv("annotations_export.csv", index=False)
        download_admin = gr.File(label="Download annotations.csv", value="annotations_export.csv")

    admin_login_btn.click(
        admin_login,
        inputs=admin_password,
        outputs=[admin_panel, results_table, download_admin, login_status]
    )

# ============================
# APP LAUNCH
# ============================

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