import gradio as gr
import csv
import os
import shutil
from datetime import datetime
from huggingface_hub import HfApi

# --- CONFIGURATION ---

# 1. SECURITY
HF_TOKEN = os.getenv("HF_TOKEN") 
BACKUP_REPO_NAME = "kinyarwanda-voice-backup"

# 2. STORAGE
DATA_ROOT = "/data" if os.path.exists("/data") else "local_data"
AUDIO_FOLDER = os.path.join(DATA_ROOT, "clips")
CSV_FILE = os.path.join(DATA_ROOT, "metadata.csv")
TSV_FILE = os.path.join(DATA_ROOT, "validated.tsv") 

os.makedirs(AUDIO_FOLDER, exist_ok=True)

# 3. SCRIPTS (Anonymized & Improved)
SCRIPTS = [
    "Nagerageje kohereza amafaranga kuri MoMo, ariko rezo yari mbi cyane sinabasha kubona mesaje yemeza ko yagiye.",
    "Mbere yo kwinjira muri mudasobwa, banza urebe niba sisitemu yakoze update, hanyuma ushyiremo ijambo ry'ibanga kugira ngo ufungure.",
    "Nta megabayiti zihagije mfite muri telefoni, reka nshakishe wifi hano hafi kugira ngo mbashe gukurura iyo porogaramu vuba.",
    "Iyi sharijeri ya smartphone iragura ibihumbi bitanu, ariko niba ushaka n'ibirahure bya ecran, turaguha byose ku icumi.",
    "Nugera kuri banki i Nyarugenge, ubwire ejenti agufashe kubikuza ayo madolari mbere y'uko ukwezi kwa Gicurasi kurangira.",
    "Wohereze imeli itarimo amakosa kugira ngo ubashe guhindura mot de passe, hanyuma ukande kuri linki iri bube kuri ecran yawe.",
    "Ugomba gusiba amafoto adakenewe kugira ngo ubone aho kubika izindi gigabayiti ebyiri, cyangwa ukoreshe flash.",
    "Koresha Mokash wishyure iyo fagitire, kode ni zeru gatatu rimwe, ubundi tujye Gasabo ku wa Gatanu.",
    "Koresha kibodi wandike ubutumwa, hanyuma uwohereze kuri WhatsApp kuko interineti yo mu Majyaruguru iragenda buhoro.",
    "Nurangiza gukoresha iyo apurikasiyo, wibuke gukora log out kugira ngo hatagira undi umukiriya ureba amabanga yawe."
]

# --- HELPER FUNCTIONS ---

def generate_tsv():
    """Generates the Standard Mozilla 'validated.tsv' for the backup."""
    cv_header = ["client_id", "path", "sentence", "up_votes", "down_votes", "age", "gender", "accent", "locale", "segment"]
    tsv_content = []
    tsv_content.append("\t".join(cv_header))
    
    if os.path.exists(CSV_FILE):
        with open(CSV_FILE, 'r', encoding='utf-8') as f:
            reader = csv.reader(f)
            try:
                next(reader, None) 
                for row in reader:
                    if len(row) < 5: continue 
                    participant_id, age, gender, accent = row[0], row[2], row[3].lower(), row[4]
                    for i in range(len(SCRIPTS)):
                        file_index = 5 + i
                        if file_index < len(row):
                            filename = row[file_index]
                            if filename: 
                                tsv_row = [
                                    participant_id, filename, SCRIPTS[i], 
                                    "1", "0", age, gender, accent, "rw", ""
                                ]
                                tsv_content.append("\t".join(tsv_row))
            except StopIteration:
                pass
    
    with open(TSV_FILE, 'w', encoding='utf-8') as f:
        f.write("\n".join(tsv_content))

def backup_to_dataset_repo(repo_name=BACKUP_REPO_NAME):
    """Backs up data to Hugging Face Dataset."""
    if not HF_TOKEN:
        print("⚠️ Backup Skipped: HF_TOKEN missing.")
        return
    
    generate_tsv()
    
    api = HfApi(token=HF_TOKEN)
    user = api.whoami()["name"]
    full_repo_id = f"{user}/{repo_name}"
    
    try:
        api.create_repo(repo_id=full_repo_id, repo_type="dataset", private=True, exist_ok=True)
        api.upload_folder(
            folder_path=DATA_ROOT,
            repo_id=full_repo_id,
            repo_type="dataset",
            path_in_repo="data", 
            commit_message=f"Auto-backup {datetime.now().strftime('%Y-%m-%d %H:%M')}"
        )
        print(f"✅ Auto-Backup successful to {full_repo_id}")
    except Exception as e:
        print(f"⚠️ Backup Failed: {e}")

