import os
from enum import StrEnum
from functools import cache
from logging import Logger
from time import sleep
from typing import cast

import requests
import yaml
from pydantic import BaseModel, Field
from pymultirole_plugins.v1.converter import ConverterBase, ConverterParameters
from pymultirole_plugins.v1.schema import Document, Sentence
from starlette.datastructures import UploadFile

logger = Logger("pymultirole")
WHISPERX_URL = os.getenv("WHISPERX_URL", None)


class InputFormat(StrEnum):
    YamlFile = "YamlFile"


class WhisperXParameters(ConverterParameters):
    whisperx_url: str = Field(
        WHISPERX_URL,
        description="Base URL of the WhisperX transcription server (e.g. 'http://localhost:8000'). "
        "Defaults to the WHISPERX_URL environment variable if set.",
    )
    input_format: InputFormat = Field(
        InputFormat.YamlFile,
        description="Format of the input file. Currently supports YAML files containing an audio URL "
        "and optional metadata (title, speakers, etc.).",
    )
    language: str = Field(
        None,
        description="ISO 639-1 language code of the audio (e.g. 'en', 'fr', 'de'). "
        "When not set, WhisperX will auto-detect the language.",
    )
    initial_prompt: str = Field(
        None,
        description="Optional text to condition the first transcription window. Useful for providing "
        "domain-specific vocabulary, acronyms, or context that helps improve recognition accuracy.",
        json_schema_extra={"extra": "advanced"},
    )
    temperature: float = Field(
        0,
        description="Sampling temperature for the transcription model. A value of 0 uses greedy decoding "
        "(most deterministic). Higher values (e.g. 0.2-0.8) increase randomness and may help with "
        "difficult audio but can reduce accuracy.",
        json_schema_extra={"extra": "advanced"},
    )
    batch_size: int = Field(
        32,
        description="Number of audio chunks to process in parallel. Higher values speed up transcription "
        "but require more GPU memory. Reduce this if you encounter out-of-memory errors.",
        json_schema_extra={"extra": "advanced"},
    )
    align_output: bool = Field(
        True,
        description="When enabled, runs a phoneme alignment model on the transcription output to produce "
        "accurate word-level timestamps. Disable to skip alignment and only get segment-level timestamps.",
        json_schema_extra={"extra": "advanced"},
    )
    diarization: bool = Field(
        True,
        description="When enabled, performs speaker diarization to identify and label different speakers "
        "in the audio. Each transcribed segment will be assigned a speaker ID (e.g. SPEAKER_00).",
        json_schema_extra={"extra": "advanced"},
    )
    group_adjacent_speakers: bool = Field(
        False,
        description="When enabled, merges consecutive transcript segments that share the same speaker "
        "into a single segment. Reduces output fragmentation for long monologues.",
        json_schema_extra={"extra": "advanced"},
    )
    min_speakers: int = Field(
        0,
        description="Minimum expected number of speakers in the audio. Set to 0 to let the diarization "
        "model determine the count automatically. Only used when diarization is enabled.",
        json_schema_extra={"extra": "advanced"},
    )
    max_speakers: int = Field(
        0,
        description="Maximum expected number of speakers in the audio. Set to 0 to let the diarization "
        "model determine the count automatically. Only used when diarization is enabled.",
        json_schema_extra={"extra": "advanced"},
    )
    return_embeddings: bool = Field(
        False,
        description="When enabled, the WhisperX server returns a representative embedding vector for each "
        "detected speaker. Useful for downstream speaker identification or clustering tasks.",
        json_schema_extra={"extra": "advanced"},
    )


class WhisperXConverter(ConverterBase):
    """WhisperX converter for audio transcription with speaker diarization.

    Processes YAML input files containing audio URLs and metadata,
    sends them to a WhisperX server for transcription, and returns
    documents with speaker-annotated text and timecodes.
    """

    def convert(self, source: UploadFile, parameters: ConverterParameters) -> list[Document]:
        params: WhisperXParameters = cast(WhisperXParameters, parameters)
        client = get_client(params.whisperx_url)
        if params.input_format == InputFormat.YamlFile:
            config = yaml.safe_load(source.file)

            bean = client.transcribe(config["url"], params)
            text = ""
            sentences = []
            for seg in bean["result"]["segments"]:
                speaker = seg.get("speaker", "SPEAKER_UNKNOWN")
                start = len(text)
                text = text + f"[{speaker}]: " + seg["text"] + "\n"
                end = len(text)
                sent = Sentence(
                    start=start,
                    end=end,
                    metadata={
                        "speaker": speaker,
                        "timecode": f"{seg['start']}:{seg['end']}",
                    },
                )
                sentences.append(sent)
            metadata = config.get("metadata", {})
            metadata["url"] = config["url"]
            doc = Document(
                identifier=config["url"],
                title=config.get("title", config["url"]),
                text=text,
                sentences=sentences,
                metadata=config.get("metadata", None),
            )
            return [doc]
        return [None]

    @classmethod
    def get_model(cls) -> type[BaseModel]:
        return WhisperXParameters


@cache
def get_client(base_url):
    client = WhisperXClient(base_url)
    return client


def is_success(job_bean):
    return job_bean and job_bean["status"] == "COMPLETED"


class WhisperXClient:
    def __init__(self, base_url):
        self.base_url = base_url[0:-1] if base_url.endswith("/") else base_url
        self.dsession = requests.Session()
        self.dsession.headers.update({"Accept": "application/json"})
        self.dsession.verify = False

    def transcribe(self, url: str, params: WhisperXParameters):
        try:
            resp = self.dsession.post(
                f"{self.base_url}/submit",
                json={
                    "options": {
                        "audio_url": url,
                        "language": params.language,
                        "initial_prompt": params.initial_prompt,
                        "temperature": params.temperature,
                        "batch_size": params.batch_size,
                        "align_output": params.align_output,
                        "diarization": params.diarization,
                        "group_adjacent_speakers": params.group_adjacent_speakers,
                        "min_speakers": params.min_speakers,
                        "max_speakers": params.max_speakers,
                        "return_embeddings": params.return_embeddings,
                    }
                },
                timeout=300,
            )
            if resp.ok:
                job_bean = resp.json()
                job_bean = self.wait_for_completion(job_bean)
                if is_success(job_bean):
                    return job_bean
                else:
                    logger.warning("Unsuccessful transcription: {url}")
                    resp.raise_for_status()
            else:
                logger.warning("Unsuccessful transcription: {url}")
                resp.raise_for_status()
        except BaseException as err:
            logger.warning("An exception was thrown!", exc_info=True)
            raise err

    def wait_for_completion(self, job_bean):
        if job_bean:
            while job_bean["status"] not in [
                "COMPLETED",
                "CANCELLED",
                "FAILED",
            ]:
                sleep(10)
                resp = self.dsession.get(f"{self.base_url}/status/{job_bean['uid']}")
                if resp.ok:
                    job_bean = resp.json()
                else:
                    logger.warning("Impossible to retrieve status of job: {job_bean['uid']}")
                    resp.raise_for_status()
        return job_bean
