import ast
from functools import cached_property
from importlib import resources
from itertools import chain
from pathlib import Path
from typing import TYPE_CHECKING, Iterable, Optional, Union

import pandas as pd
from omegaconf import OmegaConf

from bewer.configs.resolve import resolve_pipelines
from bewer.core.example import Example
from bewer.core.keyword import Keyword, get_keyword_trie
from bewer.core.text import TokenList
from bewer.metrics.base import MetricCollection

__all__ = ["Dataset", "TextList", "TextTokenList"]

if TYPE_CHECKING:
    from bewer.core.keyword import KeywordTrie
    from bewer.core.text import Text


def _is_list_literal(s):
    try:
        return isinstance(ast.literal_eval(s), list)
    except (ValueError, SyntaxError):
        return False


class Dataset(object):
    """BeWER dataset.

    Attributes:
        config (OmegaConf): The resolved configuration object.
        pipelines: The resolved preprocessing pipelines.
        examples (list[Example]): A list of Example objects.
        metrics (MetricCollection): A metrics collection for the dataset.
        refs (TextList): The reference texts in the dataset.
        hyps (TextList): The hypothesis texts in the dataset.
    """

    def __init__(self, config: str | None = None):
        """Initialize the Dataset.

        The dataset must be populated using one of the load_* methods or manually using the add() method.

        Args:
            config (str | None): Path to the configuration file. If None, uses the default configuration.
        """
        self.config_path = self.get_config_path(config)
        self.config = OmegaConf.load(self.config_path)
        self._pipelines = resolve_pipelines(self.config)
        self.examples = []
        self._dynamic_keyword_vocabs = {}
        self._static_keyword_vocabs = {}
        self._cache_keyword_tries = {}
        self.metrics = MetricCollection(self)

    @property
    def pipelines(self):
        return self._pipelines

    @cached_property
    def refs(self) -> "TextList":
        """Get the reference texts as a TextList object.

        Returns:
            TextList: The reference texts.
        """
        return TextList([example.ref for example in self.examples])

    @cached_property
    def hyps(self) -> "TextList":
        """Get the hypothesis texts as a TextList object.

        Returns:
            TextList: The hypothesis texts.
        """
        return TextList([example.hyp for example in self.examples])

    def add(self, ref: str, hyp: str, keywords: dict[str, list[str]] | None = None) -> None:
        """Add an example to the dataset."""
        if keywords is not None:
            keywords = {name: set(kw_list) for name, kw_list in keywords.items()}
            for name, kw_set in keywords.items():
                self._update_static_keyword_vocab(name, kw_set)
        example = Example(ref, hyp, keywords=keywords, src=self, index=len(self))
        self.examples.append(example)

    def load_dataset(self, dataset, ref_col="ref", hyp_col="hyp", keyword_cols: list | None = None) -> None:
        """Load a Hugging Face dataset."""
        raise NotImplementedError("load_dataset() method not implemented.")

    def load_pandas(self, df: pd.DataFrame, ref_col="ref", hyp_col="hyp", keyword_cols: list | None = None) -> None:
        """Add a pandas DataFrame to the dataset."""
        if not isinstance(df, pd.DataFrame):
            raise TypeError("df must be a pandas DataFrame")
        if keyword_cols is None:
            keyword_cols = []

        for col in keyword_cols:
            df[col] = self._infer_keyword_column(df[col])

        # Add examples to the dataset
        for row in df.itertuples(index=False):
            hyp = getattr(row, hyp_col)
            ref = getattr(row, ref_col)
            if len(keyword_cols) > 0:
                keywords = {}
                for col in keyword_cols:
                    keywords[col] = getattr(row, col)
            else:
                keywords = None
            self.add(ref, hyp, keywords=keywords)

    def load_csv(self, csv_file: str, ref_col="ref", hyp_col="hyp", keyword_cols: list | None = None, **kwargs) -> None:
        """Add a CSV file to the dataset."""
        df = pd.read_csv(csv_file, **kwargs)
        self.load_pandas(df, ref_col, hyp_col, keyword_cols)

    def load_jsonl(
        self, jsonl_file: str, ref_col="ref", hyp_col="hyp", keyword_cols: list | None = None, **kwargs
    ) -> None:
        """Add a JSONL file to the dataset."""
        df = pd.read_json(jsonl_file, lines=True, **kwargs)
        self.load_pandas(df, ref_col, hyp_col, keyword_cols)

    def add_keyword_list(self, name: str, keywords: Iterable[str]) -> None:
        """Add a named keyword vocabulary to the dataset.

        Keywords are matched against the reference text of each example (including already added examples).

        Args:
            name (str): The name of the keyword vocabulary.
            keywords (Iterable[str]): The keywords to add.
        """
        if not isinstance(keywords, Iterable) or isinstance(keywords, str):
            raise TypeError("keywords must be an iterable of strings")

        keywords = set(keywords)

        for keyword in keywords:
            if not isinstance(keyword, str):
                raise TypeError(f"keywords must be an iterable of strings, but got element of type {type(keyword)}")

        self._update_dynamic_keyword_vocab(name, keywords)

    def add_keyword_file(self, name: str, keyword_file: str) -> None:
        """Add a named keyword vocabulary to the dataset from a file.

        The file should be plain text and contain one keyword per line.

        Keywords are matched against the reference text of each example (including already added examples).

        Args:
            name (str): The name of the keyword vocabulary.
            keyword_file (str): Path to the keyword file.
        """
        if not Path(keyword_file).is_file():
            raise FileNotFoundError(f"Keyword file {keyword_file} not found")

        with open(keyword_file, "r") as f:
            keywords = f.read().strip().splitlines()

        self.add_keyword_list(name, keywords)

    def _get_keyword_trie(
        self, vocab: str, normalized: bool = True, add_capitalized: bool = False
    ) -> Optional["KeywordTrie"]:
        """Get a trie for the specified keyword vocabulary."""
        return get_keyword_trie(
            self._dynamic_keyword_vocabs,
            self._cache_keyword_tries,
            vocab,
            normalized=normalized,
            add_capitalized=add_capitalized,
        )

    @staticmethod
    def get_config_path(config_path: str | None) -> str:
        """Get the configuration path."""
        if config_path is None or not Path(config_path).is_file():
            config_path = "base" if config_path is None else config_path
            return resources.files("bewer.configs").joinpath(f"{config_path}.yml")
        return Path(config_path).resolve()

    def _infer_keyword_column(self, series: pd.Series) -> pd.Series:
        """Infer the keyword terms from a pandas Series."""
        if series.map(_is_list_literal).all():
            series = series.apply(ast.literal_eval)
            return series
        elif series.map(lambda x: isinstance(x, str)).all():
            series = series.apply(lambda x: [x])
            return series
        elif series.map(lambda x: isinstance(x, list)).all():
            return series
        else:
            raise ValueError(f"Column {series.name} is not a list (or literal) or string")

    def _update_dynamic_keyword_vocab(self, name: str, keywords: set[str]) -> None:
        """Update the keyword vocabulary with new keywords."""
        keywords = set(Keyword(keyword, src=self) for keyword in keywords)
        if name in self._dynamic_keyword_vocabs:
            self._dynamic_keyword_vocabs[name].update(keywords)
        else:
            self._dynamic_keyword_vocabs[name] = set(keywords)

    def _update_static_keyword_vocab(self, name: str, keywords: set[str]) -> None:
        """Update the static keyword vocabulary with new keywords."""
        keywords = set(Keyword(keyword, src=self) for keyword in keywords)
        if name in self._static_keyword_vocabs:
            self._static_keyword_vocabs[name].update(keywords)
        else:
            self._static_keyword_vocabs[name] = set(keywords)

    def __len__(self) -> int:
        """Get the number of examples in the dataset."""
        return len(self.examples)

    def __getitem__(self, index: int) -> Example:
        """Get an example by index."""
        return self.examples[index]

    def __iter__(self):
        """Iterate over the examples in the dataset."""
        return iter(self.examples)

    def __repr__(self):
        """Get a string representation of the dataset."""
        return f"Dataset({len(self.examples)} examples)"


