import io
import mimetypes
import re
import wave
from dataclasses import dataclass
from pathlib import Path
from typing import List, Optional, Protocol

import httpx
import numpy as np

from .audio_metrics import AudioAnalysisError, decode_audio_mono, speech_frame_mask


# Whisper was trained on 30-second clips. Sending longer audio in a single request
# makes the model enter a repetition loop (every later segment becomes the same
# sentence while timestamps keep advancing), so long audio is split into short
# chunks at silence gaps and each chunk is transcribed independently.
CHUNK_SECONDS = 30.0
MIN_CHUNK_SECONDS = 5.0
MIN_SILENCE_SECONDS = 0.5
MIN_TAIL_SECONDS = 4.0
# Whisper's own repetition-loop detector: segments that compress this well are
# almost certainly hallucinated repeated text (healthy speech is ~1.3-1.8).
COMPRESSION_RATIO_LIMIT = 2.4


@dataclass(frozen=True)
class TranscriptionSegment:
    start_seconds: float
    end_seconds: float
    text: str
    speaker: Optional[str] = None


@dataclass(frozen=True)
class Transcript:
    text: str
    segments: List[TranscriptionSegment]


class Transcriber(Protocol):
    @property
    def available(self) -> bool:
        ...

    def transcribe(self, audio_path: Path, language: Optional[str] = None) -> Transcript:
        ...


class MossTranscriber:
    def __init__(
        self,
        endpoint: str,
        model: str,
        timeout_seconds: float = 1800,
        max_new_tokens: int = 65536,
    ):
        self.endpoint = endpoint.strip()
        self.model = model
        self.timeout_seconds = timeout_seconds
        self.max_new_tokens = max_new_tokens

    @property
    def available(self) -> bool:
        return bool(self.endpoint)

    def transcribe(self, audio_path: Path, language: Optional[str] = None) -> Transcript:
        if not self.available:
            raise RuntimeError("MOSS transcription is not configured on this server.")
        duration = _wav_duration_seconds(audio_path)
        if duration is None or duration > CHUNK_SECONDS + 2.0:
            try:
                samples, sample_rate = decode_audio_mono(audio_path)
            except AudioAnalysisError:
                return self._transcribe_single(audio_path, language)
            if samples.size / sample_rate > CHUNK_SECONDS + 2.0:
                try:
                    return self._transcribe_chunked(samples, sample_rate, language)
                except AudioAnalysisError:
                    return self._transcribe_single(audio_path, language)
        return self._transcribe_single(audio_path, language)

    def _transcribe_single(
        self, audio_path: Path, language: Optional[str]
    ) -> Transcript:
        content_type = mimetypes.guess_type(audio_path.name)[0] or "application/octet-stream"
        with audio_path.open("rb") as audio_file:
            payload = self._post_audio(audio_file, audio_path.name, content_type, language)
        raw_text = str(payload.get("text") or "").strip()
        segments = _parse_json_segments(payload.get("segments"))
        if not segments:
            segments = _parse_compact_segments(raw_text)
        plain_text = " ".join(segment.text for segment in segments).strip() or raw_text
        return Transcript(text=plain_text, segments=segments)

    def _transcribe_chunked(
        self, samples: np.ndarray, sample_rate: int, language: Optional[str]
    ) -> Transcript:
        segments: List[TranscriptionSegment] = []
        for start_seconds, end_seconds in _plan_chunks(samples, sample_rate):
            chunk = samples[
                int(round(start_seconds * sample_rate)) : int(round(end_seconds * sample_rate))
            ]
            payload = self._post_audio(
                io.BytesIO(_wav_bytes(chunk, sample_rate)),
                "chunk.wav",
                "audio/wav",
                language,
            )
            chunk_segments = _parse_json_segments(
                payload.get("segments"), compression_limit=COMPRESSION_RATIO_LIMIT
            )
            if not chunk_segments:
                chunk_segments = _parse_compact_segments(str(payload.get("text") or "").strip())
            for segment in chunk_segments:
                segments.append(
                    TranscriptionSegment(
                        start_seconds=segment.start_seconds + start_seconds,
                        end_seconds=segment.end_seconds + start_seconds,
                        text=segment.text,
                        speaker=segment.speaker,
                    )
                )
        segments.sort(key=lambda item: (item.start_seconds, item.end_seconds))
        plain_text = " ".join(segment.text for segment in segments).strip()
        return Transcript(text=plain_text, segments=segments)

    def _post_audio(
        self,
        audio_file,
        filename: str,
        content_type: str,
        language: Optional[str],
    ) -> dict:
        data = {
            "model": self.model,
            "response_format": "verbose_json",
            "temperature": "0",
            "condition_on_previous_text": "false",
            "max_new_tokens": str(self.max_new_tokens),
        }
        if language:
            data["language"] = language

        try:
            response = httpx.post(
                self.endpoint,
                data=data,
                files={"file": (filename, audio_file, content_type)},
                timeout=httpx.Timeout(self.timeout_seconds, connect=30),
            )
            response.raise_for_status()
            return response.json()
        except (httpx.HTTPError, ValueError) as exc:
            raise RuntimeError(f"MOSS transcription request failed: {exc}") from exc


