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
    words: Optional[List["WordTimestamp"]] = None


@dataclass(frozen=True)
class WordTimestamp:
    start_seconds: float
    end_seconds: float
    text: str


@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:
                words = None
                if segment.words:
                    words = [
                        WordTimestamp(
                            start_seconds=word.start_seconds + start_seconds,
                            end_seconds=word.end_seconds + start_seconds,
                            text=word.text,
                        )
                        for word in segment.words
                    ]
                segments.append(
                    TranscriptionSegment(
                        start_seconds=segment.start_seconds + start_seconds,
                        end_seconds=segment.end_seconds + start_seconds,
                        text=segment.text,
                        speaker=segment.speaker,
                        words=words,
                    )
                )
        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:
            word_payload = dict(data)
            word_payload["timestamp_granularities[]"] = "word"
            response = self._post_once(
                word_payload, audio_file, filename, content_type
            )
            if response.status_code in (400, 422):
                audio_file.seek(0)
                response = self._post_once(data, audio_file, filename, content_type)
            response.raise_for_status()
            return response.json()
        except (httpx.HTTPError, ValueError) as exc:
            raise RuntimeError(f"MOSS transcription request failed: {exc}") from exc

    def _post_once(self, data: dict, audio_file, filename: str, content_type: str):
        return httpx.post(
            self.endpoint,
            data=data,
            files={"file": (filename, audio_file, content_type)},
            timeout=httpx.Timeout(self.timeout_seconds, connect=30),
        )


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,
                    words=_parse_word_timestamps(item.get("words")),
                )
            )
    return segments


def _parse_word_timestamps(raw_words) -> Optional[List[WordTimestamp]]:
    if not isinstance(raw_words, list):
        return None
    words = []
    for item in raw_words:
        if not isinstance(item, dict):
            continue
        try:
            start = float(item["start"])
            end = float(item["end"])
        except (KeyError, TypeError, ValueError):
            continue
        text = str(item.get("word") or "").strip()
        if text and end > start:
            words.append(WordTimestamp(start_seconds=start, end_seconds=end, text=text))
    return words or None


_PERIOD_BOUNDARY = re.compile(r"\.+")


def split_segment_by_periods(
    segment: TranscriptionSegment,
) -> List[TranscriptionSegment]:
    """Split one whisper segment into sentences at periods (".")."""
    text = segment.text.strip()
    if not text:
        return []
    pieces: List[str] = []
    cursor = 0
    for match in _PERIOD_BOUNDARY.finditer(text):
        piece = text[cursor:match.end()].strip()
        if piece.strip(".").strip():
            pieces.append(piece)
        cursor = match.end()
    tail = text[cursor:].strip()
    if tail:
        pieces.append(tail)
    if len(pieces) <= 1:
        return [segment] if pieces else []
    ends = _sentence_end_seconds(segment, pieces)
    sub_segments: List[TranscriptionSegment] = []
    start = segment.start_seconds
    for piece, end in zip(pieces, ends):
        if end > start:
            sub_segments.append(
                TranscriptionSegment(
                    start_seconds=start,
                    end_seconds=end,
                    text=piece,
                    speaker=segment.speaker,
                )
            )
        start = end
    return sub_segments


def _sentence_end_seconds(
    segment: TranscriptionSegment, pieces: List[str]
) -> List[float]:
    lengths = [len(piece) for piece in pieces]
    total = sum(lengths)
    if segment.words and len(segment.words) >= len(pieces):
        word_ends = _word_boundary_ends(segment, lengths)
        if _valid_boundaries(segment, word_ends):
            return word_ends + [segment.end_seconds]
    return _proportional_ends(segment, lengths) + [segment.end_seconds]


def _word_boundary_ends(
    segment: TranscriptionSegment, lengths: List[int]
) -> List[float]:
    targets = [sum(lengths[:count]) for count in range(1, len(lengths))]
    buffer = ""
    word_indexes: List[int] = []
    target_index = 0
    for word_index, word in enumerate(segment.words):
        buffer += (" " if buffer else "") + word.text
        while target_index < len(targets) and len(buffer) >= targets[target_index]:
            word_indexes.append(word_index)
            target_index += 1
    if len(word_indexes) != len(targets):
        return []
    return [segment.words[index].end_seconds for index in word_indexes]


def _valid_boundaries(
    segment: TranscriptionSegment, ends: List[float]
) -> bool:
    previous = segment.start_seconds
    for end in ends:
        if not segment.start_seconds < end < segment.end_seconds:
            return False
        if end <= previous:
            return False
        previous = end
    return True


def _proportional_ends(
    segment: TranscriptionSegment, lengths: List[int]
) -> List[float]:
    total = sum(lengths)
    ends = []
    cumulative = 0
    for length in lengths[:-1]:
        cumulative += length
        ends.append(
            segment.start_seconds
            + (cumulative / total) * (segment.end_seconds - segment.start_seconds)
        )
    return ends


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