from __future__ import annotations

from pathlib import Path
from typing import Any, Generator, Literal

from docarray import BaseDoc, DocList
from docarray.utils.filter import filter_docs
from pydantic import Field, model_validator
from typing_extensions import Self

from .audio import Audio, AudioChannel, AudioList


class AudioEvalExample(BaseDoc):
    """Audio example object that represents an example of input and target audio."""

    x: Audio = Field(..., description="Input audio.")
    y: Audio = Field(..., description="Target audio.")
    is_badcase: bool = Field(False, description="whether the example is a badcase")

    @model_validator(mode="after")
    def validate_audio_example(self):
        if len(self.x.channels) != len(self.y.channels):
            raise ValueError(
                "Input and target audio must have the same number of channels."
            )
        return self

    def char_error_rate(
        self,
        ignore_space: bool = True,
        ignore_punc: bool = True,
        normalize: bool = False,
    ) -> float:
        """Compute the character error rate (CER) between the input and target audio on all the channels.

        Args:
            ignore_space (bool, optional): Whether to ignore space. Defaults to True.
            ignore_punc (bool, optional): Whether to ignore punctuation. Defaults
            normalize (bool, optional): Whether to normalize the text. Defaults to False.
        Returns:
            float: Character error rate.

        """
        from fasr.utils.metric import compute_char_error_rate

        cer = 0
        for i in range(len(self.x.channels)):
            x = self.x.channels[i].text
            y = self.y.channels[i].text
            err = compute_char_error_rate(
                x=x,
                y=y,
                ignore_space=ignore_space,
                ignore_punc=ignore_punc,
                normalize=normalize,
            )
            if err > 0:
                self.is_badcase = True
            cer += err

        return round(cer, 4)

    def clear_waveform(self) -> Self:
        """Clear the waveform of the input and target audio."""
        self.x.clear_waveform()
        self.y.clear_waveform()
        return self

    def display(self) -> None:
        """Display character-level differences between x and y for each channel.

        Uses edit distance alignment to highlight:
        - Green: correct (matching characters)
        - Red: substitution (different characters)
        - Yellow: insertion (extra characters in x)
        - Cyan: deletion (missing characters from y)
        """
        from fasr.utils.metric import match_token

        RED = "\033[91m"
        GREEN = "\033[92m"
        YELLOW = "\033[93m"
        CYAN = "\033[96m"
        RESET = "\033[0m"
        BOLD = "\033[1m"

        for i in range(len(self.x.channels)):
            x_text = self.x.channels[i].text or ""
            y_text = self.y.channels[i].text or ""

            print(f"{BOLD}Channel {i}{RESET}")

            if x_text == y_text:
                print(f"  x: {GREEN}{x_text}{RESET}")
                print(f"  y: {GREEN}{y_text}{RESET}")
                print()
                continue

            x_chars = list(x_text)
            y_chars = list(y_text)
            matched = match_token(x_chars, y_chars)

            matched_y = {m[0] for m in matched}
            matched_x = {m[1] for m in matched}

            x_display = ""
            for idx, ch in enumerate(x_chars):
                if idx in matched_x:
                    x_display += f"{GREEN}{ch}{RESET}"
                else:
                    x_display += f"{YELLOW}{ch}{RESET}"

            y_display = ""
            for idx, ch in enumerate(y_chars):
                if idx in matched_y:
                    y_display += f"{GREEN}{ch}{RESET}"
                else:
                    y_display += f"{CYAN}{ch}{RESET}"

            subs = []
            for ref_idx, hyp_idx in matched:
                if y_chars[ref_idx] != x_chars[hyp_idx]:
                    subs.append(
                        f"{RED}{y_chars[ref_idx]}{RESET} → {RED}{x_chars[hyp_idx]}{RESET}"
                    )

            ins_chars = [
                x_chars[idx] for idx in range(len(x_chars)) if idx not in matched_x
            ]
            del_chars = [
                y_chars[idx] for idx in range(len(y_chars)) if idx not in matched_y
            ]

            print(f"  x: {x_display}")
            print(f"  y: {y_display}")
            if subs:
                print(f"  {RED}替换: {', '.join(subs)}{RESET}")
            if ins_chars:
                print(f"  {YELLOW}插入(x多余): {''.join(ins_chars)}{RESET}")
            if del_chars:
                print(f"  {CYAN}删除(y缺失): {''.join(del_chars)}{RESET}")
            print()


