"""
Signal Loom AI — HuggingFace Spaces App
=======================================
Demo surface for Signal Loom AI.
Routes all transcription through the Signal Loom API — no local ML.

API: http://localhost:18790
Auth: Bearer token (demo key hardcoded; replace with your own at signalloomai.com)

Author: Aster Vale / Signal Loom AI
"""

import asyncio
import httpx
import json
import os
import time
from typing import Literal

import gradio as gr

# ── Configuration ────────────────────────────────────────────────────────────
API_BASE_URL = os.environ.get("SIGNAL_LOOM_API_URL", "http://localhost:18790")
# Demo API key — replace with your own key from signalloomai.com
API_KEY = os.environ.get("SIGNAL_LOOM_API_KEY", "slo_oEDM0c-d")
POLL_INTERVAL_SECONDS = 3
POLL_TIMEOUT_SECONDS = 300  # 5 minutes

HEADERS = {"Authorization": f"Bearer {API_KEY}"}

# Supported formats
SUPPORTED_EXTENSIONS = [
    ".mp3", ".wav", ".m4a", ".flac", ".ogg",
    ".mp4", ".mov", ".mkv", ".avi", ".webm",
]
MAX_FILE_SIZE_MB = 500

OUTPUT_FORMAT_OPTIONS = ["txt", "json", "srt", "vtt"]


# ── API Client ─────────────────────────────────────────────────────────────────

async def _submit_url(url: str, language: str | None = None) -> str:
    """Submit a URL for transcription. Returns job_id."""
    async with httpx.AsyncClient(timeout=60.0, follow_redirects=True) as client:
        params = {"url": url}
        if language:
            params["language"] = language
        resp = await client.post(
            f"{API_BASE_URL}/v1/transcribe",
            headers=HEADERS,
            data=params,
        )
        if resp.status_code == 429:
            detail = resp.json().get("detail", "Rate limit reached")
            raise gr.Error(
                f"Monthly audio minute limit reached. "
                f"Upgrade at signalloomai.com/#pricing — {detail}"
            )
        if resp.status_code == 401 or resp.status_code == 403:
            raise gr.Error("Invalid API key. This demo's key may have been revoked.")
        if resp.status_code == 400:
            raise gr.Error(f"Bad request: {resp.json().get('detail', resp.text)}")
        resp.raise_for_status()
        return resp.json()["job_id"]


async def _submit_file(file_path: str, language: str | None = None) -> str:
    """Upload a file for transcription. Returns job_id."""
    async with httpx.AsyncClient(timeout=120.0) as client:
        with open(file_path, "rb") as f:
            files = {"file": (os.path.basename(file_path), f)}
            data = {}
            if language:
                data["language"] = language
            resp = await client.post(
                f"{API_BASE_URL}/v1/transcribe",
                headers=HEADERS,
                files=files,
                data=data,
            )
        if resp.status_code == 413:
            raise gr.Error(f"File too large. Maximum size is {MAX_FILE_SIZE_MB}MB.")
        if resp.status_code == 429:
            detail = resp.json().get("detail", "Rate limit reached")
            raise gr.Error(
                f"Monthly audio minute limit reached. "
                f"Upgrade at signalloomai.com/#pricing — {detail}"
            )
        if resp.status_code == 401 or resp.status_code == 403:
            raise gr.Error("Invalid API key. This demo's key may have been revoked.")
        if resp.status_code == 400:
            raise gr.Error(f"Bad request: {resp.json().get('detail', resp.text)}")
        resp.raise_for_status()
        return resp.json()["job_id"]


