import gradio as gr
import pandas as pd
import os
import hashlib
from datasets import Dataset, concatenate_datasets, load_dataset
from huggingface_hub import login

# ============================
# CONFIGURATION
# ============================

AUDIO_FOLDER = "audio_files"
audio_files = sorted(os.listdir(AUDIO_FOLDER))

ADMIN_PASSWORD_HASH = "eb32da077dfaf326cd6f73e0716b628da6427aa318a2d0b9fafa9ef315b5e885"

# ============================
# HF DATASET CONFIGURATION
# ============================

HF_DATASET_NAME = "MeysamSh/Annotate_Samples"  # replace with your HF dataset
HF_TOKEN = os.environ.get("HF_TOKEN")  # store your token as a secret in Spaces
login(HF_TOKEN)

# ============================
# USER FUNCTIONS
# ============================

def load_hf_dataset():
    """Load existing HF Dataset or create empty one if not exists."""
    try:
        ds = load_dataset(HF_DATASET_NAME, split="train")
    except:
        # Dataset does not exist yet
        df = pd.DataFrame(columns=["user_id", "gender", "audio_file", "score"])
        ds = Dataset.from_pandas(df)
        ds.push_to_hub(HF_DATASET_NAME, private=True)
    return ds

def submit_annotation_hf(user_id, gender, score, audio_file):
    """Append annotation to HF Dataset."""
    ds = load_hf_dataset()
    new_row = pd.DataFrame([{
        "user_id": user_id,
        "gender": gender,
        "audio_file": audio_file,
        "score": score
    }])
    ds_new = Dataset.from_pandas(new_row)
    ds_combined = concatenate_datasets([ds, ds_new])
    ds_combined.push_to_hub(HF_DATASET_NAME)

def load_next_audio(state):
    if state is None:
        state = {"index": 0}

    idx = state["index"]

    if idx >= len(audio_files):
        return None, state, {"played": 0}, "All audio files annotated."

    filepath = os.path.join(AUDIO_FOLDER, audio_files[idx])
    return filepath, state, {"played": 0}, f"Loaded {audio_files[idx]}"

def submit_annotation(user_id, gender, score, state):
    if state is None:
        state = {"index": 0}

    idx = state["index"]

    if idx >= len(audio_files):
        return state, "No more audio files."

    audio_file = audio_files[idx]

    submit_annotation_hf(user_id, gender, score, audio_file)

    state["index"] += 1
    return state, f"Saved rating for {audio_file}."

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

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

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

# ============================
# ADMIN FUNCTIONS
# ============================

def hash_password(password: str) -> str:
    return hashlib.sha256(password.encode()).hexdigest()

def admin_login(input_password):
    if hash_password(input_password) == ADMIN_PASSWORD_HASH:
        ds = load_hf_dataset()
        df = ds.to_pandas()
        # For download: write temp CSV
        temp_csv = "annotations_export.csv"
        df.to_csv(temp_csv, index=False)
        return (
            gr.update(visible=True),
            df,
            temp_csv,
            gr.update(value="Admin authentication successful.")
        )
    else:
        return (
            gr.update(visible=False),
            None,
            None,
            gr.update(value="Admin authentication failed.")
        )

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

with gr.Blocks() as demo:

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

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

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

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

    audio_player = gr.Audio(label="Audio File")

    with gr.Row():
        score = gr.Dropdown(choices=["None","1", "2", "3", "4", "5"], value="None", label="MOS Score (1–5)")
        submit_btn = gr.Button("Submit Score", interactive=False)

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

    # Load first audio
    demo.load(
        load_next_audio,
        inputs=state,
        outputs=[audio_player, state, audio_played, status]
    )

    # Mark audio as played
    audio_player.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, score],
        outputs=submit_btn
    )
    audio_played.change(
        check_submit_ready,
        inputs=[user_id, audio_played, score],
        outputs=submit_btn
    )
    score.change(
        check_submit_ready,
        inputs=[user_id, audio_played, score],
        outputs=submit_btn
    )

    # Save annotation
    submit_btn.click(
        submit_annotation,
        inputs=[user_id, gender, score, state],
        outputs=[state, status]
    )

    # Load next audio
    submit_btn.click(
        load_next_audio,
        inputs=state,
        outputs=[audio_player, state, audio_played, 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()