"""
Gherbal: FastText-based language identification model with HuggingFace integration.
"""

import os
import shutil
import tempfile
from typing import List, Optional, Union

import fasttext
import pandas as pd
from huggingface_hub import HfApi, create_repo, hf_hub_download
from tqdm import tqdm
from transformers import PreTrainedModel

from .config import GherbalConfig
from .preprocessing import preprocess_text, create_clean_script_function
from .training import train_fasttext, quantize_model


class Gherbal(PreTrainedModel):
    """
    HuggingFace-compatible FastText language identification model.

    Supports loading from:
    - A local directory containing labels.txt and model.ftz/model.bin
    - A HuggingFace Hub repository with the same structure
    - AutoModel.from_pretrained with trust_remote_code=True
    """

    config_class = GherbalConfig
    base_model_prefix = "gherbal"
    supports_gradient_checkpointing = False

    def __init__(self, config: GherbalConfig):
        super().__init__(config)
        self.config = config
        self.model = None
        self.labels = None

        if config.fasttext_path and os.path.exists(config.fasttext_path):
            self._load_fasttext(config.fasttext_path)

        if config.labels_path and os.path.exists(config.labels_path):
            self._load_labels(config.labels_path)

    # ------------------------------------------------------------------
    # Internal loaders
    # ------------------------------------------------------------------
    def _load_fasttext(self, model_path):
        if not os.path.exists(model_path):
            raise FileNotFoundError(f"Model file not found: {model_path}")
        self.model = fasttext.load_model(model_path)

    def _load_labels(self, labels_path):
        with open(labels_path, "r", encoding="utf-8") as f:
            self.labels = [line.strip() for line in f if line.strip()]

    # ------------------------------------------------------------------
    # Prediction
    # ------------------------------------------------------------------
    def predict(self, texts, k=-1):
        """
        Predict language(s) for one or more texts.

        Args:
            texts: A single string or list of strings.
            k: Number of top predictions to return per text (-1 = all).

        Returns:
            For a single string: list of (label, probability) tuples.
            For a list: list of lists of (label, probability) tuples.
        """
        single = isinstance(texts, str)
        if single:
            texts = [texts]

        preprocessed = [preprocess_text(t) for t in texts]
        raw_labels, raw_probs = self.model.predict(preprocessed, k=k)

        results = []
        for labels_pred, probs in zip(raw_labels, raw_probs):
            result = [
                (self.labels[int(lbl.replace("__label__", ""))], float(prob))
                for lbl, prob in zip(labels_pred, probs)
            ]
            results.append(result)

        return results[0] if single else results

    def predict_df(
        self,
        df,
        text_column="text",
        prediction_column="lang_predicted",
        all_labels=False,
    ):
        """
        Add language predictions to a DataFrame.

        Args:
            df: Input DataFrame.
            text_column: Column containing text to classify.
            prediction_column: Column name for predictions.
            all_labels: If True, add probability columns for every label.

        Returns:
            DataFrame with prediction column(s) added.
        """
        preprocessed = df[text_column].apply(preprocess_text).tolist()

        if all_labels:
            k = len(self.labels)
            raw_labels, raw_probs = self.model.predict(preprocessed, k=k)
            for i, label in enumerate(self.labels):
                df[f"{prediction_column}_{label}"] = [p[i] for p in raw_probs]
        else:
            raw_labels, raw_probs = self.model.predict(preprocessed, k=1)
            df[prediction_column] = [
                self.labels[int(lbl[0].replace("__label__", ""))]
                for lbl in raw_labels
            ]

        return df

    # ------------------------------------------------------------------
    # Training
    # ------------------------------------------------------------------
    @classmethod
    def train(cls, train_df, save_path, dict_props=None, train_opts=None, quantize=False):
        """
        Train a new Gherbal model from a DataFrame.

        Args:
            train_df: DataFrame with 'lang' and 'text' columns.
            save_path: Directory to save model artifacts.
            dict_props: Override FastText dictionary properties.
            train_opts: Override FastText training options.
            quantize: If True, save quantized .ftz model.

        Returns:
            Trained Gherbal instance.
        """
        ft_model, labels, model_path, labels_path = train_fasttext(
            train_df, save_path, dict_props, train_opts, quantize
        )

        config = GherbalConfig(
            fasttext_path=model_path,
            labels_path=labels_path,
        )
        instance = cls(config)
        instance.model = ft_model
        instance.labels = labels
        return instance

    @classmethod
    def from_dataset(
        cls,
        dataset_name="sawalni-ai/gherbal-multi",
        subset=None,
        save_path="gherbal-multi",
        quantize=False,
    ):
        """
        Train from a HuggingFace dataset.

        Args:
            dataset_name: HuggingFace dataset identifier.
            subset: Dataset subset/config name.
            save_path: Where to save the trained model.
            quantize: Whether to quantize.

        Returns:
            Trained Gherbal instance.
        """
        from datasets import load_dataset

        ds = load_dataset(dataset_name, subset) if subset else load_dataset(dataset_name)
        train_df = pd.DataFrame(ds["train"])
        instance = cls.train(train_df, save_path, quantize=quantize)

        if "test" in ds:
            test_df = pd.DataFrame(ds["test"])
            instance.report(test_df)

        return instance

    def quantize(self, save_path, cutoff=100000):
        """
        Quantize the current model and save to disk.

        Args:
            save_path: Directory to save the quantized model.
            cutoff: Quantization cutoff parameter.

        Returns:
            self (with quantized model loaded).
        """
        self.model, model_path, labels_path = quantize_model(
            self.model, self.labels, save_path, cutoff=cutoff
        )
        self.config.fasttext_path = model_path
        self.config.labels_path = labels_path
        return self

    # ------------------------------------------------------------------
    # Evaluation (delegates to evaluation module)
    # ------------------------------------------------------------------
    def evaluate(self, eval_df):
        """Evaluate on a labeled DataFrame. Returns (accuracy, report, cm, labels, df)."""
        from .evaluation import evaluate
        return evaluate(self, eval_df)

    def report(self, eval_df, silent=False):
        """Evaluate and print a full report. Returns (accuracy, filtered_df)."""
        from .evaluation import report
        return report(self, eval_df, silent=silent)

    # ------------------------------------------------------------------
    # Preprocessing helpers (class methods for backwards compat)
    # ------------------------------------------------------------------
    @staticmethod
    def preprocess_text(text):
        """Preprocess text for language identification."""
        return preprocess_text(text)

    @staticmethod
    def clean_text(text):
        """Clean text (remove mentions, URLs, emojis, normalize)."""
        return preprocess_text(text)

    @staticmethod
    def remove_emoji(text):
        """Remove emojis from text (delegated to unscript.clean_text)."""
        return preprocess_text(text)

    @staticmethod
    def create_script_cleaner(config=None):
        """Create a script-aware text cleaner function."""
        return create_clean_script_function(config)

    # ------------------------------------------------------------------
    # HuggingFace Serialization
    # ------------------------------------------------------------------
    def save_pretrained(self, save_directory: str, **kwargs):
        """
        Save model, labels, config, and inference code to a directory.

        The saved directory can be loaded with:
        - Gherbal.from_pretrained(save_directory)
        - AutoModel.from_pretrained(save_directory, trust_remote_code=True)
        """
        os.makedirs(save_directory, exist_ok=True)

        # Save the inference code files for trust_remote_code support
        self._save_inference_code(save_directory)

        # Determine model filename
        ft_path = self.config.fasttext_path or ""
        model_name = os.path.basename(ft_path) if ft_path else "model.bin"
        if not model_name.endswith((".ftz", ".bin")):
            model_name = "model.bin"

        # Save FastText model binary
        if self.model is not None:
            src_path = self._resolve_model_path()
            dst_path = os.path.join(save_directory, model_name)
            if src_path and os.path.abspath(src_path) != os.path.abspath(dst_path):
                shutil.copy2(src_path, dst_path)
            elif not src_path:
                # If we can't find the source file, save from memory
                self.model.save_model(dst_path)

        # Save labels
        if self.labels is not None:
            with open(os.path.join(save_directory, "labels.txt"), "w", encoding="utf-8") as f:
                for label in self.labels:
                    f.write(f"{label}\n")

        # Save config with relative paths (create a copy to avoid mutating self)
        save_config = GherbalConfig(
            fasttext_path=model_name,
            labels_path="labels.txt",
        )
        save_config.auto_map = self.config.auto_map
        save_config.save_pretrained(save_directory)

    def _resolve_model_path(self):
        """Try to find the FastText model file on disk."""
        candidates = [
            self.config.fasttext_path,
            os.path.join(os.getcwd(), self.config.fasttext_path),
            os.path.abspath(self.config.fasttext_path),
        ]
        name_or_path = getattr(self.config, "_name_or_path", "")
        if name_or_path:
            candidates.append(os.path.join(name_or_path, self.config.fasttext_path))

        for path in candidates:
            if path and os.path.exists(path):
                return path
        return None

    def _save_inference_code(self, save_directory):
        """
        Write standalone inference code into the save directory so that
        AutoModel.from_pretrained(..., trust_remote_code=True) works.

        Two files are written:
        - gherbal_config.py  (referenced by auto_map AutoConfig)
        - gherbal_model.py   (referenced by auto_map AutoModel)
        """
        config_code = _REMOTE_CONFIG_CODE
        model_code = _REMOTE_MODEL_CODE

        with open(os.path.join(save_directory, "gherbal_config.py"), "w") as f:
            f.write(config_code)
        with open(os.path.join(save_directory, "gherbal_model.py"), "w") as f:
            f.write(model_code)

    @classmethod
    def from_pretrained(
        cls,
        pretrained_model_name_or_path: str,
        cache_dir=None,
        revision: str = None,
        *model_args,
        **kwargs,
    ) -> "Gherbal":
        """
        Load a Gherbal model from a local directory or HuggingFace Hub.

        Local directories must contain:
        - model.ftz or model.bin (FastText model)
        - labels.txt (one label per line)
        - Optionally config.json

        HuggingFace repos must have the same structure.

        Args:
            pretrained_model_name_or_path: Local path or HF repo id.
            cache_dir: Cache directory for HF downloads.
            revision: Branch, tag, or commit to load from HF.

        Returns:
            Loaded Gherbal instance.
        """
        config = kwargs.pop("config", None)
        token = kwargs.pop("token", kwargs.pop("use_auth_token", None))

        if config is None:
            config = cls._build_config(
                pretrained_model_name_or_path, cache_dir, revision, token, **kwargs
            )

        model = cls(config)
        cls._load_model_files(
            model, pretrained_model_name_or_path, cache_dir, revision, token
        )
        return model

    @classmethod
    def _build_config(cls, name_or_path, cache_dir, revision, token, **kwargs):
        """Build a GherbalConfig, auto-detecting model files for local paths."""
        if os.path.isdir(name_or_path):
            model_files = [
                f for f in os.listdir(name_or_path) if f.endswith((".bin", ".ftz"))
            ]
            if model_files:
                ftz = [f for f in model_files if f.endswith(".ftz")]
                kwargs.setdefault("fasttext_path", ftz[0] if ftz else model_files[0])

            if os.path.exists(os.path.join(name_or_path, "labels.txt")):
                kwargs.setdefault("labels_path", "labels.txt")

        try:
            config = GherbalConfig.from_pretrained(
                name_or_path,
                cache_dir=cache_dir,
                revision=revision,
                token=token,
                **kwargs,
            )
        except Exception:
            # No config.json on disk/hub — build from kwargs
            config = GherbalConfig(**kwargs)

        config._name_or_path = name_or_path
        return config

    @classmethod
    def _load_model_files(cls, model, name_or_path, cache_dir, revision, token):
        """Load FastText model and labels from local path or HF Hub."""
        # --- FastText model ---
        if model.model is None and model.config.fasttext_path:
            local = cls._try_local(name_or_path, model.config.fasttext_path)
            if local:
                model._load_fasttext(local)
            else:
                path = cls._download_from_hub(
                    name_or_path, model.config.fasttext_path, cache_dir, revision, token
                )
                model._load_fasttext(path)

        # --- Labels ---
        if model.labels is None and model.config.labels_path:
            local = cls._try_local(name_or_path, model.config.labels_path)
            if local:
                model._load_labels(local)
            else:
                path = cls._download_from_hub(
                    name_or_path, model.config.labels_path, cache_dir, revision, token
                )
                model._load_labels(path)

    @staticmethod
    def _try_local(base_path, filename):
        """Return the full path if filename exists relative to base_path."""
        if os.path.isdir(base_path):
            full = os.path.join(os.path.abspath(base_path), filename)
            if os.path.exists(full):
                return full
        # Also try filename as absolute
        if os.path.exists(filename):
            return filename
        return None

    @staticmethod
    def _download_from_hub(repo_id, filename, cache_dir, revision, token):
        """Download a file from HuggingFace Hub (cached)."""
        try:
            return hf_hub_download(
                repo_id=repo_id,
                filename=filename,
                revision=revision,
                token=token,
                cache_dir=cache_dir,
                local_files_only=True,
            )
        except Exception:
            return hf_hub_download(
                repo_id=repo_id,
                filename=filename,
                revision=revision,
                token=token,
                cache_dir=cache_dir,
            )

    def push_to_hub(self, repo_id, token=None, revision=None, private=True):
        """
        Push the model to HuggingFace Hub.

        Saves model, labels, config, and inference code so that
        AutoModel.from_pretrained works with trust_remote_code=True.

        Args:
            repo_id: HuggingFace repo id (e.g. 'username/model-name').
            token: HuggingFace API token (falls back to HF_TOKEN env var).
            revision: Optional tag to create after pushing.
            private: Whether to create a private repo.
        """
        from huggingface_hub import upload_file

        token = token or os.environ.get("HF_TOKEN")

        try:
            create_repo(repo_id, token=token, private=private)
        except Exception:
            pass  # repo may already exist

        with tempfile.TemporaryDirectory() as tmp_dir:
            self.save_pretrained(tmp_dir)

            for filename in os.listdir(tmp_dir):
                upload_file(
                    path_or_fileobj=os.path.join(tmp_dir, filename),
                    path_in_repo=filename,
                    repo_id=repo_id,
                    token=token,
                    revision="main",
                )

            if revision:
                api = HfApi()
                api.create_tag(
                    repo_id=repo_id, tag=revision, revision="main", token=token
                )


