frontend/app.py

"""
ATC TTS MOS Evaluation App
Main entrypoint wiring backend + frontend pages.
"""

import time
import gradio as gr
import os
import sys

# Add parent directory (mos) to path so we can import backend
mos_dir = os.path.dirname(os.path.dirname(__file__))
sys.path.insert(0, mos_dir)

# Backend pieces
from backend.data_manager import DataManager
from backend.session_manager import SessionManager
from backend.hf_logging import push_session_to_hub

# Frontend pieces
from frontend.css import CSS
from frontend.pages.intro import build_intro_page
from frontend.pages.samples import build_sample_page
from frontend.pages.mos import build_mos_page, render_mos_html
from frontend.pages.ab_model import build_ab_model_page, render_ab_model_html
from frontend.pages.ab_gender import build_ab_gender_page, render_ab_gender_html
from frontend.pages.conclusion import build_conclusion_page
from frontend.pages.thank_you import build_thank_you_page
from frontend.utils.validation import (
    validate_mos_completion,
    validate_ab_completion,
    validate_final_submission,
)
from frontend.utils.js_collectors import (
    COLLECT_MOS_DATA_JS,
    COLLECT_AB_DATA_JS,
    COLLECT_CONCLUSION_DATA_JS,
    SCROLL_TO_TOP_JS,
)


