"""Tests for gherbal.model module — core Gherbal class."""

import os
import json
import tempfile
import pytest
import pandas as pd
from gherbal import Gherbal, GherbalConfig


class TestGherbalInit:
    """Tests for Gherbal initialization."""

    def test_init_with_empty_config(self):
        config = GherbalConfig()
        model = Gherbal(config)
        assert model.model is None
        assert model.labels is None

    def test_init_with_nonexistent_paths(self):
        config = GherbalConfig(
            fasttext_path="/nonexistent/model.ftz",
            labels_path="/nonexistent/labels.txt",
        )
        model = Gherbal(config)
        # Should not crash, just leave model/labels as None
        assert model.model is None
        assert model.labels is None


class TestGherbalPredict:
    """Tests for Gherbal.predict()."""

    def test_predict_single_string(self, trained_model):
        model, _ = trained_model
        result = model.predict("Hello how are you today")
        assert isinstance(result, list)
        assert len(result) > 0
        assert isinstance(result[0], tuple)
        assert isinstance(result[0][0], str)
        assert isinstance(result[0][1], float)

    def test_predict_list_of_strings(self, trained_model):
        model, _ = trained_model
        texts = ["Hello world", "Bonjour le monde", "مرحبا بالعالم"]
        result = model.predict(texts)
        assert isinstance(result, list)
        assert len(result) == 3
        for pred in result:
            assert isinstance(pred, list)
            assert len(pred) > 0

    def test_predict_returns_valid_labels(self, trained_model):
        model, _ = trained_model
        result = model.predict("Hello world")
        for label, prob in result:
            assert label in model.labels
            assert 0.0 <= prob <= 1.0

    def test_predict_k_parameter(self, trained_model):
        model, _ = trained_model
        result_all = model.predict("Hello world", k=-1)
        result_1 = model.predict("Hello world", k=1)
        assert len(result_1) == 1
        assert len(result_all) == len(model.labels)

    def test_predict_arabic(self, trained_model):
        model, _ = trained_model
        result = model.predict("مرحبا بالعالم كيف حالكم")
        top_label = result[0][0]
        assert top_label == "arb_Arab"

    def test_predict_english(self, trained_model):
        model, _ = trained_model
        result = model.predict("the cat sat on the mat")
        top_label = result[0][0]
        assert top_label == "eng_Latn"

    def test_predict_french(self, trained_model):
        model, _ = trained_model
        result = model.predict("bonjour comment allez vous")
        top_label = result[0][0]
        assert top_label == "fra_Latn"

    def test_predict_with_noisy_text(self, trained_model):
        model, _ = trained_model
        # Should not crash on noisy input
        result = model.predict("Hello @user https://test.com #tag 🎉")
        assert isinstance(result, list)

    def test_predict_empty_after_preprocessing(self, trained_model):
        model, _ = trained_model
        # Numbers-only or special chars might preprocess to empty string
        result = model.predict("🎉🎊🎈")
        assert isinstance(result, list)


class TestGherbalPredictDf:
    """Tests for Gherbal.predict_df()."""

    def test_predict_df_basic(self, trained_model):
        model, _ = trained_model
        df = pd.DataFrame({
            "text": ["Hello world", "Bonjour monde", "مرحبا بالعالم"],
        })
        result = model.predict_df(df)
        assert "lang_predicted" in result.columns
        assert len(result) == 3

    def test_predict_df_custom_columns(self, trained_model):
        model, _ = trained_model
        df = pd.DataFrame({"content": ["Hello world", "Bonjour"]})
        result = model.predict_df(df, text_column="content", prediction_column="lang")
        assert "lang" in result.columns

    def test_predict_df_all_labels(self, trained_model):
        model, _ = trained_model
        df = pd.DataFrame({"text": ["Hello world"]})
        result = model.predict_df(df, all_labels=True)
        # Should have columns for each label
        for label in model.labels:
            assert f"lang_predicted_{label}" in result.columns


class TestGherbalStaticMethods:
    """Tests for backwards-compatible static methods."""

    def test_preprocess_text_static(self):
        result = Gherbal.preprocess_text("Hello World @user")
        assert isinstance(result, str)

    def test_clean_text_static(self):
        result = Gherbal.clean_text("Hello World")
        assert isinstance(result, str)

    def test_remove_emoji_static(self):
        result = Gherbal.remove_emoji("Hello 🌍 World")
        assert isinstance(result, str)

    def test_create_script_cleaner(self):
        cleaner = Gherbal.create_script_cleaner()
        assert callable(cleaner)
        result = cleaner("Latn", "Hello World")
        assert isinstance(result, str)

    def test_create_script_cleaner_with_config(self):
        cleaner = Gherbal.create_script_cleaner({"numbers": True})
        assert callable(cleaner)