def save_data(age, gender, region, *audios):
    """
    Saves data locally. 
    Includes specific error checking to tell user WHICH audio is missing.
    """
    # 1. Identify missing recordings
    missing_indices = []
    for i, audio in enumerate(audios):
        if audio is None:
            missing_indices.append(str(i+1)) # Store 1-based index
            
    if missing_indices:
        missing_str = ", ".join(missing_indices)
        error_msg = f"⚠️ Error: You missed Sentence(s): {missing_str}. Please record them."
        return f"<span style='color: red'>{error_msg}</span>", gr.update(visible=True), gr.update(visible=True)

    # 2. Generate Filenames & Save
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    participant_id = f"user_{timestamp}"
    filenames = []
    
    for index, audio_path in enumerate(audios):
        filename = f"{participant_id}_s{index+1}.wav"
        save_path = os.path.join(AUDIO_FOLDER, filename)
        shutil.copy(audio_path, save_path)
        filenames.append(filename)

    # 3. Append to CSV
    file_exists = os.path.exists(CSV_FILE)
    with open(CSV_FILE, mode='a', newline='', encoding='utf-8') as f:
        writer = csv.writer(f)
        if not file_exists:
            header = ["ID", "Timestamp", "Age_Range", "Gender", "Accent_Region"] + [f"File_{i+1}" for i in range(len(audios))]
            writer.writerow(header)
        writer.writerow([participant_id, timestamp, age, gender, region] + filenames)
    
    return f"✅ Saved {participant_id}.", gr.update(visible=False), gr.update(visible=True)

def transition_to_recorder(age, gender, region):
    if not age or not gender or not region:
        return gr.update(visible=True), gr.update(visible=False), "<span style='color: red'>⚠️ Please fill in all fields.</span>"
    return gr.update(visible=False), gr.update(visible=True), ""

# --- UI BUILDER ---

# CSS: System Sans-Serif Stack (Fastest & Cleanest)
system_sans_css = """
body, .gradio-container {
    font-family: system-ui, -apple-system, "Segoe UI", Roboto, "Helvetica Neue", "Noto Sans", "Liberation Sans", Arial, sans-serif, "Apple Color Emoji", "Segoe UI Emoji", "Segoe UI Symbol", "Noto Color Emoji" !important;
}

button, input, textarea, span, div, label, p, h1, h2, h3 {
    font-family: system-ui, -apple-system, "Segoe UI", Roboto, "Helvetica Neue", "Noto Sans", "Liberation Sans", Arial, sans-serif, "Apple Color Emoji", "Segoe UI Emoji", "Segoe UI Symbol", "Noto Color Emoji" !important;
}
"""

with gr.Blocks(title="Improve Kinyarwanda AI", css=system_sans_css, theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🇷🇼 Improve Kinyarwanda AI")
    
    # PAGE 1: METADATA
    with gr.Column(visible=True) as metadata_page:
        gr.Markdown("### Intambwe 1: Umwirondoro (Step 1: Profile)")
        with gr.Group():
            age_input = gr.Dropdown(
                choices=["teens", "twenties", "thirties", "fourties", "fifties"], 
                label="Age Range (Icyiciro cy'imyaka)", 
                value="twenties"
            )
            gender_input = gr.Radio(["Male", "Female"], label="Gender", value="Female")
            region_input = gr.Dropdown(
                ["Kigali City", "Musanze", "Rubavu", "Eastern", "Southern", "Diaspora"], 
                label="Accent Origin", 
                value="Kigali City"
            )
        error_msg = gr.Markdown("") 
        next_btn = gr.Button("Komeza (Next) ➡️", variant="primary")

    # PAGE 2: RECORDING
    with gr.Column(visible=False) as recording_page:
        gr.Markdown(f"### Intambwe 2: Fata Amajwi {len(SCRIPTS)} (Step 2: Record {len(SCRIPTS)} Sentences)")
        audio_inputs = [] 
        for i, text in enumerate(SCRIPTS):
            gr.Markdown(f"**{i+1}. {text}**")
            audio = gr.Audio(sources=["microphone"], type="filepath", label=f"Sentence {i+1}", scale=3)
            audio_inputs.append(audio)
        
        submit_btn = gr.Button("Ohereza Byose (Submit All)", variant="primary")
        status_msg = gr.Label(label="Status")
    
    # PAGE 3: SUCCESS
    with gr.Column(visible=False) as success_page:
        gr.Markdown("# ✅ Murakoze!")
        gr.Button("Refresh Page to Start New").click(None, js="window.location.reload()")

    # --- WIRING ---
    next_btn.click(fn=transition_to_recorder, inputs=[age_input, gender_input, region_input], outputs=[metadata_page, recording_page, error_msg])
    
    # 1. Save Local
    save_event = submit_btn.click(
        fn=save_data, 
        inputs=[age_input, gender_input, region_input] + audio_inputs, 
        outputs=[status_msg, recording_page, success_page]
    )

    # 2. Auto-Backup to Cloud (Background)
    save_event.then(
        fn=lambda: backup_to_dataset_repo(), 
        inputs=None,
        outputs=None 
    )

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