async def _poll_job(job_id: str, status_md, timeout: int = POLL_TIMEOUT_SECONDS):
    """Poll job status, yielding status strings as it progresses."""
    start = time.monotonic()
    interval = POLL_INTERVAL_SECONDS
    last_status = "unknown"

    while time.monotonic() - start < timeout:
        await asyncio.sleep(interval)

        async with httpx.AsyncClient(timeout=30.0) as client:
            try:
                resp = await client.get(
                    f"{API_BASE_URL}/v1/status/{job_id}",
                    headers=HEADERS,
                )
                resp.raise_for_status()
            except httpx.HTTPStatusError as e:
                if e.response.status_code == 404:
                    raise gr.Error(f"Job {job_id} not found. It may have expired.")
                raise gr.Error(f"Status check failed: {e}")

        data = resp.json()
        status = data.get("status", "unknown")

        if status == last_status and status not in ("queued", "processing", "uploaded"):
            # No meaningful change; keep polling
            pass
        last_status = status

        if status == "queued":
            yield "status_md", "⏳ **Submitted — Queued for processing...**\n\n_Position unknown. Waiting for a worker slot..._", None
        elif status == "processing":
            elapsed = int(time.monotonic() - start)
            language = data.get("language") or "auto-detected"
            dur = data.get("duration_seconds")
            dur_str = f", ~{int(dur)}s audio" if dur else ""
            yield (
                "status_md",
                f"🔄 **Processing...** (`{elapsed}s elapsed`)\n\n"
                f"Model: `{data.get('model', 'whisper-large-v3-turbo')}` | "
                f"Language: `{language}`{dur_str}\n\n"
                f"_Polling every {interval}s — typically 1–3 min for short files._",
                None,
            )
            interval = min(interval * 1.2, 8.0)
        elif status == "completed":
            # Done — let the result fetch handle the final output
            return
        elif status in ("failed", "error"):
            failure = data.get("failure_code") or data.get("error") or "unknown"
            raise gr.Error(
                f"❌ Transcription failed{f' ({failure})' if failure else ''}.\n\n"
                f"Try a different file/URL, or check signalloomai.com for support."
            )
        elif status in ("cancelled", "expired"):
            raise gr.Error(
                f"Job {status}. Results are retained for 7 days after completion. "
                f"Please submit a new transcription."
            )
        else:
            # Unknown / created / uploaded — keep waiting
            yield (
                "status_md",
                f"🔄 **Status: `{status}`** — waiting for processing...\n\n"
                f"_Still starting up... ({(time.monotonic() - start):.0f}s elapsed)_",
                None,
            )

    raise gr.Error(
        f"⏱ Transcription timed out after {timeout}s. "
        f"The file may be too long or the service is under load. "
        f"Check job status at signalloomai.com or try again shortly."
    )


async def _fetch_result(job_id: str) -> dict:
    """Fetch the final transcript result."""
    async with httpx.AsyncClient(timeout=30.0) as client:
        resp = await client.get(
            f"{API_BASE_URL}/v1/result/{job_id}",
            headers=HEADERS,
        )
        if resp.status_code == 202:
            raise gr.Error(f"Job {job_id} returned unexpected 202 (still processing). Try again shortly.")
        if resp.status_code == 404:
            raise gr.Error(f"Result for job {job_id} not found. It may have expired.")
        if resp.status_code == 410:
            raise gr.Error(f"Result for job {job_id} has expired. Results are retained for 7 days after completion.")
        resp.raise_for_status()
        return resp.json()


def _build_outputs(result: dict, requested_formats: list[str]) -> dict:
    """Build the output dict matching the expected format options."""
    outputs = {}

    transcript = result.get("transcript") or {}
    raw_text = result.get("raw_text") or ""
    srt_text = result.get("srt") or ""
    vtt_text = result.get("vtt") or ""

    # TXT
    if "txt" in requested_formats:
        if raw_text:
            outputs["txt"] = raw_text.strip()
        elif transcript:
            # Reconstruct from segments
            text = " ".join(seg.get("text", "").strip() for seg in transcript.get("segments", []) if seg.get("text"))
            outputs["txt"] = text.strip()
        else:
            outputs["txt"] = raw_text or "(no transcript available)"

    # JSON
    if "json" in requested_formats:
        outputs["json"] = {
            "status": "success",
            "job_id": result.get("job_id"),
            "language": transcript.get("language", "unknown") if isinstance(transcript, dict) else "unknown",
            "duration_seconds": transcript.get("duration_seconds") if isinstance(transcript, dict) else None,
            "segments": [
                {
                    "start_seconds": round(seg.get("start", seg.get("start_time", 0)), 2),
                    "end_seconds": round(seg.get("end", seg.get("end_time", 0)), 2),
                    "speaker": seg.get("speaker", "UNKNOWN"),
                    "text": seg.get("text", "").strip(),
                }
                for seg in (transcript.get("segments", []) if isinstance(transcript, dict) else [])
                if seg.get("text", "").strip()
            ],
        }

    # SRT
    if "srt" in requested_formats:
        outputs["srt"] = srt_text

    # VTT
    if "vtt" in requested_formats:
        outputs["vtt"] = vtt_text

    return outputs


# ── Pipeline (generator — yields (output_key, value, _) for Gradio streaming) ──

