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])"
)


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.
    """
    sentences: List[TranscriptionSegment] = []
    pieces: List[str] = []
    start_seconds: Optional[float] = None
    last_segment_end: Optional[float] = None

    def flush(end_seconds: float) -> None:
        nonlocal pieces, start_seconds
        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,
                )
            )
        pieces = []
        start_seconds = None

    for segment in sorted(
        segments, key=lambda item: (item.start_seconds, item.end_seconds)
    ):
        text = segment.text.strip()
        if not text:
            continue
        if not pieces:
            start_seconds = segment.start_seconds
        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
    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