class TestGherbalSaveLoad:
    """Tests for save_pretrained / from_pretrained."""

    def test_save_pretrained_creates_files(self, trained_model, tmp_dir):
        model, _ = trained_model
        model.save_pretrained(tmp_dir)

        assert os.path.exists(os.path.join(tmp_dir, "config.json"))
        assert os.path.exists(os.path.join(tmp_dir, "labels.txt"))
        assert os.path.exists(os.path.join(tmp_dir, "gherbal_config.py"))
        assert os.path.exists(os.path.join(tmp_dir, "gherbal_model.py"))

        # Check model binary exists
        model_files = [f for f in os.listdir(tmp_dir) if f.endswith((".ftz", ".bin"))]
        assert len(model_files) > 0

    def test_save_pretrained_config_json(self, trained_model, tmp_dir):
        model, _ = trained_model
        model.save_pretrained(tmp_dir)

        with open(os.path.join(tmp_dir, "config.json")) as f:
            config_data = json.load(f)

        assert config_data["model_type"] == "gherbal"
        assert "auto_map" in config_data
        assert "labels_path" in config_data

    def test_save_pretrained_labels(self, trained_model, tmp_dir):
        model, _ = trained_model
        model.save_pretrained(tmp_dir)

        with open(os.path.join(tmp_dir, "labels.txt")) as f:
            saved_labels = [line.strip() for line in f if line.strip()]

        assert saved_labels == model.labels

    def test_from_pretrained_local(self, trained_model, tmp_dir):
        model, _ = trained_model
        model.save_pretrained(tmp_dir)

        loaded = Gherbal.from_pretrained(tmp_dir)
        assert loaded.model is not None
        assert loaded.labels is not None
        assert loaded.labels == model.labels

    def test_roundtrip_predictions_match(self, trained_model, tmp_dir):
        model, _ = trained_model
        model.save_pretrained(tmp_dir)

        loaded = Gherbal.from_pretrained(tmp_dir)
        texts = ["Hello world", "مرحبا بالعالم", "Bonjour le monde"]

        for text in texts:
            orig = model.predict(text)
            reloaded = loaded.predict(text)
            assert orig[0][0] == reloaded[0][0]  # same top prediction

    def test_from_pretrained_no_config_json(self, trained_model, tmp_dir):
        """Loading should work even without config.json (auto-detect files)."""
        model, _ = trained_model
        model.save_pretrained(tmp_dir)

        # Remove config.json to force auto-detection
        config_path = os.path.join(tmp_dir, "config.json")
        if os.path.exists(config_path):
            os.remove(config_path)

        loaded = Gherbal.from_pretrained(tmp_dir)
        assert loaded.model is not None
        assert loaded.labels is not None

    def test_save_pretrained_remote_code_files(self, trained_model, tmp_dir):
        """Verify the standalone inference code files are valid Python."""
        model, _ = trained_model
        model.save_pretrained(tmp_dir)

        for filename in ("gherbal_config.py", "gherbal_model.py"):
            path = os.path.join(tmp_dir, filename)
            assert os.path.exists(path)
            with open(path) as f:
                code = f.read()
            # Verify it compiles
            compile(code, filename, "exec")

    def test_from_pretrained_with_original_dir(self, trained_model):
        """Load directly from the training output directory."""
        model, model_dir = trained_model
        loaded = Gherbal.from_pretrained(model_dir)
        assert loaded.model is not None
        assert loaded.labels == model.labels


class TestGherbalTrain:
    """Tests for Gherbal.train()."""

    def test_train_basic(self, train_df, tmp_dir):
        model = Gherbal.train(train_df, tmp_dir, train_opts={"epoch": 2, "dim": 20})
        assert model.model is not None
        assert model.labels is not None
        assert len(model.labels) == 3

    def test_train_creates_files(self, train_df, tmp_dir):
        Gherbal.train(train_df, tmp_dir, train_opts={"epoch": 2, "dim": 20})
        assert os.path.exists(os.path.join(tmp_dir, "labels.txt"))
        model_files = [f for f in os.listdir(tmp_dir) if f.endswith((".bin", ".ftz"))]
        assert len(model_files) > 0

    def test_train_quantized(self, train_df, tmp_dir):
        model = Gherbal.train(
            train_df, tmp_dir,
            train_opts={"epoch": 2, "dim": 20},
            quantize=True,
        )
        assert model.model is not None
        assert os.path.exists(os.path.join(tmp_dir, "model.ftz"))

    def test_train_custom_hyperparams(self, train_df, tmp_dir):
        model = Gherbal.train(
            train_df,
            tmp_dir,
            dict_props={"wordNgrams": 2, "minn": 2, "maxn": 4},
            train_opts={"epoch": 2, "dim": 20, "lr": 0.5},
        )
        assert model.model is not None

    def test_train_cleans_temp_file(self, train_df, tmp_dir):
        Gherbal.train(train_df, tmp_dir, train_opts={"epoch": 2, "dim": 20})
        temp_files = [f for f in os.listdir(tmp_dir) if f.startswith("_train_")]
        assert len(temp_files) == 0

    def test_train_labels_sorted(self, train_df, tmp_dir):
        model = Gherbal.train(train_df, tmp_dir, train_opts={"epoch": 2, "dim": 20})
        assert model.labels == sorted(model.labels)


class TestGherbalQuantize:
    """Tests for Gherbal.quantize()."""

    def test_quantize(self, train_df, tmp_dir):
        """Train a fresh model and quantize it."""
        import os
        train_dir = os.path.join(tmp_dir, "train")
        quant_dir = os.path.join(tmp_dir, "quant")
        model = Gherbal.train(train_df, train_dir, train_opts={"epoch": 2, "dim": 20})
        model.quantize(quant_dir)
        assert os.path.exists(os.path.join(quant_dir, "model.ftz"))
        assert os.path.exists(os.path.join(quant_dir, "labels.txt"))

    def test_quantize_still_predicts(self, train_df, tmp_dir):
        """Quantized model should still be able to predict."""
        import os
        train_dir = os.path.join(tmp_dir, "train")
        quant_dir = os.path.join(tmp_dir, "quant")
        model = Gherbal.train(train_df, train_dir, train_opts={"epoch": 2, "dim": 20})
        model.quantize(quant_dir)
        result = model.predict("Hello world")
        assert isinstance(result, list)
        assert len(result) > 0