async def run_pipeline(
    audio_file: str | None,
    url_input: str | None,
    language: str | None,
    output_formats: list[str],
) -> tuple:
    """
    Full transcription pipeline via Signal Loom API.
    Accepts file upload OR URL ingestion. Yields status/outputs for Gradio streaming.
    Returns tuple: (status_md_text, txt_output, json_output, srt_output, vtt_output)
    """
    requested_formats = output_formats or ["txt", "json"]

    # ── Step 0: validate inputs ───────────────────────────────────────────────
    if url_input and url_input.strip():
        source_label = f"URL"
    elif audio_file:
        source_label = f"File: {os.path.basename(audio_file)}"
        # Check file size (best-effort — actual limit enforced by API)
        try:
            size_mb = os.path.getsize(audio_file) / (1024 * 1024)
            if size_mb > MAX_FILE_SIZE_MB:
                raise gr.Error(
                    f"File too large ({size_mb:.0f}MB). Maximum is {MAX_FILE_SIZE_MB}MB.\n\n"
                    f"For larger files, use signalloomai.com API."
                )
        except OSError:
            pass
    else:
        raise gr.Error("Please paste a URL or upload a file first.")

    # ── Step 1: submit ────────────────────────────────────────────────────────
    yield (
        "status_md",
        f"🚀 **Submitting...** `{source_label}`",
        None, None, None, None,
    )

    try:
        if url_input and url_input.strip():
            job_id = await _submit_url(url_input.strip(), language)
        elif audio_file:
            job_id = await _submit_file(audio_file, language)
        else:
            return
    except gr.Error:
        raise
    except httpx.HTTPStatusError as e:
        try:
            detail = e.response.json().get("detail", e.response.text)
        except Exception:
            detail = str(e)
        raise gr.Error(f"Submission failed: {detail}")
    except Exception as e:
        raise gr.Error(f"Submission failed: {e}")

    # Immediate acknowledgement state
    yield (
        "status_md",
        f"✅ **Submitted — Transcription in progress...**\n\n"
        f"Job ID: `{job_id}`\n\n"
        f"⏱ Processing typically takes 1–3 minutes depending on length.\n\n"
        f"_Results are not stored long-term — this is a stateless demo._",
        None, None, None, None,
    )

    # ── Step 2: poll ───────────────────────────────────────────────────────────
    try:
        async for update in _poll_job(job_id, None):
            if update is None:
                continue
            key, val = update[0], update[1]
            if key == "status_md":
                yield val, None, None, None, None, None
    except gr.Error:
        raise
    except Exception as e:
        raise gr.Error(f"Polling failed: {e}")

    # ── Step 3: fetch result ──────────────────────────────────────────────────
    yield (
        "status_md",
        f"📥 **Fetching transcript...**",
        None, None, None, None,
    )

    try:
        result = await _fetch_result(job_id)
    except gr.Error:
        raise
    except httpx.HTTPStatusError as e:
        try:
            detail = e.response.json().get("detail", e.response.text)
        except Exception:
            detail = str(e)
        raise gr.Error(f"Failed to fetch result: {detail}")
    except Exception as e:
        raise gr.Error(f"Failed to fetch result: {e}")

    # ── Step 4: format and return ───────────────────────────────────────────────
    outputs = _build_outputs(result, requested_formats)

    transcript = result.get("transcript") or {}
    segments = transcript.get("segments", []) if isinstance(transcript, dict) else []
    num_speakers = len(set(s.get("speaker") for s in segments if s.get("speaker") and s.get("speaker") != "UNKNOWN")) or "N/A"
    lang = transcript.get("language", "unknown") if isinstance(transcript, dict) else "unknown"
    dur = transcript.get("duration_seconds") if isinstance(transcript, dict) else None

    summary = (
        f"✅ **Done!** `{len(segments)}` segments"
        f" | Lang: `{lang}`"
        f" | Speakers: `{num_speakers}`"
        + (f" | Duration: `{int(dur)}s`" if dur else "")
    )
    cta = (
        "\n\n---\n"
        "🚀 **Want API access?** "
        "[Get your free key at signalloomai.com](https://signalloomai.com/signup.html) "
        "— 100 min/month free, no credit card required."
    )

    yield (
        "status_md",
        summary + cta,
        outputs.get("txt"),
        outputs.get("json"),
        outputs.get("srt"),
        outputs.get("vtt"),
    )


# ── Gradio Interface ─────────────────────────────────────────────────────────

_DARK_CSS = """
    /* Dark ambient background */
    body { background: #0f1117; }
    /* Minimal dark chrome */
    .gradio-container { background: transparent; }
    /* Card-like panels */
    .panel-box { background: #161b27; border: 1px solid #1e2535; border-radius: 8px; padding: 16px; }
    /* Status highlight */
    .status-ready { color: #34d399; font-weight: 600; }
    .status-processing { color: #fbbf24; font-weight: 600; }
    .status-error { color: #f87171; font-weight: 600; }
    /* Compact markdown in output */
    .json-output pre { font-size: 12px; }
    /* Links */
    a { color: #818cf8; text-decoration: none; }
    a:hover { text-decoration: underline; }
"""

