"""
Training logic for Gherbal FastText-based language identification models.
"""

import os
from tqdm import tqdm
from .preprocessing import preprocess_text


# Default FastText dictionary properties
DEFAULT_DICT_PROPS = {
    "minCount": 1,
    "minCountLabel": 0,
    "wordNgrams": 3,
    "bucket": 2000000,
    "minn": 3,
    "maxn": 6,
    "t": 0.0001,
    "label": "__label__",
}

# Default FastText training options
DEFAULT_TRAIN_OPTS = {
    "lr": 1.0,
    "lrUpdateRate": 100,
    "dim": 600,
    "ws": 5,
    "epoch": 20,
    "neg": 5,
    "loss": "softmax",
    "thread": 12,
}


def train_fasttext(
    train_df,
    save_path,
    dict_props=None,
    train_opts=None,
    quantize=False,
):
    """
    Train a FastText supervised model for language identification.

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

    Returns:
        Tuple of (fasttext_model, labels_list, model_path, labels_path).
    """
    import fasttext

    merged_dict = DEFAULT_DICT_PROPS.copy()
    if dict_props:
        merged_dict.update(dict_props)

    merged_opts = DEFAULT_TRAIN_OPTS.copy()
    if train_opts:
        merged_opts.update(train_opts)

    labels = sorted(train_df["lang"].unique())
    label_to_index = {label: idx for idx, label in enumerate(labels)}

    label_prefix = merged_dict["label"]
    train_df = train_df.copy()
    train_df["lang_fasttext"] = train_df["lang"].apply(
        lambda x: f"{label_prefix}{label_to_index[x]}"
    )
    train_df["text"] = (
        train_df["text"]
        .apply(preprocess_text)
        .str.replace("\n", " ", regex=False)
    )

    os.makedirs(save_path, exist_ok=True)
    train_file = os.path.join(save_path, "_train_fasttext.txt")

    # Write training file in chunks for memory efficiency
    with open(train_file, "w", encoding="utf-8") as f:
        chunk_size = 100_000
        for start in range(0, len(train_df), chunk_size):
            chunk = train_df.iloc[start : start + chunk_size]
            lines = chunk["lang_fasttext"].values + " " + chunk["text"].values
            f.write("\n".join(lines) + "\n")

    ft_dict_args = {k: v for k, v in merged_dict.items() if k != "label"}
    model = fasttext.train_supervised(
        input=train_file,
        **ft_dict_args,
        **merged_opts,
    )

    # Save labels
    labels_path = os.path.join(save_path, "labels.txt")
    with open(labels_path, "w", encoding="utf-8") as f:
        for label in labels:
            f.write(f"{label}\n")

    # Save model (quantized or regular)
    if quantize:
        model.quantize(qnorm=True, retrain=False, cutoff=100000)
        model_path = os.path.join(save_path, "model.ftz")
    else:
        model_path = os.path.join(save_path, "model.bin")
    model.save_model(model_path)

    # Clean up temp file
    if os.path.exists(train_file):
        os.unlink(train_file)

    return model, labels, model_path, labels_path


def quantize_model(model, labels, save_path, cutoff=100000):
    """
    Quantize an existing FastText model and save to disk.

    Args:
        model: Loaded fasttext model object.
        labels: List of label strings.
        save_path: Directory to save the quantized model and labels.
        cutoff: Quantization cutoff parameter.

    Returns:
        Tuple of (quantized_model, model_path, labels_path).
    """
    os.makedirs(save_path, exist_ok=True)

    model.quantize(qnorm=True, retrain=False, cutoff=cutoff)
    model_path = os.path.join(save_path, "model.ftz")
    model.save_model(model_path)

    labels_path = os.path.join(save_path, "labels.txt")
    with open(labels_path, "w", encoding="utf-8") as f:
        for label in labels:
            f.write(f"{label}\n")

    return model, model_path, labels_path
