"""
Transcript processing using Groq Whisper API
"""

from __future__ import annotations

import asyncio
import os
from dataclasses import dataclass
from typing import Callable, Literal, Optional, TypedDict

from .media_manager import MediaFileManager
from .settings import settings
from .utils import (
    ChunkSlot,
    MediaChunk,
    RetryConfig,
    TempPath,
    delete_tmp_files,
    logger,
    process_time,
    retry,
    to_sexagesimal,
)


class TranscriptSegment(TypedDict):
    """Transcript segment from Whisper"""

    avg_logprob: float
    compression_ratio: float
    end: float
    id: int
    no_speech_prob: float
    seek: int
    start: float
    temperature: float
    text: str
    tokens: list[int]


@dataclass(frozen=True)
class ProcessTranscriptResult:
    """Transcript result with timing"""

    start: float
    end: float
    transcript: str

    def __repr__(self):
        return f"""
        {self.start} -> {self.end}

        {self.transcript}
        """


ReturnValue = Literal["text", "vtt", "srt"]


def get_groq_api_key() -> str:
    """Get Groq API key from environment"""
    api_key = os.environ.get("GROQ_API_KEY")
    if not api_key:
        raise ValueError("GROQ_API_KEY environment variable is required for transcription")
    return api_key


class ProcessTranscript(MediaFileManager):
    """Process transcripts from media files using Groq Whisper"""

    chunk_duration = 60 * 3  # 3 minutes
    return_value: ReturnValue | None

    def __init__(self, file_path: str, return_value: ReturnValue = "text"):
        super().__init__(file_path)

        from groq import Groq

        api_key = get_groq_api_key()
        self.groq_client = Groq(api_key=api_key)
        self.return_value = return_value

        self.concurrency = self.max_workers
        self.ffmpeg_concurrency = min(4, self.max_workers // 2)

    async def __aenter__(self) -> "ProcessTranscript":
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb) -> bool:
        return False

    async def _split_media_file_to_chunks(self) -> list[MediaChunk]:
        """Split media file into chunks for processing"""
        if self.chunk_duration > self.duration:
            return [MediaChunk(path=self.file_path, start=0, end=self.duration)]

        semaphore = asyncio.Semaphore(self.ffmpeg_concurrency)

        async def _handler(slot: ChunkSlot) -> MediaChunk:
            async with semaphore:
                file_path = TempPath.get_path(ext=self.file_ext)
                chunk_path = await self._clip_media_async(
                    slot.start,
                    slot.end,
                    file_path,
                    use_audio=True,
                )
                return MediaChunk(path=chunk_path, start=slot.start, end=slot.end)

        slices = self._slice_media_file(self.chunk_duration)
        results = await asyncio.gather(
            *[_handler(s) for s in slices],
            return_exceptions=True,
        )

        clean_results: list[MediaChunk] = []
        errors: list[BaseException] = []
        for r in results:
            if isinstance(r, MediaChunk):
                clean_results.append(r)
            elif isinstance(r, BaseException):
                errors.append(r)
                logger.error("Chunk split failed: %s", r)

        if len(slices) != len(clean_results):
            detail = "; ".join(str(e) for e in errors) if errors else "unknown"
            raise Exception(f"Failed to split media into chunks: {detail}")

        return clean_results

    async def process(
        self,
        on_chunk: Optional[Callable[[str], None]] = None,
    ) -> str:
        """
        Extract transcript from media file with controlled concurrency
        """
        logger.info(f"Processing transcript for {self.file_path}")

        semaphore = asyncio.Semaphore(self.concurrency)

        async def _process_chunk(chunk: MediaChunk) -> ProcessTranscriptResult:
            async with semaphore:
                try:
                    transcript = await self._get_transcript(chunk.path, chunk.start)
                    result = ProcessTranscriptResult(chunk.start, chunk.end, transcript)
                    if on_chunk:
                        on_chunk(str(result))
                    return result
                except Exception as e:
                    logger.error(f"Error processing chunk: {e}")
                    raise e
                finally:
                    delete_tmp_files([chunk.path])

        # Split the media file into chunks and sort them
        results = await self._split_media_file_to_chunks()
        sorted_results = sorted(results, key=lambda v: v.start)

        # Process chunks concurrently
        gathered_results = await asyncio.gather(
            *[_process_chunk(c) for c in sorted_results],
            return_exceptions=True,
        )
        valid_results: list[ProcessTranscriptResult] = []
        chunk_errors: list[BaseException] = []
        for r in gathered_results:
            if isinstance(r, ProcessTranscriptResult):
                valid_results.append(r)
            elif isinstance(r, BaseException):
                chunk_errors.append(r)

        if len(valid_results) != len(sorted_results):
            detail = "; ".join(str(e) for e in chunk_errors) if chunk_errors else "unknown"
            raise Exception(f"Transcript extraction failed for {len(chunk_errors)} chunk(s): {detail}")

        # Sort results by start time and combine transcripts
        valid_results = sorted(valid_results, key=lambda v: v.start)
        transcript = " ".join(v.transcript for v in valid_results).strip()
        if not transcript:
            raise Exception("No transcript content generated")

        if self.return_value == "srt":
            return self._vtt_to_srt(transcript)

        return transcript

    @process_time()
    async def _run_transcription_groq(
        self,
        file_path: str,
        response_format: Literal["text", "json", "verbose_json"] = "verbose_json",
    ):
        """
        Run transcription using Groq Whisper API
        """

        def _read_audio_file_sync(fp: str) -> tuple[str, bytes]:
            with open(fp, "rb") as audio_file:
                return (fp, audio_file.read())

        file_ref = await asyncio.to_thread(_read_audio_file_sync, file_path)

        @retry(RetryConfig(max_retries=2, delay=3, backoff=1.5))
        def _transcribe():
            try:
                return self.groq_client.audio.transcriptions.create(
                    file=file_ref,
                    model=settings.whisper_models[0],
                    response_format=response_format,
                )
            except Exception as e:
                logger.info("Retrying transcription with fallback model due to error: %s", e)
                return self.groq_client.audio.transcriptions.create(
                    file=file_ref,
                    model=settings.whisper_models[1],
                    response_format=response_format,
                )

        return await asyncio.to_thread(_transcribe)

    async def _get_transcript(self, file_path: str, time_offset=0.0) -> str:
        """
        Get transcript for a media file chunk
        """
        transcription = await self._run_transcription_groq(
            file_path,
            response_format="text" if self.return_value == "text" else "verbose_json",
        )

        if self.return_value == "text":
            if isinstance(transcription, str) and transcription:
                return transcription
            elif transcription.text:
                return transcription.text
            raise Exception("text field is None in transcription result")

        result = transcription.model_dump()
        segments = result.get("segments")
        if not segments:
            raise Exception("No segments returned")

        return self._segment_to_vtt(segments, time_offset)

    def _segment_to_vtt(
        self,
        segments: list[TranscriptSegment],
        time_offset=0.0,
    ) -> str:
        """Convert segments to VTT format"""
        vtt = "" if time_offset > 0.0 else "WEBVTT\n\n"
        for segment in segments:
            start = to_sexagesimal(round(segment["start"] + time_offset, 3))
            end = to_sexagesimal(round(segment["end"] + time_offset, 3))
            text = segment["text"]
            vtt += f"{start} --> {end}\n{text}\n\n"
        return vtt

    def _vtt_to_srt(self, vtt: str) -> str:
        """Convert VTT to SRT format"""
        vtt = vtt.strip()
        if vtt.startswith("WEBVTT"):
            vtt = vtt.replace("WEBVTT", "").strip()

        srt = ""
        for i, line in enumerate(vtt.split("\n\n")):
            if not line.strip():
                continue
            time_line, text = line.split("\n", 1)
            start, end = time_line.split(" --> ")
            start = start.replace(".", ",")
            end = end.replace(".", ",")
            srt += f"{i + 1}\n{start} --> {end}\n{text.strip()}\n\n"

        return srt.strip()


@process_time()
async def get_video_transcript(
    video_path: str,
    format: ReturnValue = "text",
    on_chunk: Optional[Callable[[str], None]] = None,
) -> str:
    """Extract transcript from a video file, with option to invoke *on_chunk* with text as it arrives.

    Args:
        video_path: Path to the video or audio file.
        format: Output format ('text', 'vtt', or 'srt').
        on_chunk: Optional callback called with each text chunk as it becomes available.

    Returns:
        The full transcript string.
    """
    return_value = format if format in ("text", "vtt", "srt") else "text"
    async with ProcessTranscript(video_path, return_value) as proc:
        return await proc.process(on_chunk)
    return ""