class AudioEvalResult(DocList):
    @property
    def x(self) -> AudioList:
        return AudioList(docs=[example.x for example in self])

    @property
    def y(self) -> AudioList:
        return AudioList(docs=[example.y for example in self])

    def add_example(self, x: Audio | AudioList, y: Audio | AudioList) -> Self:
        """Add an audio example to the evaluation."""
        if isinstance(x, Audio):
            x = AudioList([x])
        if isinstance(y, Audio):
            y = AudioList([y])
        if len(x) != len(y):
            raise ValueError(
                "Input and target audio must have the same number of channels."
            )
        for x_audio, y_audio in zip(x, y):
            self.append(AudioEvalExample(x=x_audio, y=y_audio))
        return self

    def char_error_rate(
        self,
        ignore_space: bool = True,
        ignore_punc: bool = True,
        normalize: bool = False,
    ) -> float:
        """Compute the character error rate (CER) between the input and target audio on all the channels.
        Args:
            ignore_space (bool, optional): Whether to ignore space. Defaults to True.
            ignore_punc (bool, optional): Whether to ignore punctuation. Defaults to True.
            normalize (bool, optional): Whether to normalize the text. Defaults to False.
        Returns:
            float: Character error rate.
        """

        cer = 0
        for audio_example in self:
            audio_example: AudioEvalExample
            cer += audio_example.char_error_rate(
                ignore_space=ignore_space,
                ignore_punc=ignore_punc,
                normalize=normalize,
            )

        return round(cer / len(self), 6)

    def save_binary(
        self,
        file: str | Path,
        protocol: Literal[
            "protobuf", "pickle", "json", "json-array", "protobuf-array", "pickle-array"
        ] = "protobuf",
        compress: Literal["lz4", "bz2", "lzma", "zlib", "gzip"] | None = None,
        show_progress: bool = False,
        clear_x_waveform: bool = True,
    ) -> None:
        """Save the audio list to a binary file.

        Args:
            file (str | Path): Path to the binary file.
            protocol (Literal[ &quot;protobuf&quot;, &quot;pickle&quot;, &quot;json&quot;, &quot;json, optional): the protocol to use. Defaults to &quot;protobuf&quot;.
            compress (Literal[&quot;lz4&quot;, &quot;bz2&quot;, &quot;lzma&quot;, &quot;zlib&quot;, &quot;gzip&quot;] | None, optional): the compression method to use. Defaults to None.
            show_progress (bool, optional): whether to show progress bar. Defaults to False.
            filter_bad (bool, optional): whether to filter bad audio files. Defaults to False.
            clear_x_waveform: whether clear the waveform of the input audio. Defaults to True.
        """
        for example in self:
            example: AudioEvalExample
            x_texts = []
            for channel in example.x.channels:
                channel: AudioChannel
                x_texts.append(channel.text)
            if clear_x_waveform:
                example.x.clear_waveform()
            example.x.clear_text()
            for i, channel in enumerate(example.x.channels):
                channel.raw_text = x_texts[i]
        return super().save_binary(file, protocol, compress, show_progress)

    @classmethod
    def load_binary(
        cls,
        file: str | bytes | Path,
        protocol: Literal[
            "protobuf", "pickle", "json", "json-array", "protobuf-array", "pickle-array"
        ] = "protobuf",
        compress: Literal["lz4", "bz2", "lzma", "zlib", "gzip"] | None = None,
        show_progress: bool = False,
        streaming: bool = False,
    ) -> AudioEvalResult | Generator[AudioEvalExample, None, None]:
        """Load the audio list from a binary file.

        Args:
            file (str | Path): Path to the binary file.
            show_progress (bool, optional): whether to show progress bar. Defaults to False.

        Returns:
            AudioEvalResult: Audio eval result.
        """
        docs = DocList[AudioEvalExample].load_binary(
            file=file,
            protocol=protocol,
            compress=compress,
            show_progress=show_progress,
            streaming=streaming,
        )
        if isinstance(docs, Generator):
            return docs
        return AudioEvalResult(docs=docs)

    def filter_by_cer(
        self,
        min_cer: float = 0.0,
        ignore_space: bool = True,
        ignore_punc: bool = True,
        normalize: bool = False,
    ) -> AudioEvalResult:
        """Filter examples whose character error rate is greater than min_cer.

        Args:
            min_cer: CER threshold, only examples with CER > min_cer are kept.
            ignore_space: Whether to ignore space when computing CER.
            ignore_punc: Whether to ignore punctuation when computing CER.
            normalize: Whether to normalize the text when computing CER.
        """
        filtered = []
        for example in self:
            example: AudioEvalExample
            cer = example.char_error_rate(
                ignore_space=ignore_space,
                ignore_punc=ignore_punc,
                normalize=normalize,
            )
            if cer > min_cer:
                filtered.append(example)
        return AudioEvalResult(docs=filtered)

    def get_badcase(self) -> AudioEvalResult:
        query = {"is_badcase": {"$eq": True}}
        return AudioEvalResult(docs=filter_docs(docs=self, query=query))