# ======================================================================
# Standalone inference code embedded into save_pretrained output.
# These strings are written as gherbal_config.py and gherbal_model.py
# so that `AutoModel.from_pretrained(..., trust_remote_code=True)` works
# without requiring the gherbal package to be installed.
# ======================================================================

_REMOTE_CONFIG_CODE = '''\
"""GherbalConfig — standalone config for trust_remote_code loading."""
from transformers import PretrainedConfig


class GherbalConfig(PretrainedConfig):
    model_type = "gherbal"

    def __init__(self, fasttext_path=None, labels_path=None, **kwargs):
        self.fasttext_path = fasttext_path
        self.labels_path = labels_path
        self.auto_map = {
            "AutoModel": "gherbal_model.Gherbal",
            "AutoConfig": "gherbal_config.GherbalConfig",
        }
        super().__init__(**kwargs)
'''

_REMOTE_MODEL_CODE = '''\
"""Gherbal — standalone model for trust_remote_code loading."""
import os
import re
import unicodedata

import fasttext
from huggingface_hub import hf_hub_download
from transformers import PreTrainedModel

from .gherbal_config import GherbalConfig


def _preprocess_text(text):
    """Lightweight preprocessing matching the full gherbal pipeline."""
    if not isinstance(text, str) or not text:
        return ""
    # Remove emojis
    text = re.sub(r"[\\U0001F000-\\U0001FFFF]", "", text)
    text = re.sub(r"[\\u200D\\u20E3\\uFE0F\\u2640-\\u2642\\u2600-\\u26FF\\u2700-\\u27BF]", "", text)
    # Remove mentions, hashtags, URLs
    text = re.sub(r"@{1,2}[a-zA-Z0-9_]+", "", text)
    text = re.sub(r"[+][a-zA-Z0-9_]+", "", text)
    text = re.sub(r"#[a-zA-Z0-9_]+", "", text)
    text = re.sub(r"https?://\\S+", "", text)
    # Remove punctuation and special chars
    text = re.sub(r"[^\\w\\s\\u0600-\\u06FF]", " ", text)
    # Normalize
    text = unicodedata.normalize("NFD", text).lower()
    text = "".join(c for c in text if c.isalnum() or c.isspace() or "\\u0600" <= c <= "\\u06ff")
    text = re.sub(r"(.)\\1{2,}", r"\\1", text)
    text = " ".join(text.split())
    if re.match(r"^\\d+$", text):
        return ""
    return text


try:
    from unscript import clean_text as _unscript_clean_text
    _preprocess = lambda t: _unscript_clean_text(t, lowercase=True) if isinstance(t, str) else ""
except ImportError:
    _preprocess = _preprocess_text


class Gherbal(PreTrainedModel):
    config_class = GherbalConfig
    base_model_prefix = "gherbal"
    supports_gradient_checkpointing = False

    def __init__(self, config):
        super().__init__(config)
        self.config = config
        self.model = None
        self.labels = None
        if config.fasttext_path and os.path.exists(config.fasttext_path):
            self.model = fasttext.load_model(config.fasttext_path)
        if config.labels_path and os.path.exists(config.labels_path):
            with open(config.labels_path, "r", encoding="utf-8") as f:
                self.labels = [l.strip() for l in f if l.strip()]

    def predict(self, texts, k=-1):
        single = isinstance(texts, str)
        if single:
            texts = [texts]
        preprocessed = [_preprocess(t) for t in texts]
        raw_labels, raw_probs = self.model.predict(preprocessed, k=k)
        results = []
        for lp, pp in zip(raw_labels, raw_probs):
            results.append([
                (self.labels[int(l.replace("__label__", ""))], float(p))
                for l, p in zip(lp, pp)
            ])
        return results[0] if single else results

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, cache_dir=None,
                        revision=None, *args, **kwargs):
        config = kwargs.pop("config", None)
        token = kwargs.pop("token", kwargs.pop("use_auth_token", None))

        if config is None:
            if os.path.isdir(pretrained_model_name_or_path):
                model_files = [f for f in os.listdir(pretrained_model_name_or_path)
                               if f.endswith((".bin", ".ftz"))]
                if model_files:
                    ftz = [f for f in model_files if f.endswith(".ftz")]
                    kwargs.setdefault("fasttext_path", ftz[0] if ftz else model_files[0])
                if os.path.exists(os.path.join(pretrained_model_name_or_path, "labels.txt")):
                    kwargs.setdefault("labels_path", "labels.txt")
            try:
                config = GherbalConfig.from_pretrained(
                    pretrained_model_name_or_path, cache_dir=cache_dir,
                    revision=revision, token=token, **kwargs)
            except Exception:
                config = GherbalConfig(**kwargs)
            config._name_or_path = pretrained_model_name_or_path

        model = cls(config)

        def _try_load(filename):
            if os.path.isdir(pretrained_model_name_or_path):
                p = os.path.join(os.path.abspath(pretrained_model_name_or_path), filename)
                if os.path.exists(p):
                    return p
            try:
                return hf_hub_download(repo_id=pretrained_model_name_or_path,
                                       filename=filename, revision=revision,
                                       token=token, cache_dir=cache_dir,
                                       local_files_only=True)
            except Exception:
                return hf_hub_download(repo_id=pretrained_model_name_or_path,
                                       filename=filename, revision=revision,
                                       token=token, cache_dir=cache_dir)

        if model.model is None and config.fasttext_path:
            model.model = fasttext.load_model(_try_load(config.fasttext_path))
        if model.labels is None and config.labels_path:
            with open(_try_load(config.labels_path), "r", encoding="utf-8") as f:
                model.labels = [l.strip() for l in f if l.strip()]
        return model
'''
