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

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


_PUNCTUATION_BOUNDARY = re.compile(r"[.?]+|,+")
MIN_WORDS_FOR_COMMA_SPLIT = 50
_ABBREVIATION_RE = re.compile(
    r"(?i)\b(?:a\.m\.|p\.m\.|mr\.|mrs\.|ms\.|dr\.|prof\.|st\.|vs\.|etc\.|"
    r"e\.g\.|i\.e\.|no\.|approx\.|fig\.|inc\.|ltd\.|jr\.|sr\.|u\.s\.|u\.k\.)"
    r"(?![A-Za-z0-9])"
)
_SPEAKER_TAG_RE = re.compile(r"\[(S\d+)\]")


def _expand_speaker_segments(
    segment: TranscriptionSegment,
) -> List[TranscriptionSegment]:
    """Split a segment at inline ``[S01]`` speaker tags and strip the tags.

    MOSS sometimes embeds speaker tags directly in the text (rather than in
    the ``speaker`` JSON field). Each tag marks the start of a new speaker
    turn, so the text is split there and the tags are removed from the
    visible transcription.
    """
    text = segment.text
    matches = list(_SPEAKER_TAG_RE.finditer(text))
    if not matches:
        return [segment]

    sub_segments: List[TranscriptionSegment] = []
    total = len(matches)
    duration = segment.end_seconds - segment.start_seconds
    for index, match in enumerate(matches):
        start_pos = match.end()
        if index + 1 < total:
            end_pos = matches[index + 1].start()
        else:
            end_pos = len(text)
        piece = text[start_pos:end_pos].strip()
        if not piece:
            continue
        # Allocate time proportionally across sub-segments.
        frac_start = index / total
        frac_end = (index + 1) / total
        sub_start = segment.start_seconds + duration * frac_start
        sub_end = segment.start_seconds + duration * frac_end
        words = None
        if segment.words:
            sub_words = [
                w
                for w in segment.words
                if w.start_seconds >= sub_start and w.end_seconds <= sub_end
            ]
            words = sub_words or None
        sub_segments.append(
            TranscriptionSegment(
                start_seconds=sub_start,
                end_seconds=sub_end,
                text=piece,
                speaker=match.group(1),
                words=words,
            )
        )
    return sub_segments or [segment]


def split_sentences_at_punctuation(
    segments: List[TranscriptionSegment],
) -> List[TranscriptionSegment]:
    """Merge whisper segments and cut sentences at sentence punctuation.

    Whisper's own segment breaks often fall mid-sentence, so text is accumulated
    across segments and a sentence is closed at periods/question marks. Commas
    are also eligible boundaries, but only when the text before the comma has
    enough words; otherwise the comma stays with the following text.

    Inline ``[S01]`` speaker tags are stripped and each speaker turn always
    starts a new sentence, even when no sentence-ending punctuation exists.
    """
    sentences: List[TranscriptionSegment] = []
    pieces: List[str] = []
    start_seconds: Optional[float] = None
    last_segment_end: Optional[float] = None
    current_speaker: Optional[str] = None

    def flush(end_seconds: float) -> None:
        nonlocal pieces, start_seconds, current_speaker
        text = " ".join(piece for piece in pieces if piece).strip()
        if text and start_seconds is not None and end_seconds > start_seconds:
            sentences.append(
                TranscriptionSegment(
                    start_seconds=start_seconds,
                    end_seconds=end_seconds,
                    text=text,
                    speaker=current_speaker,
                )
            )
        pieces = []
        start_seconds = None
        current_speaker = None

    expanded: List[TranscriptionSegment] = []
    for seg in segments:
        expanded.extend(_expand_speaker_segments(seg))

    for segment in sorted(
        expanded, key=lambda item: (item.start_seconds, item.end_seconds)
    ):
        text = segment.text.strip()
        if not text:
            continue
        # Each speaker-tagged sub-segment always starts a new sentence,
        # regardless of whether the speaker actually changed — MOSS emits
        # a tag per phrase/line even for the same speaker.
        if pieces and segment.speaker is not None:
            flush(segment.start_seconds)
        if not pieces:
            start_seconds = segment.start_seconds
        current_speaker = segment.speaker
        seg_pieces: List[Tuple[str, Optional[str]]] = []
        cursor = 0
        for match in _PUNCTUATION_BOUNDARY.finditer(text):
            boundary_type: Optional[str] = None
            if "," in match.group(0):
                boundary_type = "comma"
            elif _is_sentence_end_run(text, match.start(), match.end()):
                boundary_type = "sentence"
            if boundary_type is None:
                continue
            piece = text[cursor:match.end()].strip()
            cursor = match.end()
            if _has_word_token(piece):
                seg_pieces.append((piece, boundary_type))
        tail = text[cursor:].strip()
        if tail:
            seg_pieces.append((tail, False))
        if not seg_pieces:
            continue
        lengths = [len(piece) for piece, _ in seg_pieces]
        ends: List[float] = []
        if segment.words and len(segment.words) >= len(seg_pieces):
            ends = _word_boundary_ends(segment, lengths)
            if not _valid_boundaries(segment, ends):
                ends = []
        if not ends:
            ends = _proportional_ends(segment, lengths)
        last_segment_end = segment.end_seconds
        for index, (piece, boundary_type) in enumerate(seg_pieces):
            pieces.append(piece)
            if boundary_type == "sentence" or (
                boundary_type == "comma"
                and _word_count(" ".join(pieces)) >= MIN_WORDS_FOR_COMMA_SPLIT
            ):
                end_seconds = ends[index] if index < len(ends) else segment.end_seconds
                flush(end_seconds)
                start_seconds = end_seconds
                current_speaker = segment.speaker
    if pieces:
        flush(last_segment_end if last_segment_end is not None else start_seconds)
    return sentences


def _word_count(text: str) -> int:
    return len(text.split())


def _has_word_token(text: str) -> bool:
    return any(char.isalnum() for char in text)


def _is_sentence_end_run(text: str, run_start: int, run_end: int) -> bool:
    """Decide whether a "." / "?" run really ends a sentence."""
    punctuation = text[run_start:run_end]
    if "?" in punctuation:
        return True
    if (
        run_start > 0
        and run_end < len(text)
        and text[run_start - 1].isdigit()
        and text[run_end].isdigit()
    ):
        return False  # decimal point, e.g. "0.4"
    if len(punctuation) > 1:
        return True  # ellipsis or "!?" etc.
    if _is_abbreviation_run(text, run_start, run_end):
        return False  # e.g. "a.m.", "p.m.", "Mr."
    rest = text[run_end:].lstrip()
    if rest and rest[0].islower():
        return False  # period followed by a lowercase word is an abbreviation
    return True


def _is_abbreviation_run(text: str, run_start: int, run_end: int) -> bool:
    for match in _ABBREVIATION_RE.finditer(text):
        if not (match.start() < run_end <= match.end()):
            continue
        if match.group(0).lower() == "no.":
            rest = text[match.end():].lstrip()
            if not rest or not rest[0].isdigit():
                return False  # "No." as a reply ends the sentence
        return True
    return False


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