with gr.Blocks(
    title="Signal Loom AI",
    css=_DARK_CSS,
) as demo:

    gr.Markdown(
        "## ✦ Signal Loom AI\n"
        "**Structured media intelligence — transcription, speaker diarization, and machine-readable output.**\n\n"
        "⚡ **Powered by Signal Loom API** — MLX Whisper on Apple Silicon.\n\n"
        "| | |\n"
        "|---|---|\n"
        "| **Free tier** | 100 min/month |\n"
        "| **Output** | TXT · JSON · SRT · VTT |\n"
        "| **Sources** | YouTube, direct upload, 30+ platforms |\n\n"
        "📌 **This is a demo surface.** For API access and production use → **[signalloomai.com](https://signalloomai.com)**\n\n"
        "---\n"
    )

    with gr.Row(equal_height=False):
        # ── LEFT: Inputs ────────────────────────────────────────────────────────
        with gr.Column(scale=1, min_width=320):
            gr.Markdown("### 🔗 Option 1: Paste a URL")
            gr.Markdown(
                "_YouTube, direct links, Google Drive, Dropbox, and more._\n\n"
                "<span style='font-size: 12px; color: #6b7280;'>"
                "Note: Some YouTube videos may be blocked by platform anti-bot protections. "
                "Direct uploads are the most reliable."
                "</span>"
            )
            url_input = gr.Textbox(
                label="Media URL",
                placeholder="https://www.youtube.com/watch?v=...",
                lines=2,
                info="We'll download and transcribe automatically",
            )

            gr.Markdown("### 📁 Option 2: Upload a File")
            gr.Markdown(
                f"**Supported:** mp3 · wav · m4a · flac · ogg · mp4 · mov · mkv · avi · webm  \n"
                f"**Max size:** {MAX_FILE_SIZE_MB}MB  \n"
                "<span style='font-size: 12px; color: #6b7280;'>"
                "Results are not stored long-term — stateless demo surface."
                "</span>"
            )
            audio_input = gr.File(
                label="Upload Audio or Video",
                file_types=SUPPORTED_EXTENSIONS,
                type="filepath",
            )

            with gr.Accordion("⚙️ Options", open=False):
                language = gr.Dropdown(
                    label="Language (auto-detect if blank)",
                    choices=["", "en", "es", "fr", "de", "it", "pt", "zh", "ja", "ko", "ar", "ru"],
                    value="",
                )
                output_formats = gr.CheckboxGroup(
                    label="Output Formats",
                    choices=OUTPUT_FORMAT_OPTIONS,
                    value=["txt", "json"],
                )

            run_btn = gr.Button("▶  Run Transcription", variant="primary", size="lg")

        # ── RIGHT: Outputs ──────────────────────────────────────────────────────
        with gr.Column(scale=2, min_width=400):
            status_md = gr.Markdown(
                "**Ready.** Paste a URL or upload a file, then click **Run Transcription**.\n\n"
                "⏱ Processing typically takes 1–3 minutes depending on length and queue depth.",
            )

            gr.Markdown("#### 📄 Transcript (TXT)")
            txt_output = gr.Textbox(label=None, lines=8, interactive=False, show_label=False)

            gr.Markdown("#### 📊 Structured JSON")
            json_output = gr.JSON(label=None, show_label=False)

            gr.Markdown("#### 📝 Subtitles (SRT)")
            srt_output = gr.Textbox(label=None, lines=6, interactive=False, show_label=False)

            gr.Markdown("#### 🎬 Captions (VTT)")
            vtt_output = gr.Textbox(label=None, lines=6, interactive=False, show_label=False)

    # Footer
    gr.Markdown(
        "---\n"
        "**Signal Loom AI** — _Media In. Machine Intelligence Out._  \n"
        "This HuggingFace Space is a **demo surface**. For full API access, billing, and production use → "
        "[**signalloomai.com**](https://signalloomai.com) | "
        "[Get free API key](https://signalloomai.com/signup.html) | "
        "[HuggingFace Org](https://huggingface.co/signalloom-ai)"
    )

    # ── Event binding ─────────────────────────────────────────────────────────
    run_btn.click(
        fn=run_pipeline,
        inputs=[audio_input, url_input, language, output_formats],
        outputs=[status_md, txt_output, json_output, srt_output, vtt_output],
        queue=True,
        concurrency_limit=None,
    )


# ── Local launch ─────────────────────────────────────────────────────────────
if __name__ == "__main__":
    import sys
    port = int(sys.argv[1]) if len(sys.argv) > 1 else 18791
    demo.launch(server_port=port, server_name="0.0.0.0")