# -*- coding: utf-8 -*-
"""Module with a class to detect dominant languages in text data"""

import logging
import sys
from typing import List, AnyStr
from concurrent.futures import ThreadPoolExecutor

import pandas as pd
from langid.langid import LanguageIdentifier, model
from fastcore.utils import store_attr

# pycld3 is not available for Python >= 3.10, use pycld2 instead
if sys.version_info >= (3, 10):
    import pycld2
    USE_PYCLD2 = True
else:
    import cld3
    USE_PYCLD2 = False

from language_support import (
    SUPPORTED_LANGUAGES_CLD,
    SUPPORTED_LANGUAGES_CLD_NOT_LANGID,
    LANGUAGE_REMAPPING_CLD_LANGID,
)

from plugin_io_utils import generate_unique, truncate_text_list


class LanguageDetector:
    """Language detection wrapper class on top of CLD (pycld2/cld3) and `langid`

    Uses pycld2 for Python >= 3.10, cld3 for older versions.

    Additional features compared to using CLD or `langid` directly:
    - Use CLD for documents with more than 140 characters, else `langid`
        * This proved quite valuable in our benchmarks
        * CLD is very good for long documents but not for short ones
        * `langid` is more accurate for short documents
    - Harmonize small differences between CLD and langid language scopes
    - Add filter on language scope and minimum confidence score, else replace detection by fallback

    """

    LANGID_CLD3_NUM_CHAR_THRESHOLD = 140
    NUM_THREADS = 4
    COLUMN_DESCRIPTIONS = {
        "language_code": "Language code in ISO 639-1 format",
        "language_name": "Language name in ISO 639-1 format",
        "language_score": "Probability score of the predicted language from 0 to 1",
    }

    def __init__(
        self,
        language_scope: List = SUPPORTED_LANGUAGES_CLD.keys(),
        minimum_score: float = 0.0,
        fallback_language: AnyStr = "",
    ):
        store_attr()
        self.column_descriptions = self.COLUMN_DESCRIPTIONS.copy()  # may be changed by detect_languages_df
        self._langid_identifier = LanguageIdentifier.from_modelstring(model, norm_probs=True)
        self._langid_identifier.set_languages(
            [l for l in self.language_scope if l not in SUPPORTED_LANGUAGES_CLD_NOT_LANGID]
        )

    def _langid_detection(self, doc: AnyStr) -> (AnyStr, float):
        """Detect the language of a string using the `langid` library"""
        language_detection_object = self._langid_identifier.classify(doc)
        lang_id = language_detection_object[0][:2]
        lang_probability = float(language_detection_object[1])
        return (lang_id, lang_probability)

    def _cld_detection(self, doc: AnyStr) -> (AnyStr, float):
        """Detect the language using pycld2 (Python >= 3.10) or cld3 (Python < 3.10)"""
        if USE_PYCLD2:
            lang_id, lang_probability = self._pycld2_detection(doc)
        else:
            lang_id, lang_probability = self._cld3_detection(doc)
        # Remap language codes to match langid conventions
        for original_code, new_code in LANGUAGE_REMAPPING_CLD_LANGID.items():
            lang_id = lang_id.replace(original_code, new_code)
        return (lang_id, lang_probability)

    def _cld3_detection(self, doc: AnyStr) -> (AnyStr, float):
        """Detect the language of a string using the `cld3` library (Python < 3.10)"""
        language_detection_object = cld3.get_language(doc)
        lang_id = language_detection_object.language[:2]
        lang_probability = float(language_detection_object.probability)
        return (lang_id, lang_probability)

    def _pycld2_detection(self, doc: AnyStr) -> (AnyStr, float):
        """Detect the language of a string using the `pycld2` library (Python >= 3.10)"""
        try:
            _, _, details = pycld2.detect(doc)
            if details[0][1] == "un":  # unknown
                return ("un", 0.0)
            lang_id = details[0][1]
            lang_probability = details[0][2] / 100.0  # Convert percentage (0-100) to probability (0-1)
            return (lang_id, lang_probability)
        except pycld2.error:
            return ("un", 0.0)

    def _detection_filter(self, doc: AnyStr, lang_id: AnyStr, lang_probability: float) -> (AnyStr, float):
        """Filter the detected language of a string using the `language_scope` and `minimum_score` attributes

        If the detected language is not in `self.language_scope` or if its probability is below `self.minimum_score`,
        then the function returns `self.fallback_language` and replaces the detected probability by `None`

        """
        if lang_probability < self.minimum_score or lang_id not in self.language_scope:
            warning_msg = f"Problem encountered for document: '{truncate_text_list([doc])[0]}'\n"
            if lang_id not in self.language_scope:
                warning_msg += f"Detected language: '{lang_id}' not within language scope: {self.language_scope}\n"
            else:
                warning_msg += f"Confidence score: {lang_probability:.2f} below minimum: {self.minimum_score:.2f}\n"
            warning_msg += f"Replacing detected language: '{lang_id}' by fallback: '{self.fallback_language}'"
            logging.warning(warning_msg)
            lang_id, lang_probability = self.fallback_language, None
        return (lang_id, lang_probability)

    def detect_language_doc(self, doc: AnyStr) -> (AnyStr, AnyStr, float):
        """Detect the language of a string using CLD (pycld2 or cld3) or `langid` libraries

        Use CLD if the string has more than `self.LANGID_CLD3_NUM_CHAR_THRESHOLD` characters, else `langid`
        Apply the filtering method `_detection_filter` and round language probability to 3 decimals

        """
        # Route to langid or CLD depending on number of characters
        if not doc:
            return ("", "", None)
        if len(doc) <= self.LANGID_CLD3_NUM_CHAR_THRESHOLD:
            lang_id, lang_probability = self._langid_detection(doc)
        else:
            lang_id, lang_probability = self._cld_detection(doc)
        # Filters for language scope and minimum scores
        lang_id, lang_probability = self._detection_filter(doc, lang_id, lang_probability)
        # Enrich with language human name
        lang_name = SUPPORTED_LANGUAGES_CLD.get(lang_id, "")
        # Round probability to 3 decimals
        lang_probability = round(lang_probability, 3) if lang_probability else None
        return (lang_id, lang_name, lang_probability)

    def detect_languages_df(self, df: pd.DataFrame, text_column: AnyStr) -> pd.DataFrame:
        """Apply the `detect_language_doc` method to a pandas DataFrame with a text column, with multithreading"""
        self.column_descriptions = {}
        for k, v in self.COLUMN_DESCRIPTIONS.items():
            self.column_descriptions[generate_unique(k, df.keys(), text_column)] = v
        doc_iterator = (doc for _, doc in df[text_column].astype(str).items())
        output_df = df.copy()
        with ThreadPoolExecutor(max_workers=self.NUM_THREADS) as executor:
            lang_output_tuple_list = list(executor.map(self.detect_language_doc, doc_iterator))
        for i, col in enumerate(self.column_descriptions):
            output_df[col] = [t[i] for t in lang_output_tuple_list]
        return output_df
