import gradio as gr
import os
import re
import base64
import io
from pathlib import Path
import pypdfium2 as pdfium
from huggingface_hub import snapshot_download, HfApi

HF_TOKEN = os.environ.get("HF_TOKEN")
PRIVATE_DATA = os.environ.get("PRIVATE_DATA")

# FIX: removed duplicate `api = HfApi(...)` line that appeared twice
api = HfApi(token=HF_TOKEN)

# Check access
info = api.repo_info(repo_id=PRIVATE_DATA, repo_type="dataset")
print(info.id)

# Download
dataset_path = snapshot_download(
    repo_id=PRIVATE_DATA,
    repo_type="dataset",
    token=HF_TOKEN,
    local_dir="data",
)
print(dataset_path)

# --- Configuration ---
AUDIO_FILES_INFO = {}

LANGUAGES = {
    "Arabic": "ar",
    "English": "en",
    "French": "fr",
    "Spanish": "es",
    "German": "de",
    "Italian": "it",
}


def discover_audio_files(root_dir=None):
    """
    Scans the demo_audio directory for wav files from a specific list.
    FIX: default arg now uses None to avoid capturing dataset_path at import time.
    """
    # FIX: resolve root_dir lazily so dataset_path is always defined
    if root_dir is None:
        root_dir = f"{dataset_path}/demo_audio"

    allowed_files = [
        "A02997", "A03046", "A03520", "A03858", "A04616", "A04957", "A04960",
        "A04965", "A05012", "A05195", "A05857", "A06010", "A06118", "A06622", "A06850",
        "A06888", "A06918", "A06936", "A07167", "A07251", "A07296", "A07452", "A07462",
        "A07679", "A07768", "A07791", "A08133", "A08649", "A08957", "A08958", "A08960",
        "A08966", "A09304"
    ]

    audio_files = {}
    if not os.path.isdir(root_dir):
        print(f"Warning: Audio source directory '{root_dir}' not found.")
        return audio_files

    for file_id in allowed_files:
        expected_wav_path = os.path.join(root_dir, file_id, f"{file_id}.wav")
        if os.path.exists(expected_wav_path):
            audio_files[file_id] = file_id

    if not audio_files:
        print(f"Warning: No allowed audio files found in '{root_dir}'")

    return audio_files


# --- Gradio Application Logic ---
def load_audio_and_transcript(audio_selection):
    """
    Callback: loads audio, transcript, enables language dropdown, defaults to English PDF.
    FIX: all component returns now use gr.update() — Gradio 5+ requires this,
         returning gr.Audio(...) / gr.Dropdown(...) instances is silently ignored.
    """
    if not audio_selection or audio_selection not in AUDIO_FILES_INFO:
        return (
            None,
            "Please select an audio file to see its transcript.",
            gr.update(interactive=False, value=None),
            gr.update(value="<p style='color:#64748b; padding:16px;'>Select an audio file to load the PDF summary.</p>", visible=True),
        )

    base_filename = AUDIO_FILES_INFO[audio_selection]

    audio_path = os.path.join(
        f"{dataset_path}/demo_audio", base_filename, f"{base_filename}.wav"
    )

    transcript_path = os.path.join(
        f"{dataset_path}/meeting_reports", base_filename, "final_transcript_text.txt"
    )
    try:
        with open(transcript_path, "r", encoding="utf-8") as f:
            transcript_content = f.read()
    except FileNotFoundError:
        transcript_content = f"Transcript file not found at: {transcript_path}"

    english_pdf_html = load_summary_pdf(audio_selection, "English")
    pdf_update = gr.update(
        value=english_pdf_html or "<p>No English PDF found.</p>",
        visible=True,
    )

    return (
        gr.update(value=audio_path, label=f"Playing: {audio_selection}"),
        transcript_content,
        gr.update(interactive=True, value="English"),
        pdf_update,
    )


