"""MeetPilot Whisper + PyAnnote Diarization Gradio Space (ZeroGPU Compatible).

Exposes a Gradio web interface and named API endpoint (`/transcribe`)
powered by faster-whisper and pyannote.audio with WhisperX alignment.
"""

import logging
import os
from typing import Any, Dict, Optional

import gradio as gr
from handler import EndpointHandler

# Configure logging
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
logger = logging.getLogger("meetpilot-space")

# Optional ZeroGPU import (graceful fallback if running locally on standard CPU/GPU)
try:
    import spaces
    has_spaces = True
except ImportError:
    spaces = None
    has_spaces = False

# Initialize EndpointHandler in lazy mode so models are NOT loaded during import/startup
handler = EndpointHandler(lazy=True)


def _run_transcription(
    audio_path,
    min_speakers=None,
    max_speakers=None,
    language=None,
):
    """Core transcription worker called inside the active execution context."""
    if not audio_path or not os.path.exists(audio_path):
        raise gr.Error("Please provide a valid audio file.")

    hf_token = (
        os.getenv("HFTOKEN")
        or os.getenv("HF_TOKEN")
        or os.getenv("HUGGINGFACE_TOKEN")
        or os.getenv("HF_API_TOKEN")
    )
    if not hf_token:
        logger.warning("HFTOKEN is not set in Space Secrets. PyAnnote gated models may fail.")
    else:
        os.environ["HF_TOKEN"] = hf_token
        os.environ["HUGGING_FACE_HUB_TOKEN"] = hf_token

    # Sanitize speaker counts
    min_spk = int(min_speakers) if min_speakers is not None and min_speakers > 0 else None
    max_spk = int(max_speakers) if max_speakers is not None and max_speakers > 0 else None
    lang = language.strip() if language and language.strip() else None

    payload = {
        "inputs": audio_path,
        "parameters": {
            "min_speakers": min_spk,
            "max_speakers": max_spk,
            "language": lang,
        },
    }

    try:
        result = handler(payload)
        return result
    except Exception as exc:
        logger.exception("Error during meeting transcription: %s", exc)
        raise gr.Error(f"Transcription failed: {str(exc)}")


# Apply @spaces.GPU decorator if running in a Hugging Face ZeroGPU environment
if has_spaces and spaces is not None:
    @spaces.GPU(duration=120)
    def transcribe(audio_file, min_speakers=None, max_speakers=None, language=None):
        return _run_transcription(audio_file, min_speakers, max_speakers, language)
else:
    def transcribe(audio_file, min_speakers=None, max_speakers=None, language=None):
        return _run_transcription(audio_file, min_speakers, max_speakers, language)


# Build Gradio Interface
with gr.Blocks(title="MeetPilot - Whisper + PyAnnote Speaker Diarization") as demo:
    gr.Markdown(
        """
        # 🎙️ MeetPilot AI: Speech Recognition & Speaker Diarization
        Transcribe audio and identify distinct speakers in a single pass using **faster-whisper** (`large-v3`) + **pyannote.audio** (`community-1`) with WhisperX alignment on **Hugging Face ZeroGPU**.
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            audio_input = gr.Audio(
                type="filepath",
                label="Audio Recording (WAV, MP3, M4A, etc.)",
            )
            with gr.Accordion("Advanced Diarization Options", open=False):
                min_spk_input = gr.Number(label="Minimum Speakers", precision=0, value=None)
                max_spk_input = gr.Number(label="Maximum Speakers", precision=0, value=None)
                lang_input = gr.Textbox(label="Language Code (Optional, e.g. 'en')", placeholder="auto")

            submit_btn = gr.Button("Transcribe & Diarize", variant="primary")

        with gr.Column(scale=1):
            output_json = gr.JSON(label="Diarized Transcript Segments")

    submit_btn.click(
        fn=transcribe,
        inputs=[audio_input, min_spk_input, max_spk_input, lang_input],
        outputs=output_json,
        api_name="transcribe",
    )

    gr.Markdown(
        """
        ### API Usage from MeetPilot Backend
        This Space exposes a named API endpoint `transcribe`. You can call it programmatically via `gradio_client` or HTTP POST:
        ```python
        from gradio_client import Client, handle_file

        client = Client("Subham05x/meetpilot-whisper-diarization-space", hf_token=os.environ["HF_TOKEN"])
        result = client.predict(
            audio_file=handle_file("meeting.wav"),
            min_speakers=None,
            max_speakers=None,
            language=None,
            api_name="/transcribe"
        )
        print(result) # {'segments': [...], 'language': 'en', 'duration': 12.5}
        ```
        """
    )

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