class TextList(tuple["Text", ...]):
    """An immutable sequence of Text objects."""

    def __new__(cls, iterable=()):
        return super().__new__(cls, iterable)

    @property
    def raw(self) -> list[str]:
        """Get the raw texts as a regular Python list.

        Returns:
            list[str]: The raw text.
        """
        return [text.raw for text in self]

    @property
    def standardized(self) -> list[str]:
        """Get the standardized texts as a regular Python list.

        Returns:
            list[str]: The standardized texts.
        """
        return [text.standardized for text in self]

    @property
    def tokens(self) -> "TextTokenList":
        """Get the tokens as a TextTokenList object.

        Returns:
            TextTokenList: The tokens.
        """
        return TextTokenList([text.tokens for text in self])

    def __getitem__(self, index: int | slice) -> Union["Text", "TextList"]:
        if isinstance(index, slice):
            return TextList(super().__getitem__(index))
        return super().__getitem__(index)

    def __add__(self, other: "TextList") -> "TextList":
        return TextList(super().__add__(other))

    def __repr__(self):
        texts = self[:60]
        texts_str = ",\n ".join([repr(text) for text in texts])
        if len(self) > 60:
            texts_str += ",\n ..."
        return f"TextList([\n {texts_str}]\n)"


class TextTokenList(tuple["TokenList", ...]):
    """An immutable sequence of TokenList objects."""

    def __new__(cls, iterable=()):
        return super().__new__(cls, iterable)

    @property
    def raw(self) -> list[list[str]]:
        """Get the raw tokens as a regular Python list.

        Returns:
            list[list[str]]: The raw tokens.
        """
        return [tokens.raw for tokens in self]

    @property
    def normalized(self) -> list[list[str]]:
        """Get the normalized tokens as a regular Python list.

        Returns:
            list[list[str]]: The normalized tokens.
        """
        return [tokens.normalized for tokens in self]

    @property
    def flat(self) -> "TokenList":
        """Flatten the TextTokenList into a TokenList.

        Returns:
            TokenList: The flattened TokenList.
        """
        return TokenList(chain(*self))

    def __getitem__(self, index: int | slice) -> Union["TokenList", "TextTokenList"]:
        if isinstance(index, slice):
            return TextTokenList(super().__getitem__(index))
        return super().__getitem__(index)

    def __add__(self, other: "TextTokenList") -> "TextTokenList":
        return TextTokenList(super().__add__(other))

    def __repr__(self):
        text_tokens = self[:60]
        text_tokens_str = ",\n ".join([tokens._sub_repr() for tokens in text_tokens])
        if len(self) > 60:
            text_tokens_str += ",\n ..."
        return f"TextTokenList([\n {text_tokens_str}]\n)"