def pdf_path_to_html(pdf_path):
    """Render PDF pages as high-quality embedded PNG images inside HTML."""
    try:
        pdf = pdfium.PdfDocument(pdf_path)
        page_html = []

        # 300 DPI ~= 300 / 72 scale in PDFium
        render_scale = 300 / 72

        for i in range(len(pdf)):
            page = pdf[i]
            pil_image = page.render(scale=render_scale).to_pil()

            buffer = io.BytesIO()
            pil_image.save(buffer, format="PNG", optimize=True)
            img_b64 = base64.b64encode(buffer.getvalue()).decode("utf-8")

            page_html.append(f"""
                <div style="margin-bottom: 20px; text-align: center;">
                    <div style="font-size: 14px; color: #64748b; margin-bottom: 8px;">
                        Page {i + 1}
                    </div>
                    <img
                        src="data:image/png;base64,{img_b64}"
                        style="
                            width: 100%;
                            height: auto;
                            border: none;
                            border-radius: 8px;
                            box-shadow: 0 1px 4px rgba(0,0,0,0.08);
                        "
                    />
                </div>
            """)

        return f"""
            <div style="width: 100%; height: 1600px; overflow-y: auto; padding-right: 8px;">
                {''.join(page_html)}
            </div>
        """

    except Exception as e:
        return f'<p style="color:red;">Could not load PDF: {e}<br>Path: {pdf_path}</p>'


def load_summary_pdf(audio_selection, language_selection):
    """Returns HTML preview for the PDF."""
    if not audio_selection or not language_selection or audio_selection not in AUDIO_FILES_INFO:
        return None

    base_filename = AUDIO_FILES_INFO[audio_selection]
    lang_name = language_selection

    report_dir = os.path.join(f"{dataset_path}/meeting_reports", base_filename)
    if not os.path.isdir(report_dir):
        print(f"[PDF] ERROR: Report directory not found: {report_dir}")
        return None

    all_pdfs = [f for f in os.listdir(report_dir) if f.lower().endswith(".pdf")]
    print(f"[PDF] Looking for language='{lang_name}' in {report_dir}")
    print(f"[PDF] Available PDFs: {all_pdfs}")

    pdf_path = None

    # Strategy 1: strict pattern, case-insensitive
    pattern = re.compile(
        f"meeting_analysis_{re.escape(base_filename)}-.*?{re.escape(lang_name)}\\.pdf",
        re.IGNORECASE,
    )
    for filename in all_pdfs:
        if pattern.match(filename):
            pdf_path = os.path.join(report_dir, filename)
            print(f"[PDF] Matched: {filename}")
            break

    # Strategy 2: any PDF containing the language name
    if not pdf_path:
        for filename in all_pdfs:
            if lang_name.lower() in filename.lower():
                pdf_path = os.path.join(report_dir, filename)
                print(f"[PDF] Fallback match: {filename}")
                break

    # Strategy 3: first PDF available
    if not pdf_path and all_pdfs:
        pdf_path = os.path.join(report_dir, all_pdfs[0])
        print(f"[PDF] Last resort: {all_pdfs[0]}")

    if not pdf_path:
        print(f"[PDF] ERROR: No PDFs found for {base_filename}")
        return None

    return pdf_path_to_html(pdf_path)


def load_summary_pdf_with_visibility(audio_selection, language_selection):
    html = load_summary_pdf(audio_selection, language_selection)
    if html:
        return gr.update(value=html, visible=True)
    return gr.update(value="<p>No PDF found for this selection.</p>", visible=True)

