"""Tests for HuggingFace hub integration (save/load/push)."""

import os
import json
import tempfile
import pytest
from gherbal import Gherbal, GherbalConfig


class TestSavePretrained:
    """Tests for save_pretrained output structure."""

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

        required = ["config.json", "labels.txt", "gherbal_config.py", "gherbal_model.py"]
        for fname in required:
            assert os.path.exists(os.path.join(tmp_dir, fname)), f"Missing {fname}"

        binaries = [f for f in os.listdir(tmp_dir) if f.endswith((".ftz", ".bin"))]
        assert len(binaries) >= 1

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

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

        assert data["auto_map"]["AutoModel"] == "gherbal_model.Gherbal"
        assert data["auto_map"]["AutoConfig"] == "gherbal_config.GherbalConfig"

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

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

        # Paths should be just filenames, not absolute
        assert "/" not in data.get("labels_path", "")
        ft_path = data.get("fasttext_path", "")
        assert "/" not in ft_path or ft_path.startswith("model")

    def test_labels_match_model(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 = [line.strip() for line in f if line.strip()]
        assert saved == model.labels


class TestFromPretrainedLocal:
    """Tests for from_pretrained with local directories."""

    def test_roundtrip(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 == model.labels

    def test_predictions_stable_after_roundtrip(self, trained_model, tmp_dir):
        model, _ = trained_model
        model.save_pretrained(tmp_dir)
        loaded = Gherbal.from_pretrained(tmp_dir)

        for text in ["hello world", "مرحبا", "bonjour le monde"]:
            a = model.predict(text, k=1)
            b = loaded.predict(text, k=1)
            assert a[0][0] == b[0][0]

    def test_double_roundtrip(self, trained_model, tmp_dir):
        """Save -> Load -> Save -> Load should still work."""
        model, _ = trained_model

        dir1 = os.path.join(tmp_dir, "save1")
        dir2 = os.path.join(tmp_dir, "save2")

        model.save_pretrained(dir1)
        loaded1 = Gherbal.from_pretrained(dir1)
        loaded1.save_pretrained(dir2)
        loaded2 = Gherbal.from_pretrained(dir2)

        assert loaded2.labels == model.labels
        result = loaded2.predict("hello world", k=1)
        assert result[0][0] == model.predict("hello world", k=1)[0][0]

    def test_load_from_training_dir(self, trained_model):
        """Should load from the raw 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

    def test_load_without_config_json(self, trained_model, tmp_dir):
        """Auto-detect model files even without config.json."""
        model, _ = trained_model
        model.save_pretrained(tmp_dir)
        os.remove(os.path.join(tmp_dir, "config.json"))

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

    def test_prefers_ftz_over_bin(self, trained_model, tmp_dir):
        """When auto-detecting (no config.json), .ftz should be preferred over .bin."""
        model, _ = trained_model
        model.save_pretrained(tmp_dir)

        # Remove config.json to trigger auto-detection
        os.remove(os.path.join(tmp_dir, "config.json"))

        # Ensure both model.ftz and model.bin exist (copy the real binary to both names)
        import shutil
        model_files = [f for f in os.listdir(tmp_dir) if f.endswith((".ftz", ".bin"))]
        real_model = model_files[0]
        for name in ("model.ftz", "model.bin"):
            dst = os.path.join(tmp_dir, name)
            if not os.path.exists(dst):
                shutil.copy2(os.path.join(tmp_dir, real_model), dst)

        loaded = Gherbal.from_pretrained(tmp_dir)
        # Should load .ftz when both exist
        assert loaded.model is not None
        assert loaded.config.fasttext_path == "model.ftz"


class TestRemoteCodeFiles:
    """Tests that the standalone inference code files are valid."""

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

        with open(os.path.join(tmp_dir, "gherbal_config.py")) as f:
            code = f.read()
        compile(code, "gherbal_config.py", "exec")

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

        with open(os.path.join(tmp_dir, "gherbal_model.py")) as f:
            code = f.read()
        compile(code, "gherbal_model.py", "exec")

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

        with open(os.path.join(tmp_dir, "gherbal_config.py")) as f:
            code = f.read()
        assert "class GherbalConfig" in code
        assert "PretrainedConfig" in code

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

        with open(os.path.join(tmp_dir, "gherbal_model.py")) as f:
            code = f.read()
        assert "class Gherbal" in code
        assert "PreTrainedModel" in code
        assert "def predict" in code
        assert "def from_pretrained" in code

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

        with open(os.path.join(tmp_dir, "gherbal_model.py")) as f:
            code = f.read()
        # Should have fallback for when unscript is not installed
        assert "ImportError" in code
        assert "_preprocess_text" in code
