from __future__ import annotations

import json
import mimetypes
import os
from pathlib import Path
from typing import Any

import requests
from fastapi.responses import HTMLResponse, FileResponse
from gradio import Server
from gradio.data_classes import FileData

APP_ROOT = Path(__file__).parent
EXAMPLES_MANIFEST_PATH = APP_ROOT / "examples" / "manifest.json"
EXAMPLES_AUDIO_DIR = APP_ROOT / "examples" / "audio"
INDEX_HTML_PATH = APP_ROOT / "index.html"
TRANSCRIPTIONS_API_URL = "https://api.cohere.ai/compatibility/v1/audio/transcriptions"
DEFAULT_MODEL_ID = "cohere-transcribe-arabic"
DEFAULT_TIMEOUT_SECONDS = 180.0
DEFAULT_LANGUAGE = "ar"
MAX_AUDIO_FILE_BYTES = 25 * 1024 * 1024
MAX_AUDIO_FILE_LABEL = "25 MB"
SUPPORTED_LANGUAGES = {"ar", "en"}

LANGUAGE_OPTIONS = [
    ("English", "en"),
    ("Arabic", "ar"),
]
LANGUAGE_LABELS = {code: label for label, code in LANGUAGE_OPTIONS}


# ──────────────────────────────────────────────────────────────────────────────
# Backend helpers (unchanged behavior from the original Gradio Blocks app)
# ──────────────────────────────────────────────────────────────────────────────


class TranscriptionError(Exception):
    pass


def get_api_key() -> str | None:
    api_key = os.getenv("COHERE_API_KEY", "").strip()
    return api_key or None


def get_model_id() -> str:
    model_id = os.getenv("COHERE_ASR_MODEL", DEFAULT_MODEL_ID).strip()
    return model_id or DEFAULT_MODEL_ID


def get_timeout_seconds() -> float:
    raw_value = os.getenv("COHERE_API_TIMEOUT_SECONDS", str(DEFAULT_TIMEOUT_SECONDS))
    try:
        return max(30.0, float(raw_value))
    except ValueError:
        return DEFAULT_TIMEOUT_SECONDS


def get_debug_raw_response_enabled() -> bool:
    value = os.getenv("COHERE_ASR_DEBUG_RAW_RESPONSE", "").strip().lower()
    return value in {"1", "true", "yes", "on"}


def display_language(language_code: str | None) -> str:
    if not language_code:
        return "Not returned"
    normalized = language_code.strip()
    base_code = normalized.split("-", 1)[0].lower()
    label = LANGUAGE_LABELS.get(base_code)
    if not label:
        return normalized
    if base_code == normalized.lower():
        return label
    return f"{label} ({normalized})"


def _coerce_float(value: Any) -> float | None:
    """Best-effort conversion of a timestamp field to seconds."""
    if value is None:
        return None
    try:
        return float(value)
    except (TypeError, ValueError):
        return None


def extract_segments(payload: Any) -> list[dict[str, Any]]:
    """Pull timed `segments` from a Cohere-style payload, if present.

    Returns a list of {start, end, text} dicts. Anything that doesn't have
    timing is dropped — the frontend falls back to a flat transcript when
    the list is empty.
    """
    out: list[dict[str, Any]] = []

    def _from(obj: Any) -> None:
        if not isinstance(obj, dict):
            return
        text = extract_transcript_text(obj)
        if not text:
            return
        start = _coerce_float(
            obj.get("start")
            or obj.get("start_time")
            or obj.get("startTime")
            or obj.get("begin")
        )
        end = _coerce_float(
            obj.get("end")
            or obj.get("end_time")
            or obj.get("endTime")
            or obj.get("finish")
        )
        if start is None and end is None:
            return
        if start is not None and end is not None and end < start:
            start, end = end, start
        out.append({"start": start, "end": end, "text": text})

    if isinstance(payload, dict):
        segments = payload.get("segments")
        if isinstance(segments, list):
            for seg in segments:
                _from(seg)

    return out


def extract_transcript_text(payload: Any) -> str:
    if isinstance(payload, str):
        return payload.strip()

    if isinstance(payload, dict):
        for key in ("text", "transcript"):
            value = payload.get(key)
            if isinstance(value, str) and value.strip():
                return value.strip()

        segments = payload.get("segments")
        if isinstance(segments, list):
            parts = [extract_transcript_text(segment) for segment in segments]
            joined = " ".join(part for part in parts if part)
            if joined:
                return joined.strip()

        for key in ("result", "data", "output"):
            nested = payload.get(key)
            text = extract_transcript_text(nested)
            if text:
                return text

    if isinstance(payload, list):
        parts = [extract_transcript_text(item) for item in payload]
        joined = " ".join(part for part in parts if part)
        return joined.strip()

    return ""


