import uuid
from pathlib import Path
from typing import Callable, Generator, List, Optional

import numpy as np
import soundfile as sf
from resource_segmentation import Resource, Segment, split
from scipy.signal import resample
from spacy.lang.xx import MultiLanguage
from spacy.language import Language
from spacy.tokens import Span

from .text_normalizer import TEXT_NORMALIZATION_LEVELS, normalize_text_for_tts
from .tts.protocol import TextToSpeechProtocol

SEGMENT_LEVEL = 1
SENTENCE_LEVEL = 2


class ChapterTTS:
    def __init__(
        self,
        tts_protocol: TextToSpeechProtocol,
        max_segment_length: int = 500,
        language_model: str | None = None,
        text_normalization_level: str = "basic",
    ):
        if max_segment_length < 1:
            raise ValueError("max_segment_length must be at least 1")
        if text_normalization_level not in TEXT_NORMALIZATION_LEVELS:
            raise ValueError(f"Unsupported text normalization level: {text_normalization_level}")
        self._tts_protocol = tts_protocol
        self._max_segment_length = max_segment_length
        self._text_normalization_level = text_normalization_level
        self._nlp = self._load_language_model(language_model)

    def _load_language_model(self, language_model: Optional[str]) -> Language:
        if language_model:
            try:
                import spacy

                return spacy.load(language_model)
            except OSError:
                pass

        nlp: Language = MultiLanguage()
        nlp.add_pipe("sentencizer")
        return nlp

    def process_chapter(
        self,
        text: str,
        output_path: Path,
        workspace_path: Path,
        voice: str,
        progress_callback: Optional[Callable[[int, int], None]] = None,
    ) -> None:
        # TODO: 重构 temp 逻辑，不需要 temp 这个概念了
        segments = list(self.split_text_into_segments(text))
        if not segments:
            return

        audio_segments: list[tuple[np.ndarray, int]] = []
        for i, segment in enumerate(segments):
            if progress_callback:
                progress_callback(i + 1, len(segments))

            session_id = str(uuid.uuid4())[:8]
            temp_audio_path = workspace_path / f"{session_id}_segment_{i:04d}.wav"

            self._tts_protocol.convert_text_to_audio(
                text=segment,
                output_path=temp_audio_path,
                voice=voice,
            )
            if not temp_audio_path.exists():
                continue

            audio_data: np.ndarray
            sample_rate: int
            audio_data, sample_rate = sf.read(temp_audio_path)
            audio_segments.append((audio_data, sample_rate))

        if audio_segments:
            _, first_sample_rate = audio_segments[0]
            resampled_segments = []
            for audio_data, sample_rate in audio_segments:
                if sample_rate != first_sample_rate:
                    resampled_length = int(len(audio_data) * first_sample_rate / sample_rate)
                    resampled_audio = resample(audio_data, resampled_length)
                    resampled_segments.append(resampled_audio)
                else:
                    resampled_segments.append(audio_data)

            final_audio = np.concatenate(resampled_segments)
            sf.write(output_path, final_audio, first_sample_rate)

    def split_text_into_segments(self, text: str) -> Generator[str, None, None]:
        text = normalize_text_for_tts(text, level=self._text_normalization_level)
        text = text.strip()
        if not text:
            return

        all_resources = []
        doc = self._nlp(text)
        next_start_incision = 2

        for sent in doc.sents:
            segment_text = sent.text.strip()
            if not segment_text:
                continue

            resources = list(self._build_segment_internal_structure(sent))
            if not resources:
                continue

            resources[0].start_incision = next_start_incision
            resources[-1].end_incision = 2
            next_start_incision = 2

            all_resources.extend(resources)

        if not all_resources:
            text_resource = Resource(
                count=len(text), start_incision=SENTENCE_LEVEL, end_incision=SENTENCE_LEVEL, payload=text
            )
            all_resources.append(text_resource)

        yield from self._split_by_resource_segmentation(all_resources)

    def _build_segment_internal_structure(self, sent: Span) -> Generator[Resource, None, None]:
        current_fragment: list[str] = []
        for token in sent:
            if token.is_punct:
                if current_fragment:
                    fragment_text = "".join(current_fragment)
                    fragment_resource = Resource(
                        count=len(fragment_text),
                        start_incision=SEGMENT_LEVEL,
                        end_incision=SEGMENT_LEVEL,
                        payload=fragment_text,
                    )
                    yield fragment_resource
                    current_fragment = []

                punct_resource = Resource(
                    count=len(token.text), start_incision=SEGMENT_LEVEL, end_incision=SEGMENT_LEVEL, payload=token.text
                )
                yield punct_resource
            else:
                current_fragment.append(token.text_with_ws)

        if current_fragment:
            fragment_text = "".join(current_fragment)
            fragment_resource = Resource(
                count=len(fragment_text),
                start_incision=SEGMENT_LEVEL,
                end_incision=SEGMENT_LEVEL,
                payload=fragment_text,
            )
            yield fragment_resource

    def _split_by_resource_segmentation(self, resources: List[Resource]) -> Generator[str, None, None]:
        groups = list(
            split(
                iter(resources),
                max_segment_count=self._max_segment_length,
                border_incision=1,
                gap_rate=0.0,
                tail_rate=0.0,
            )
        )
        for group in groups:
            segment_chars = []
            for item in group.body:
                if isinstance(item, Segment):
                    for resource in item.resources:
                        segment_chars.append(resource.payload)
                elif isinstance(item, Resource):
                    segment_chars.append(item.payload)

            combined_text = "".join(segment_chars).strip()
            if combined_text:
                yield from self._enforce_max_segment_length(combined_text)

    def _enforce_max_segment_length(self, text: str) -> Generator[str, None, None]:
        remaining = text.strip()
        while remaining:
            if len(remaining) <= self._max_segment_length:
                yield remaining
                return

            split_at = self._find_split_index(remaining)
            segment = remaining[:split_at].strip()
            if not segment:
                split_at = self._max_segment_length
                segment = remaining[:split_at]

            trailing = remaining[split_at:].lstrip()
            if (
                trailing
                and self._is_punctuation_only(trailing)
                and len(segment) + len(trailing) <= self._max_segment_length
            ):
                yield segment + trailing
                return

            yield segment
            remaining = remaining[split_at:].lstrip()

    def _find_split_index(self, text: str) -> int:
        max_index = min(len(text), self._max_segment_length)
        preferred_break_chars = "。！？!?；;：:\n，,、)]】》」』 "

        for index in range(max_index - 1, -1, -1):
            if text[index] in preferred_break_chars:
                return index + 1

        return max_index

    def _is_punctuation_only(self, text: str) -> bool:
        punctuation_chars = set("。！？!?；;：:\n，,、)]】》」』 ")
        return all(char in punctuation_chars for char in text)