def create_interface():
    # --- Default initial values ---
    default_audio = list(AUDIO_FILES_INFO.keys())[0] if AUDIO_FILES_INFO else None
    default_lang = "English" if default_audio else None

    if default_audio:
        default_audio_path = os.path.join(
            f"{dataset_path}/demo_audio", default_audio, f"{default_audio}.wav"
        )

        transcript_path = os.path.join(
            f"{dataset_path}/meeting_reports", default_audio, "final_transcript_text.txt"
        )
        try:
            with open(transcript_path, "r", encoding="utf-8") as f:
                default_transcript = f.read()
        except FileNotFoundError:
            default_transcript = f"Transcript file not found at: {transcript_path}"

        default_pdf_html = load_summary_pdf(default_audio, default_lang) or "<p>No English PDF found.</p>"
        default_audio_label = f"Playing: {default_audio}"
    else:
        default_audio_path = None
        default_transcript = "Please select an audio file to see its transcript."
        default_pdf_html = "<p style='color:#64748b; padding:16px;'>Select an audio file to load the PDF summary.</p>"
        default_audio_label = "Audio Player"

    css = """
        :root {
            --primary-color: #2563eb;
            --primary-hover: #1d4ed8;
            --secondary-color: #8b5cf6;
            --accent-color: #06b6d4;
            --success-color: #10b981;
            --bg-primary: #ffffff;
            --bg-secondary: #f8fafc;
            --bg-tertiary: #f1f5f9;
            --border-color: #e2e8f0;
            --text-primary: #1e293b;
            --text-secondary: #64748b;
            --shadow-sm: 0 1px 2px 0 rgba(0, 0, 0, 0.05);
            --shadow-md: 0 4px 6px -1px rgba(0, 0, 0, 0.1);
            --shadow-lg: 0 10px 15px -3px rgba(0, 0, 0, 0.1);
            color-scheme: light !important;
        }
        * { color-scheme: light !important; }

        .gradio-container {
            background: linear-gradient(135deg, #f0f9ff 0%, #e0f2fe 50%, #faf5ff 100%) !important;
            color-scheme: light !important;
        }
        .gradio-container * { color-scheme: light !important; }

        .box-style {
            background: rgba(255, 255, 255, 0.95) !important;
            border-radius: 16px !important;
            border: 1px solid var(--border-color) !important;
            padding: 24px !important;
            box-shadow: var(--shadow-md) !important;
        }

        
        #main-content-row {
            min-height: 9vh;
        }
    """

    theme = gr.themes.Soft(
        primary_hue="blue",
        secondary_hue="purple",
        neutral_hue="slate",
        font=gr.themes.GoogleFont("Inter"),
        font_mono=gr.themes.GoogleFont("JetBrains Mono"),
    ).set(
        body_background_fill="*neutral_50",
        body_background_fill_dark="*neutral_50",
        background_fill_primary="white",
        background_fill_primary_dark="white",
        background_fill_secondary="*neutral_50",
        background_fill_secondary_dark="*neutral_50",
        block_background_fill="white",
        block_background_fill_dark="white",
        input_background_fill="white",
        input_background_fill_dark="white",
        panel_background_fill="white",
        panel_background_fill_dark="white",
        body_text_color="*neutral_800",
        body_text_color_dark="*neutral_800",
    )

    with gr.Blocks(
        theme=theme,
        css=css,
        title="InterPARES-Audio",
        js="""
        function() {
            document.documentElement.style.colorScheme = 'light';
            document.body.style.colorScheme = 'light';
            document.documentElement.classList.remove('dark');
            document.body.classList.remove('dark');
            document.documentElement.classList.add('light');
            document.body.classList.add('light');
            document.documentElement.setAttribute('data-theme', 'light');
            document.body.setAttribute('data-theme', 'light');
        }
        """,
    ) as demo:

        gr.HTML("""
            <div style="background: linear-gradient(135deg, #2563eb 0%, #8b5cf6 50%, #06b6d4 100%);
                        padding: 48px 32px; border-radius: 20px; margin-bottom: 32px;
                        box-shadow: 0 20px 25px -5px rgba(0,0,0,0.1), 0 10px 10px -5px rgba(0,0,0,0.04);
                        position: relative; overflow: hidden;">
                <div style="position: absolute; top:0; left:0; right:0; bottom:0;
                            background: url('data:image/svg+xml,<svg width=&quot;100&quot; height=&quot;100&quot; xmlns=&quot;http://www.w3.org/2000/svg&quot;><rect width=&quot;100&quot; height=&quot;100&quot; fill=&quot;none&quot;/><circle cx=&quot;50&quot; cy=&quot;50&quot; r=&quot;40&quot; fill=&quot;rgba(255,255,255,0.05)&quot;/></svg>');
                            opacity: 0.3;"></div>
                <div style="display: flex; align-items: center; justify-content: space-between; position: relative;">
                    <div style="display: flex; align-items: center;">
                        <div>
                            <h1 style="margin:0; font-size:3em; color:white; font-weight:800; letter-spacing:-1px;">
                                Multilingual Audio Analysis with InterPARES-Audio
                            </h1>
                            <p style="margin:8px 0 0 0; color:rgba(255,255,255,0.95); font-size:1.15em; font-weight:500;">
                                Offline Audio Analysis and Summarization Demo
                            </p>
                        </div>
                    </div>
                    <img src="https://dlnlp.ai/img/InterPARES_Audio.jpg"
                         alt="InterPARES Audio"
                         style="height:80px; border-radius:12px; box-shadow:0 4px 6px rgba(0,0,0,0.2);"/>
                </div>
            </div>
        """)

        gr.HTML(
    """
    <div style="
        padding: 30px 34px;
        border-radius: 20px;
        background: white;
        border: none;
        box-shadow: 0 8px 24px rgba(37, 99, 235, 0.08);
        margin-bottom: 24px;
    ">
        <h2 style="
            margin: 0 0 22px 0;
            font-size: 2.1rem;
            font-weight: 800;
            color: #1e293b;
            line-height: 1.2;
        ">
            About This Demo
        </h2>

        <p style="
            margin: 0 0 20px 0;
            font-size: 1.08rem;
            line-height: 1.9;
            color: #334155;
        ">
            This demo showcases the results of a powerful pipeline designed to process long audio files with multiple speakers and languages:
        </p>

        <ol style="
            margin: 0;
            padding-left: 28px;
            color: #334155;
            font-size: 1.06rem;
            line-height: 1.9;
        ">
            <li style="margin-bottom: 14px;">
                <strong style="color:#1e293b;">Speaker Diarization:</strong>
                Identifies <em>who</em> spoke and <em>when</em>.
            </li>
            <li style="margin-bottom: 14px;">
                <strong style="color:#1e293b;">Multilingual Transcription:</strong>
                Transcribes each speaker's segment, automatically detecting the language.
            </li>
            <li style="margin-bottom: 0;">
                <strong style="color:#1e293b;">LLM Analysis:</strong>
                Uses a Large Language Model to summarize the transcript, extract action items, and generate key insights.
            </li>
        </ol>
    </div>
    """,
    elem_id="about-demo-box",
    container=False,
    padding=False,
    apply_default_css=False,
)


        with gr.Row():
            placeholder = "-- Please select an audio file --"
            audio_choices = list(AUDIO_FILES_INFO.keys()) if AUDIO_FILES_INFO else [placeholder]

            audio_selector = gr.Dropdown(
                choices=audio_choices,
                value=default_audio if default_audio else placeholder,
                label="Select Audio File",
            )

        with gr.Row(elem_classes="box-style", elem_id="main-content-row"):
            with gr.Column(scale=1):
                audio_player = gr.Audio(
                    label=default_audio_label,
                    value=default_audio_path,
                    interactive=False,
                    elem_id="audio_player",
                )
                transcript_display = gr.Textbox(
                    label="Audio Transcript",
                    max_lines=60,
                    interactive=False,
                    value=default_transcript,
                    elem_id="transcript_display",
                )

            with gr.Column(scale=2):
                lang_selector = gr.Dropdown(
                    choices=list(LANGUAGES.keys()),
                    label="Select Summary Language",
                    interactive=True if default_audio else False,
                    value=default_lang,
                    elem_id="lang_selector",
                )
                pdf_display = gr.HTML(
                    value=default_pdf_html,
                    label="PDF Summary",
                    elem_id="pdf_display",
                )

        outputs_list = [audio_player, transcript_display, lang_selector, pdf_display]

        audio_selector.change(
            fn=load_audio_and_transcript,
            inputs=audio_selector,
            outputs=outputs_list,
        )

        lang_selector.change(
            fn=load_summary_pdf_with_visibility,
            inputs=[audio_selector, lang_selector],
            outputs=pdf_display,
        )

    return demo


# --- Main Execution Block ---
if __name__ == "__main__":
    # Use a relative temp dir so Gradio can always serve it
    temp_dir = os.path.join(os.getcwd(), "gradio_temp")
    os.makedirs(temp_dir, exist_ok=True)
    os.environ["GRADIO_TEMP_DIR"] = temp_dir

    cache_dir = os.path.join(os.getcwd(), "gradio_cache")
    os.makedirs(cache_dir, exist_ok=True)

    AUDIO_FILES_INFO = discover_audio_files()

    if not AUDIO_FILES_INFO:
        print("=" * 50)
        print("ERROR: No valid audio files were found.")
        print("Expected format: demo_audio/<file_id>/<file_id>.wav")
        print("The app will start but the dropdown will be empty.")
        print("=" * 50)

    app_interface = create_interface()
    app_interface.launch(
        share=True,
        allowed_paths=[temp_dir, cache_dir],
    )