def format_api_error(response: requests.Response) -> str:
    message = ""
    try:
        payload = response.json()
    except ValueError:
        payload = response.text.strip()

    if isinstance(payload, dict):
        error = payload.get("error")
        if isinstance(error, dict):
            message = str(error.get("message") or error.get("type") or "")
        elif isinstance(error, str):
            message = error

        if not message:
            for key in ("message", "detail", "error_description"):
                value = payload.get(key)
                if isinstance(value, str) and value.strip():
                    message = value.strip()
                    break
    elif isinstance(payload, str):
        message = payload

    if not message:
        message = "Unexpected API response."

    return f"Transcription request failed ({response.status_code}): {message}"


def validate_audio_file(audio_path: str) -> Path:
    audio_file_path = Path(audio_path)
    try:
        file_size = audio_file_path.stat().st_size
    except OSError as exc:
        raise TranscriptionError("The selected audio file could not be read.") from exc

    if file_size > MAX_AUDIO_FILE_BYTES:
        raise TranscriptionError(
            f"Audio files must be {MAX_AUDIO_FILE_LABEL} or smaller. "
            "Please upload a shorter clip or a more compressed file."
        )

    return audio_file_path


def call_transcriptions_api(audio_path: str, language: str) -> dict[str, Any]:
    api_key = get_api_key()
    if not api_key:
        raise TranscriptionError(
            "COHERE_API_KEY is not configured. Add it locally or in Hugging Face Space secrets."
        )

    audio_file_path = validate_audio_file(audio_path)
    mime_type = mimetypes.guess_type(audio_file_path.name)[0] or "application/octet-stream"
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Accept": "application/json",
    }

    with audio_file_path.open("rb") as audio_file:
        multipart_fields = [
            ("model", (None, get_model_id())),
            ("language", (None, language)),
            ("file", (audio_file_path.name, audio_file, mime_type)),
        ]
        response = requests.post(
            TRANSCRIPTIONS_API_URL,
            headers=headers,
            files=multipart_fields,
            timeout=get_timeout_seconds(),
        )

    if not response.ok:
        raise TranscriptionError(format_api_error(response))

    try:
        payload = response.json()
    except ValueError as exc:
        raise TranscriptionError("The transcription API returned a non-JSON response.") from exc

    if not isinstance(payload, dict):
        raise TranscriptionError("The transcription API returned an unexpected payload shape.")

    return payload


# ──────────────────────────────────────────────────────────────────────────────
# Example manifest helpers (frontend + serve recorded audio files)
# ──────────────────────────────────────────────────────────────────────────────


def load_example_manifest() -> list[dict[str, Any]]:
    if not EXAMPLES_MANIFEST_PATH.exists():
        return []

    try:
        payload = json.loads(EXAMPLES_MANIFEST_PATH.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError):
        return []

    raw_examples = payload.get("examples", [])
    if not isinstance(raw_examples, list):
        return []

    return [item for item in raw_examples if isinstance(item, dict)]


def resolve_example_audio_path(raw_path: str | None) -> Path | None:
    if not raw_path:
        return None

    candidate = Path(raw_path)
    if not candidate.is_absolute():
        candidate = APP_ROOT / candidate

    if candidate.exists():
        return candidate
    return None


def build_public_examples(examples: list[dict[str, Any]]) -> list[dict[str, Any]]:
    public: list[dict[str, Any]] = []
    for example in examples:
        language = str(example.get("language") or DEFAULT_LANGUAGE)
        base_language = language.split("-", 1)[0].lower()
        if base_language not in SUPPORTED_LANGUAGES:
            continue

        audio_path = resolve_example_audio_path(example.get("audio_path"))
        if example.get("enabled") is False or audio_path is None:
            continue

        public.append(
            {
                "id": str(example.get("id") or audio_path.stem),
                "label": str(example.get("label") or audio_path.stem),
                "language": base_language,
                "description": str(example.get("description") or ""),
                "audio_url": f"/audio/{audio_path.name}",
            }
        )
    return public


def resolve_uploaded_path(file_data: FileData | Any) -> str:
    """Pull a local filesystem path out of an uploaded FileData payload."""
    if isinstance(file_data, dict):
        path = file_data.get("path") or file_data.get("url")
    else:
        path = getattr(file_data, "path", None) or getattr(file_data, "url", None)

    if not path:
        raise TranscriptionError("No audio file was received. Please record or upload a clip.")
    return str(path)


# ──────────────────────────────────────────────────────────────────────────────
# gradio.Server app (extends FastAPI)
# ──────────────────────────────────────────────────────────────────────────────