def _parse_json_segments(
    raw_segments, compression_limit: Optional[float] = None
) -> List[TranscriptionSegment]:
    if not isinstance(raw_segments, list):
        return []
    segments = []
    for item in raw_segments:
        if not isinstance(item, dict):
            continue
        try:
            start = float(item["start"])
            end = float(item["end"])
        except (KeyError, TypeError, ValueError):
            continue
        compression_ratio = item.get("compression_ratio")
        if (
            compression_limit is not None
            and isinstance(compression_ratio, (int, float))
            and compression_ratio > compression_limit
        ):
            continue
        text = str(item.get("text") or "").strip()
        if text and end > start >= 0:
            speaker = item.get("speaker")
            segments.append(
                TranscriptionSegment(
                    start_seconds=start,
                    end_seconds=end,
                    text=text,
                    speaker=str(speaker) if speaker is not None else None,
                )
            )
    return segments


def _wav_duration_seconds(path: Path) -> Optional[float]:
    try:
        with wave.open(str(path), "rb") as wav:
            if wav.getframerate() > 0:
                return wav.getnframes() / wav.getframerate()
    except (wave.Error, OSError):
        pass
    return None


def _wav_bytes(samples: np.ndarray, sample_rate: int) -> bytes:
    pcm = np.clip(np.round(samples * 32767), -32768, 32767).astype("<i2")
    buffer = io.BytesIO()
    with wave.open(buffer, "wb") as wav:
        wav.setnchannels(1)
        wav.setsampwidth(2)
        wav.setframerate(sample_rate)
        wav.writeframes(pcm.tobytes())
    return buffer.getvalue()


def _plan_chunks(samples: np.ndarray, sample_rate: int) -> List[tuple]:
    speech = speech_frame_mask(samples, sample_rate)
    frame_seconds = 0.03
    gap_frames = max(1, int(round(MIN_SILENCE_SECONDS / frame_seconds)))

    gaps: List[tuple] = []
    run_start: Optional[int] = None
    for index, is_speech in enumerate(speech):
        if not is_speech and run_start is None:
            run_start = index
        elif is_speech and run_start is not None:
            if index - run_start >= gap_frames:
                gaps.append((run_start * frame_seconds, index * frame_seconds))
            run_start = None
    if run_start is not None and len(speech) - run_start >= gap_frames:
        gaps.append((run_start * frame_seconds, len(speech) * frame_seconds))

    duration = samples.size / sample_rate
    cuts: List[float] = []
    cursor = 0.0
    while duration - cursor > CHUNK_SECONDS + 1e-6:
        target = cursor + CHUNK_SECONDS
        best_gap = None
        for gap_start, gap_end in gaps:
            if gap_start < cursor + MIN_CHUNK_SECONDS:
                continue
            if gap_start > target:
                break
            best_gap = (gap_start, gap_end)
        if best_gap is None:
            cut = min(target, duration)
        else:
            gap_start, gap_end = best_gap
            cut = min(target, (gap_start + gap_end) / 2.0)
        cuts.append(cut)
        cursor = cut

    boundaries = [0.0] + cuts + [duration]
    chunks = [
        (boundaries[index], boundaries[index + 1])
        for index in range(len(boundaries) - 1)
        if boundaries[index + 1] > boundaries[index] + 0.05
    ]
    if (
        len(chunks) > 1
        and chunks[-1][1] - chunks[-1][0] < MIN_TAIL_SECONDS
        and chunks[-1][1] - chunks[-2][0] <= CHUNK_SECONDS + MIN_TAIL_SECONDS
    ):
        chunks[-2] = (chunks[-2][0], chunks[-1][1])
        chunks.pop()
    return chunks


_COMPACT_SEGMENT = re.compile(
    r"\[(?P<start>\d+(?:\.\d+)?)\]"
    r"\[(?P<speaker>S\d+)\]"
    r"(?P<text>.*?)"
    r"\[(?P<end>\d+(?:\.\d+)?)\]",
    flags=re.DOTALL,
)


def _parse_compact_segments(text: str) -> List[TranscriptionSegment]:
    segments = []
    for match in _COMPACT_SEGMENT.finditer(text):
        start = float(match.group("start"))
        end = float(match.group("end"))
        segment_text = match.group("text").strip()
        if segment_text and end > start:
            segments.append(
                TranscriptionSegment(
                    start_seconds=start,
                    end_seconds=end,
                    text=segment_text,
                    speaker=match.group("speaker"),
                )
            )
    return segments