def create_app(data_manager, session_manager):
    with gr.Blocks(title="ATC TTS Evaluation", css=CSS) as app:
        # Global state
        session_state = gr.State()
        
        # Persistent data state - stores user inputs across page transitions
        mos_data_state = gr.State(value="{}")
        ab_data_state = gr.State(value="{}")

        # Hidden textboxes used by the JS collectors on MOS / AB pages
        mos_data_textbox = gr.Textbox(
            value="", visible=False, elem_id="mos_data_store"
        )
        ab_data_textbox = gr.Textbox(
            value="", visible=False, elem_id="ab_data_store"
        )

        # Build all pages
        intro_page, start_btn = build_intro_page()

        (
            sample_page,
            back_to_intro_btn,
            continue_samples_btn,
            sample_audio_container,
        ) = build_sample_page()

        (
            mos_page,
            continue_mos_btn,
            mos_dynamic_content
        ) = build_mos_page()

        (
            ab_model_page,
            continue_ab_model_btn,
            ab_model_dynamic_content,
        ) = build_ab_model_page()

        (
            ab_gender_page,
            continue_ab_gender_btn,
            ab_gender_dynamic_content,
        ) = build_ab_gender_page()

        (
            conclusion_page,
            export_btn,
            overall_feedback,
            final_comments,
        ) = build_conclusion_page()

        thank_you_page = build_thank_you_page()

        # =========================
        # Event handlers
        # =========================

        def start_evaluation():
            """Create a new session and populate sample page with sample clips."""
            try:
                session = session_manager.create_session()
                print(f"[INFO] Session created: {session['session_id']}")
                print(f"[INFO] Sample clips: {len(session.get('sample_clips', []))}")
            except Exception as e:
                print(f"[ERROR] Error creating session: {e}")
                session = {
                    "session_id": "fallback",
                    "sample_clips": [],
                    "mos_clips": [],
                    "ab_model_pairs": [],
                    "ab_gender_pairs": [],
                    "created_at": time.time(),
                    "completed": False,
                }

            return (
                gr.update(visible=False),  # hide intro
                gr.update(visible=True),   # show sample page
                session,                   # update session_state
            )

        def show_mos_page(session):
            """Navigate from sample page to MOS page and render MOS HTML."""
            if not session:
                return gr.update(), gr.update(), ""

            mos_html = render_mos_html(session)
            return (
                gr.update(visible=False),   # hide sample page
                gr.update(visible=True),    # show MOS page
                mos_html,                   # set MOS HTML content
            )

        def continue_from_mos_to_ab_model(session, mos_data_json):
            """Process MOS data and move to Model-vs-Model page."""
            print(f"[DEBUG] continue_from_mos_to_ab_model called")
            print(f"[DEBUG] MOS data JSON received: '{mos_data_json}'")
            
            # Validate MOS completion
            expected_clips = len(session.get("mos_clips", []))
            is_valid, error_msg = validate_mos_completion(mos_data_json, expected_clips)
            if not is_valid:
                gr.Warning(error_msg)
                return gr.update(), gr.update(), ""
            
            # Process and navigate
            session_manager.process_mos_data(session, mos_data_json)
            ab_model_html = render_ab_model_html(session)
            return (
                gr.update(visible=False),   # hide MOS page
                gr.update(visible=True),    # show AB model page
                ab_model_html,              # set AB model HTML content
            )

        def continue_from_ab_model_to_ab_gender(session, ab_data_json):
            """Process A/B model comparison data and move to Gender A/B page."""
            print(f"[DEBUG] continue_from_ab_model_to_ab_gender called")
            print(f"[DEBUG] AB data JSON received: '{ab_data_json}'")
            
            # Validate AB model completion
            expected_comparisons = len(session.get("ab_model_pairs", []))
            is_valid, error_msg = validate_ab_completion(ab_data_json, expected_comparisons, "model comparisons")
            if not is_valid:
                gr.Warning(error_msg)
                return gr.update(), gr.update(), ""
            
            # Process and navigate
            session_manager.process_ab_data(session, ab_data_json)
            ab_gender_html = render_ab_gender_html(session)
            return (
                gr.update(visible=False),   # hide AB model page
                gr.update(visible=True),    # show AB gender page
                ab_gender_html,             # set AB gender HTML content
            )

        def complete_evaluation_with_ab_gender_data(session, ab_data_json):
            """Process A/B gender comparison data and move to conclusion page."""
            print(f"[DEBUG] complete_evaluation_with_ab_gender_data called")
            print(f"[DEBUG] AB gender data JSON received: '{ab_data_json}'")
            
            # Validate AB gender completion
            expected_comparisons = len(session.get("ab_gender_pairs", []))
            is_valid, error_msg = validate_ab_completion(ab_data_json, expected_comparisons, "gender comparisons")
            if not is_valid:
                gr.Warning(error_msg)
                return gr.update(), gr.update()
            
            # Process and navigate
            session_manager.process_ab_data(session, ab_data_json)
            if session:
                session["completed"] = True
            return (
                gr.update(visible=False),   # hide AB gender page
                gr.update(visible=True),    # show conclusion page
            )

        def export_results_with_data_processing(
            session,
            overall_preference,
            comments,
            mos_data_json,
            ab_data_json,
        ):
            """Final submit: process data, validate overall preference, save, and push to HF."""
            print(f"[DEBUG] export_results_with_data_processing called")
            print(f"[DEBUG] overall_preference: '{overall_preference}'")
            print(f"[DEBUG] comments: '{comments}'")
            
            if not session:
                return gr.update(), gr.update()

            # Validate that overall preference is selected
            is_valid, error_msg = validate_final_submission(overall_preference)
            if not is_valid:
                print(f"[ERROR] {error_msg}")
                gr.Warning(error_msg)
                return gr.update(), gr.update()

            # Process all MOS / AB ratings
            session_manager.process_mos_data(session, mos_data_json)
            session_manager.process_ab_data(session, ab_data_json)

            # Save overall feedback
            feedback_response = {
                "session_id": session["session_id"],
                "overall_preference": overall_preference,
                "final_comments": comments or "",
                "timestamp": time.time(),
            }
            session_manager.save_response("feedback", feedback_response)
            print(f"[INFO] Saved feedback response: {feedback_response}")

            # Export and push to Hub
            export_data = session_manager.export_session(session["session_id"])
            try:
                push_session_to_hub(session["session_id"], export_data)
                print(f"[INFO] Successfully submitted session {session['session_id']} to Hub")
            except Exception as e:
                print(f"[WARN] Failed to push session {session['session_id']} to Hub: {e}")

            return (
                gr.update(visible=False),  # hide conclusion page
                gr.update(visible=True),   # show thank you page
            )

        # =========================
        # Wire buttons to handlers
        # =========================

        # Intro → Samples
        start_btn.click(
            start_evaluation,
            inputs=[],
            outputs=[intro_page, sample_page, session_state],
            js=SCROLL_TO_TOP_JS
        )

        # Samples → MOS
        continue_samples_btn.click(
            show_mos_page,
            inputs=[session_state],
            outputs=[sample_page, mos_page, mos_dynamic_content],
            js=SCROLL_TO_TOP_JS
        )

        # MOS → AB Model (process MOS first)
        continue_mos_btn.click(
            fn=continue_from_mos_to_ab_model,
            inputs=[session_state, mos_data_textbox],
            outputs=[mos_page, ab_model_page, ab_model_dynamic_content],
            js=COLLECT_MOS_DATA_JS
        )

        # AB Model → AB Gender (process AB model first)
        continue_ab_model_btn.click(
            continue_from_ab_model_to_ab_gender,
            inputs=[session_state, ab_data_textbox],
            outputs=[ab_model_page, ab_gender_page, ab_gender_dynamic_content],
            js=COLLECT_AB_DATA_JS
        )

        # AB Gender → Conclusion (process AB gender data)
        continue_ab_gender_btn.click(
            complete_evaluation_with_ab_gender_data,
            inputs=[session_state, ab_data_textbox],
            outputs=[ab_gender_page, conclusion_page],
            js=COLLECT_AB_DATA_JS
        )

        # Conclusion → Thank You (process any remaining data + export)
        export_btn.click(
            export_results_with_data_processing,
            inputs=[
                session_state,
                overall_feedback,
                final_comments,
                mos_data_textbox,
                ab_data_textbox,
            ],
            outputs=[conclusion_page, thank_you_page],
            js=COLLECT_CONCLUSION_DATA_JS
        )

        # Back navigation
        back_to_intro_btn.click(
            lambda: (gr.update(visible=True), gr.update(visible=False)),
            inputs=[],
            outputs=[intro_page, sample_page],
            js=SCROLL_TO_TOP_JS
        )


    return app


if __name__ == "__main__":
    print("Starting ATC TTS MOS Evaluation...")

    try:
        data_manager = DataManager()
        session_manager = SessionManager(data_manager)
        clips = data_manager.load_clips()
        print(f"Dataset ready with {len(clips)} clips")

        models = {clip.model for clip in clips}
        speakers = {clip.speaker for clip in clips}
        exercises = {clip.exercise for clip in clips}

        print(f"Models: {', '.join(models)}")
        print(f"Speakers: {', '.join(speakers)}")
        print(f"Exercises: {len(exercises)} unique exercises")
    except Exception as e:
        print(f"[ERROR] Error loading data: {e}")
        print("Please check your dataset access and try again.")
        raise SystemExit(1)

    app = create_app(data_manager, session_manager)
    app.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        show_error=True,
    )