app = Server()


def build_frontend_config_payload() -> dict[str, Any]:
    """Frontend bootstrap configuration, serialized into the HTML page.

    Rendered inline on '/' to avoid registering a custom ``GET /config``
    FastAPI route, which can collide with the deferred router that
    ``gradio.Server`` registers during ``app.launch()`` for the Gradio
    JS client's ``view_api`` / SSE handshake.
    """
    return {
        "api_ready": get_api_key() is not None,
        "debug_raw_response": get_debug_raw_response_enabled(),
        "max_file_size_label": MAX_AUDIO_FILE_LABEL,
        "default_language": DEFAULT_LANGUAGE,
        "languages": [{"label": label, "value": code} for label, code in LANGUAGE_OPTIONS],
        "examples": build_public_examples(load_example_manifest()),
    }


@app.get("/audio/{filename}")
async def serve_example_audio(filename: str):
    """Serve bundled FLEURS sample clips without using StaticFiles.

    Using a single-path FastAPI route avoids earlier-route-registration
    surprises with the locally-mounted gr.SSE prefixes — order is
    deterministic and ``FileResponse`` does less than ``StaticFiles``.
    The path is restricted to ``EXAMPLES_AUDIO_DIR`` so requests can't
    read arbitrary files off the Space.
    """
    if ".." in filename or filename.startswith("/") or "/" in filename:
        return HTMLResponse("Bad path", status_code=400)
    candidate = (EXAMPLES_AUDIO_DIR / filename).resolve()
    examples_root = EXAMPLES_AUDIO_DIR.resolve()
    if examples_root not in candidate.parents and candidate != examples_root:
        return HTMLResponse("Bad path", status_code=400)
    if not candidate.is_file():
        return HTMLResponse("Not found", status_code=404)
    return FileResponse(candidate)


@app.api(name="transcribe")
def transcribe(audio: FileData, language: str) -> dict[str, Any]:
    """Transcribe an audio clip via Cohere's compatibility transcription API.

    Goes through Gradio's queue, so concurrent users are serialized and the
    endpoint is callable from ``gradio_client``. Returns the full transcript
    plus any timed segments surfaced by the API — when segments carry
    timestamps, the frontend can sync-highlight words as the audio plays.
    When the API doesn't provide timing, ``segments`` is ``[]``.
    """
    base_language = (language or "").strip().split("-", 1)[0].lower()
    if base_language not in SUPPORTED_LANGUAGES:
        base_language = DEFAULT_LANGUAGE

    audio_path = resolve_uploaded_path(audio)

    try:
        payload = call_transcriptions_api(audio_path, base_language)
    except TranscriptionError as exc:
        raise RuntimeError(str(exc)) from exc
    except requests.Timeout as exc:
        raise RuntimeError(
            "The transcription request timed out. Try a shorter clip or increase the timeout."
        ) from exc
    except requests.RequestException as exc:
        raise RuntimeError(f"Unable to reach the transcription API: {exc}") from exc

    transcript = extract_transcript_text(payload)
    if not transcript:
        raise RuntimeError("The API returned a response, but no transcript text was found.")

    raw_json = json.dumps(payload, indent=2, ensure_ascii=False)

    return {
        "transcript": transcript,
        "segments": extract_segments(payload),
        "raw_json": raw_json,
    }


@app.api(name="status")
def status() -> dict[str, Any]:
    """Liveness/health probe — also demonstrates the second API endpoint."""
    return {
        "api_ready": get_api_key() is not None,
        "model": get_model_id(),
        "default_language": DEFAULT_LANGUAGE,
    }


@app.get("/", response_class=HTMLResponse)
async def homepage() -> HTMLResponse:
    if not INDEX_HTML_PATH.exists():
        return HTMLResponse(
            "<h1>index.html missing</h1><p>Place index.html next to app.py.</p>",
            status_code=500,
        )
    html = INDEX_HTML_PATH.read_text(encoding="utf-8")
    # Inline bootstrap config so the browser doesn't need to round-trip
    # through an extra FastAPI route at boot — fewer chances of colliding
    # with the deferred handlers ``gradio.Server`` wires up on launch.
    config_payload = json.dumps(build_frontend_config_payload(), ensure_ascii=False)
    bootstrap = (
        '<script id="app-config" type="application/json">'
        + config_payload.replace("</", "<\\/")  # defensive </script> escape
        + "</script>"
    )
    return HTMLResponse(html.replace("</head>", bootstrap + "</head>", 1))


if __name__ == "__main__":
    # Show stack traces instead of swallowing them.
    app.launch(show_error=True, server_port=7860)