"""
Module for handling multilingual datasets for the CapibaraGPT model.

This module provides a class for creating and managing multilingual
datasets, including language detection and translation of
unsupported texts.

Classes:
    MultilingualDataset: Manages a multilingual dataset.

Dependencies:
    - transformers: For the AutoTokenizer.
    - langdetect: For language detection.
    - googletrans: For text translation.
    - dotenv: For loading environment variables.
"""

# Standard library imports
import logging
import os
import json
import pickle
from typing import List, Dict, Any, Optional, Tuple

# Third-party imports
from transformers import AutoTokenizer
from functools import lru_cache
from langdetect import detect
from googletrans import Translator
from dotenv import load_dotenv

logger = logging.getLogger(__name__)


class MultilingualDataset:
    """
    A dataset class for handling multilingual text data.

    This class processes and prepares multilingual text data for use in language models,
    including language detection, translation, and tokenization.

    Attributes:
        data (List[Dict[str, Any]]): List of data items, each containing at least a 'text' field.
        supported_languages (List[str]): List of supported language codes.
        tokenizer: Tokenizer for encoding the text.
        max_length (int): Maximum length for tokenization.
        translation_cache (Dict[int, tuple]): Cache for storing translated texts.
    """

    @classmethod
    def from_directory(cls, directory: str, supported_languages: List[str], **kwargs):
        """
        Create a MultilingualDataset from a directory of JSON files.

        Args:
            directory (str): Path to the directory containing JSON files.
            supported_languages (List[str]): List of supported language codes.
            **kwargs: Additional arguments to pass to the MultilingualDataset constructor.

        Returns:
            MultilingualDataset: A new instance of MultilingualDataset.
        """
        def load_file(filename: str) -> List[Dict[str, Any]]:
            with open(os.path.join(directory, filename), 'r', encoding='utf-8') as f:
                return json.load(f)

        data = []
        for filename in [f for f in os.listdir(directory) if f.endswith('.json')]:
            data.extend(load_file(filename))

        return cls(data, supported_languages, **kwargs)

    def __init__(self, data: List[Dict[str, Any]], supported_languages: List[str],
                 tokenizer: Optional[Any] = None, max_length: Optional[int] = None):
        """
        Initialize the MultilingualDataset.

        Args:
            data (List[Dict[str, Any]]): List of data items.
            supported_languages (List[str]): List of supported language codes.
            tokenizer (Optional[Any]): Pre-initialized tokenizer. If None, a default one will be used.
            max_length (Optional[int]): Maximum length for tokenization. If None, it will be set to 512.

        Raises:
            ValueError: If data is empty or not in the correct format.
            KeyError: If any data item is missing the 'text' field.
        """
        load_dotenv()
        self.max_length = max_length or int(
            os.getenv('CAPIBARA_MAX_LENGTH', 512))
        self._validate_input(data)
        self.data = data
        self.supported_languages = supported_languages
        self.tokenizer = tokenizer or AutoTokenizer.from_pretrained(
            "xlm-roberta-base")
        self.translation_cache = {}
        self.translator = Translator()

        logger.info(f"Created MultilingualDataset with {len(data)} items.")

    def _validate_input(self, data: List[Dict[str, Any]]) -> None:
        """Validate the input data."""
        if not isinstance(data, list) or len(data) == 0:
            raise ValueError("Data must be a non-empty list of dictionaries.")
        if not all(isinstance(item, dict) and 'text' in item for item in data):
            raise ValueError(
                "All items in data must be dictionaries with a 'text' key.")

    def __len__(self) -> int:
        """Return the number of items in the dataset."""
        return len(self.data)

    @lru_cache(maxsize=1000)
    def _detect_and_translate(self, idx: int) -> Tuple[str, str]:
        """Detect language and translate if necessary, with caching."""
        item = self.data[idx]
        text = item['text']
        lang = self._detect_language(text)

        if lang not in self.supported_languages:
            try:
                text = self._translate_text(text, self.supported_languages[0])
                lang = self.supported_languages[0]
                logger.info(f"Translated item {idx} to {
                            self.supported_languages[0]}")
            except Exception as e:
                logger.error(f"Error translating text at index {idx}: {e}")
                lang = self.supported_languages[0]

        return text, lang

    def __getitem__(self, idx: int) -> Dict[str, Any]:
        """
        Get a single item from the dataset.

        Args:
            idx (int): Index of the item to retrieve.

        Returns:
            Dict[str, Any]: A dictionary containing 'input_ids', 'attention_mask', and 'lang'.

        Raises:
            IndexError: If the index is out of bounds.
        """
        if not 0 <= idx < len(self.data):
            raise IndexError(f"Index {idx} is out of bounds for dataset of length {
                             len(self.data)}.")

        text, lang = self._detect_and_translate(idx)
        encoding = self.tokenizer(
            text,
            max_length=self.max_length,
            padding='max_length',
            truncation=True,
            return_tensors='jax'
        )

        return {
            'input_ids': encoding['input_ids'][0],
            'attention_mask': encoding['attention_mask'][0],
            'lang': self.supported_languages.index(lang)
        }

    @staticmethod
    def _detect_language(text: str) -> str:
        """Detect the language of the given text."""
        try:
            return detect(text)
        except Exception as e:
            logger.error(f"Language detection failed for text: '{
                         text}'. Error: {e}")
            return "en"  # default to English if detection fails

    def _translate_text(self, text: str, target_lang: str) -> str:
        """Translate the given text to the target language."""
        try:
            return self.translator.translate(text, dest=target_lang).text
        except Exception as e:
            logger.error(f"Translation failed for text: '{text}'. Error: {e}")
            return text  # return original text if translation fails

    def get_language_distribution(self) -> Dict[str, int]:
        """Get the distribution of languages in the dataset."""
        lang_dist = {}
        for idx in range(len(self)):
            _, lang = self._detect_and_translate(idx)
            lang_dist[lang] = lang_dist.get(lang, 0) + 1
        return lang_dist

    def save_translation_cache(self, filepath: str) -> None:
        """Save the translation cache to a file."""
        with open(filepath, 'wb') as f:
            pickle.dump(self.translation_cache, f)

    def load_translation_cache(self, filepath: str) -> None:
        """Load the translation cache from a file."""
        if os.path.exists(filepath):
            with open(filepath, 'rb') as f:
                self.translation_cache = pickle.load(f)

    def split_dataset(self, train_ratio: float = 0.8, val_ratio: float = 0.1) -> Tuple['MultilingualDataset', 'MultilingualDataset', 'MultilingualDataset']:
        """
        Split the dataset into train, validation, and test sets.

        Args:
            train_ratio (float): Ratio of data to use for training.
            val_ratio (float): Ratio of data to use for validation.

        Returns:
            Tuple[MultilingualDataset, MultilingualDataset, MultilingualDataset]: Train, validation, and test datasets.
        """
        total = len(self)
        train_size = int(total * train_ratio)
        val_size = int(total * val_ratio)
        test_size = total - train_size - val_size

        train_data = self.data[:train_size]
        val_data = self.data[train_size:train_size+val_size]
        test_data = self.data[train_size+val_size:]

        return (
            MultilingualDataset(
                train_data, self.supported_languages, self.tokenizer, self.max_length),
            MultilingualDataset(
                val_data, self.supported_languages, self.tokenizer, self.max_length),
            MultilingualDataset(
                test_data, self.supported_languages, self.tokenizer, self.max_length)
        )


# Example usage
if __name__ == "__main__":
    sample_data = [
        {"text": "Hello, world!"},
        {"text": "Bonjour le monde!"},
        {"text": "¡Hola mundo!"}
    ]
    supported_langs = ["en", "fr", "es"]

    dataset = MultilingualDataset(sample_data, supported_langs)
    print(f"Dataset size: {len(dataset)}")
    print(f"First item: {dataset[0]}")
    print(f"Language distribution: {dataset.get_language_distribution()}")
