import json
import os
import struct
from argparse import Namespace
from pathlib import Path
from typing import Any, Optional, TypeVar
from unittest.mock import Mock

import pytest
from pytest_mock import MockerFixture

from huggingface_hub import HfApi, ModelCard, constants, hf_hub_download
from huggingface_hub.errors import RemoteEntryNotFoundError
from huggingface_hub.hub_mixin import ModelHubMixin, PyTorchModelHubMixin
from huggingface_hub.serialization._torch import storage_ptr
from huggingface_hub.utils import SoftTemporaryDirectory, is_torch_available

from .testing_constants import TOKEN, USER
from .testing_utils import repo_name


DUMMY_OBJECT = object()

DUMMY_MODEL_CARD_TEMPLATE = """
---
{{ card_data }}
---

This is a dummy model card.
Arxiv ID: 1234.56789
"""

DUMMY_MODEL_CARD_TEMPLATE_WITH_CUSTOM_KWARGS = """
---
{{ card_data }}
---

This is a dummy model card with kwargs.

{{ custom_data }}

- Code: {{ repo_url }}
- Paper: {{ paper_url }}
- Docs: {{ docs_url }}
"""

if is_torch_available():
    import torch
    import torch.nn as nn

    CONFIG = {"num": 10, "act": "gelu_fast"}

    class DummyModel(nn.Module, PyTorchModelHubMixin):
        def __init__(self, **kwargs):
            super().__init__()
            self.config = kwargs.pop("config", None)
            self.l1 = nn.Linear(2, 2)

        def forward(self, x):
            return self.l1(x)

    class DummyModelWithModelCard(
        nn.Module,
        PyTorchModelHubMixin,
        model_card_template=DUMMY_MODEL_CARD_TEMPLATE,
        language=["en", "zh"],
        library_name="my-dummy-lib",
        license="apache-2.0",
        tags=["tag1", "tag2"],
        pipeline_tag="text-classification",
    ):
        def __init__(self, linear_layer: int = 4):
            super().__init__()
            self.l1 = nn.Linear(linear_layer, linear_layer)

        def forward(self, x):
            return self.l1(x)

    class DummyModelNoConfig(nn.Module, PyTorchModelHubMixin):
        def __init__(
            self,
            num_classes: int = 42,
            state: str = "layernorm",
            not_jsonable: Any = DUMMY_OBJECT,
        ):
            super().__init__()
            self.num_classes = num_classes
            self.state = state
            self.not_jsonable = not_jsonable

    class DummyModelWithConfigAndKwargs(nn.Module, PyTorchModelHubMixin):
        def __init__(self, num_classes: int = 42, state: str = "layernorm", config: Optional[dict] = None, **kwargs):
            super().__init__()

    class DummyModelWithModelCardAndCustomKwargs(
        nn.Module,
        PyTorchModelHubMixin,
        model_card_template=DUMMY_MODEL_CARD_TEMPLATE_WITH_CUSTOM_KWARGS,
        docs_url="https://hf.co/docs/my-repo",
        paper_url="https://arxiv.org/abs/2304.12244",
        repo_url="https://hf.co/my-repo",
    ):
        def __init__(self, linear_layer: int = 4):
            super().__init__()

    class DummyModelWithEncodedConfig(
        nn.Module,
        PyTorchModelHubMixin,
        coders={
            Namespace: (
                lambda x: vars(x),
                lambda data: Namespace(**data),
            )
        },
    ):
        # Regression test for https://github.com/huggingface/huggingface_hub/issues/2334
        def __init__(self, config: Namespace):
            super().__init__()
            self.config = config

    class DummyModelWithTag1(nn.Module, PyTorchModelHubMixin, tags=["tag1"]):
        """Used to test tags not shared between sibling classes (only inheritance)."""

    class DummyModelWithTag2(nn.Module, PyTorchModelHubMixin, tags=["tag2"]):
        """Used to test tags not shared between sibling classes (only inheritance)."""

else:
    DummyModel = None
    DummyModelWithModelCard = None
    DummyModelNoConfig = None
    DummyModelWithConfigAndKwargs = None
    DummyModelWithModelCardAndCustomKwargs = None
    DummyModelWithTag1 = None
    DummyModelWithTag2 = None


@pytest.mark.skipif(not is_torch_available(), reason="Test requires torch")
class TestPytorchHubMixin:
    def test_save_pretrained_basic(self, tmp_path):
        DummyModel().save_pretrained(tmp_path)
        files = os.listdir(tmp_path)
        assert set(files) == {"README.md", "model.safetensors"}

    def test_save_pretrained_with_config(self, tmp_path):
        DummyModel().save_pretrained(tmp_path, config=CONFIG)
        files = os.listdir(tmp_path)
        assert set(files) == {"README.md", "config.json", "model.safetensors"}

    def test_save_as_safetensors(self, tmp_path):
        DummyModel().save_pretrained(tmp_path, config=TOKEN)
        modelFile = tmp_path / "model.safetensors"
        # check for safetensors header to ensure we are saving the model in safetensors format
        # while an implementation detail, assert as this has safety implications
        # https://github.com/huggingface/safetensors?tab=readme-ov-file#format
        with open(modelFile, "rb") as f:
            header_size = struct.unpack("<Q", f.read(8))[0]
            assert header_size == 128

    def test_save_pretrained_with_push_to_hub(self, tmp_path):
        repo_id = repo_name("save")
        save_directory = tmp_path / repo_id

        config = {"hello": "world"}
        mocked_model = DummyModel()
        mocked_model.push_to_hub = Mock()
        mocked_model._save_pretrained = Mock()  # disable _save_pretrained to speed-up

        # Not pushed to hub
        mocked_model.save_pretrained(save_directory)
        mocked_model.push_to_hub.assert_not_called()

        # Push to hub with repo_id
        mocked_model.save_pretrained(save_directory, push_to_hub=True, repo_id="CustomID", config=config)
        mocked_model.push_to_hub.assert_called_with(repo_id="CustomID", config=config, model_card_kwargs={})

        # Push to hub with default repo_id (based on dir name)
        mocked_model.save_pretrained(save_directory, push_to_hub=True, config=config)
        mocked_model.push_to_hub.assert_called_with(repo_id=repo_id, config=config, model_card_kwargs={})

    def test_from_pretrained_model_id_only(self, mocker: MockerFixture) -> None:
        from_pretrained_mock = mocker.patch.object(DummyModel, "_from_pretrained")
        model = DummyModel.from_pretrained("namespace/repo_name")
        from_pretrained_mock.assert_called_once()
        assert model is from_pretrained_mock.return_value

    def pretend_file_download(self, tmp_dir: Path, **kwargs):
        if kwargs.get("filename") == "config.json":
            raise RemoteEntryNotFoundError("no config", response=Mock())
        DummyModel().save_pretrained(tmp_dir)
        return tmp_dir / "model.safetensors"

    def test_from_pretrained_model_from_hub_prefer_safetensor(self, mocker: MockerFixture, tmp_path) -> None:
        hf_hub_download_mock = mocker.patch("huggingface_hub.hub_mixin.hf_hub_download")
        hf_hub_download_mock.side_effect = lambda **kwargs: self.pretend_file_download(tmp_path, **kwargs)
        model = DummyModel.from_pretrained("namespace/repo_name")
        hf_hub_download_mock.assert_any_call(
            repo_id="namespace/repo_name",
            filename="model.safetensors",
            revision=None,
            cache_dir=None,
            force_download=False,
            token=None,
            local_files_only=False,
        )
        assert model is not None

    def pretend_file_download_fallback(self, tmp_dir: Path, **kwargs):
        filename = kwargs.get("filename")
        if filename == "model.safetensors" or filename == "config.json":
            raise RemoteEntryNotFoundError("not found", response=Mock())

        class TestMixin(ModelHubMixin):
            def _save_pretrained(self, save_directory: Path) -> None:
                torch.save(DummyModel().state_dict(), save_directory / constants.PYTORCH_WEIGHTS_NAME)

        TestMixin().save_pretrained(tmp_dir)
        return tmp_dir / constants.PYTORCH_WEIGHTS_NAME

    def test_from_pretrained_model_from_hub_fallback_pickle(self, mocker: MockerFixture, tmp_path) -> None:
        hf_hub_download_mock = mocker.patch("huggingface_hub.hub_mixin.hf_hub_download")
        hf_hub_download_mock.side_effect = lambda **kwargs: self.pretend_file_download_fallback(tmp_path, **kwargs)
        model = DummyModel.from_pretrained("namespace/repo_name")
        hf_hub_download_mock.assert_any_call(
            repo_id="namespace/repo_name",
            filename="model.safetensors",
            revision=None,
            cache_dir=None,
            force_download=False,
            token=None,
            local_files_only=False,
        )
        hf_hub_download_mock.assert_any_call(
            repo_id="namespace/repo_name",
            filename="pytorch_model.bin",
            revision=None,
            cache_dir=None,
            force_download=False,
            token=None,
            local_files_only=False,
        )
        assert model is not None

    def test_from_pretrained_model_id_and_revision(self, mocker: MockerFixture) -> None:
        """Regression test for #1313.
        See https://github.com/huggingface/huggingface_hub/issues/1313."""
        from_pretrained_mock = mocker.patch.object(DummyModel, "_from_pretrained")
        model = DummyModel.from_pretrained("namespace/repo_name", revision="123456789")
        from_pretrained_mock.assert_called_once_with(
            model_id="namespace/repo_name",
            revision="123456789",  # Revision is passed correctly!
            cache_dir=None,
            force_download=False,
            local_files_only=False,
            token=None,
        )
        assert model is from_pretrained_mock.return_value

    def test_from_pretrained_to_relative_path(self):
        with SoftTemporaryDirectory(dir=Path(".")) as tmp_relative_dir:
            relative_save_directory = Path(tmp_relative_dir) / "model"
            DummyModel().save_pretrained(relative_save_directory, config=CONFIG)
            model = DummyModel.from_pretrained(relative_save_directory)
            assert model._hub_mixin_config == CONFIG

    def test_from_pretrained_to_absolute_path(self, tmp_path):
        save_directory = tmp_path / "subfolder"
        DummyModel().save_pretrained(save_directory, config=CONFIG)
        model = DummyModel.from_pretrained(save_directory)
        assert model._hub_mixin_config == CONFIG

    def test_from_pretrained_to_absolute_string_path(self, tmp_path):
        save_directory = str(tmp_path / "subfolder")
        DummyModel().save_pretrained(save_directory, config=CONFIG)
        model = DummyModel.from_pretrained(save_directory)
        assert model._hub_mixin_config == CONFIG

    def test_return_type_hint_from_pretrained(self):
        assert "return" in DummyModel.from_pretrained.__annotations__, (
            "`PyTorchModelHubMixin.from_pretrained` does not set a return type annotation."
        )
        assert isinstance(
            DummyModel.from_pretrained.__annotations__["return"],
            TypeVar,
        ), "`PyTorchModelHubMixin.from_pretrained` return type annotation is not a TypeVar."
        assert DummyModel.from_pretrained.__annotations__["return"].__bound__.__forward_arg__ == "ModelHubMixin", (
            "`PyTorchModelHubMixin.from_pretrained` return type annotation is not a TypeVar bound by `ModelHubMixin`."
        )

    def test_push_to_hub(self, api: HfApi, tmp_path):
        repo_id = f"{USER}/{repo_name('push_to_hub')}"
        DummyModel().push_to_hub(repo_id=repo_id, token=TOKEN, config=CONFIG)

        # Test model id exists
        assert api.model_info(repo_id).id == repo_id

        # Test config has been pushed to hub
        tmp_config_path = hf_hub_download(repo_id=repo_id, filename="config.json", token=TOKEN, cache_dir=tmp_path)
        with open(tmp_config_path) as f:
            assert json.load(f) == CONFIG

        # Delete repo
        api.delete_repo(repo_id=repo_id)

    def test_generate_model_card(self, tmp_path):
        model = DummyModelWithModelCard()
        card = model.generate_model_card()
        assert card.data.language == ["en", "zh"]
        assert card.data.library_name == "my-dummy-lib"
        assert card.data.license == "apache-2.0"
        assert card.data.pipeline_tag == "text-classification"
        assert card.data.tags == ["model_hub_mixin", "pytorch_model_hub_mixin", "tag1", "tag2"]
        # Model card template has been used
        assert "This is a dummy model card" in str(card)

        model.save_pretrained(tmp_path)
        card_reloaded = ModelCard.load(tmp_path / "README.md")

        assert str(card) == str(card_reloaded)
        assert card.data == card_reloaded.data

    def test_load_no_config(self, tmp_path):
        config_file = tmp_path / "config.json"

        # Test creating model => auto-generated config
        model = DummyModelNoConfig(num_classes=50)
        assert model._hub_mixin_config == {"num_classes": 50, "state": "layernorm"}

        # Test saving model => auto-generated config is saved
        model.save_pretrained(tmp_path)
        assert config_file.exists()
        assert json.loads(config_file.read_text()) == {"num_classes": 50, "state": "layernorm"}

        # Reload model => config is reloaded
        reloaded = DummyModelNoConfig.from_pretrained(tmp_path)
        assert reloaded.num_classes == 50
        assert reloaded.state == "layernorm"
        assert reloaded._hub_mixin_config == {"num_classes": 50, "state": "layernorm"}

        # Reload model with custom config => custom config is used
        reloaded_with_default = DummyModelNoConfig.from_pretrained(tmp_path, state="other")
        assert reloaded_with_default.num_classes == 50
        assert reloaded_with_default.state == "other"
        assert reloaded_with_default._hub_mixin_config == {"num_classes": 50, "state": "other"}

        config_file.unlink()  # Remove config file
        reloaded_with_default.save_pretrained(tmp_path)
        assert json.loads(config_file.read_text()) == {"num_classes": 50, "state": "other"}

    def test_save_with_non_jsonable_config(self, tmp_path):
        # Save with a non-jsonable value
        my_object = object()
        model = DummyModelNoConfig(not_jsonable=my_object)
        assert model.not_jsonable is my_object
        assert "not_jsonable" not in model._hub_mixin_config

        # Reload with default value
        model.save_pretrained(tmp_path)
        reloaded_model = DummyModelNoConfig.from_pretrained(tmp_path)
        assert reloaded_model.not_jsonable is DUMMY_OBJECT
        assert "not_jsonable" not in model._hub_mixin_config

        # If jsonable value passed by user, it's saved in the config
        (tmp_path / "config.json").unlink()
        new_model = DummyModelNoConfig(not_jsonable=123)
        new_model.save_pretrained(tmp_path)
        assert new_model._hub_mixin_config["not_jsonable"] == 123

        reloaded_new_model = DummyModelNoConfig.from_pretrained(tmp_path)
        assert reloaded_new_model.not_jsonable == 123
        assert reloaded_new_model._hub_mixin_config["not_jsonable"] == 123

    def test_save_model_with_shared_tensors(self, tmp_path):
        """
        Regression test for #2086. Shared tensors should be saved correctly.

        See https://github.com/huggingface/huggingface_hub/pull/2086 for more details.
        """

        class ModelWithSharedTensors(nn.Module, PyTorchModelHubMixin):
            def __init__(self):
                super().__init__()
                self.a = nn.Linear(100, 100)
                self.b = self.a

            def forward(self, x):
                return self.b(self.a(x))

        # Save and reload model
        model = ModelWithSharedTensors()
        model.save_pretrained(tmp_path)
        reloaded = ModelWithSharedTensors.from_pretrained(tmp_path)

        # Linear layers should share weights and biases in memory
        state_dict = reloaded.state_dict()
        a_weight_ptr = storage_ptr(state_dict["a.weight"])
        b_weight_ptr = storage_ptr(state_dict["b.weight"])
        a_bias_ptr = storage_ptr(state_dict["a.bias"])
        b_bias_ptr = storage_ptr(state_dict["b.bias"])
        assert a_weight_ptr == b_weight_ptr
        assert a_bias_ptr == b_bias_ptr

    def test_save_pretrained_when_config_and_kwargs_are_passed(self, tmp_path):
        # Test creating model with config and kwargs => all values are saved together in config.json
        model = DummyModelWithConfigAndKwargs(num_classes=50, state="layernorm", config={"a": 1}, b=2, c=3)
        model.save_pretrained(tmp_path)
        assert model._hub_mixin_config == {"num_classes": 50, "state": "layernorm", "a": 1, "b": 2, "c": 3}

        reloaded = DummyModelWithConfigAndKwargs.from_pretrained(tmp_path)
        assert reloaded._hub_mixin_config == model._hub_mixin_config

    def test_model_card_with_custom_kwargs(self, tmp_path):
        model_card_kwargs = {"custom_data": "This is a model custom data: 42."}

        # Test creating model with custom kwargs => custom data is saved in model card
        model = DummyModelWithModelCardAndCustomKwargs()
        card = model.generate_model_card(**model_card_kwargs)
        assert model_card_kwargs["custom_data"] in str(card)
        assert "Code: https://hf.co/my-repo" in str(card)
        assert "Paper: https://arxiv.org/abs/2304.12244" in str(card)
        assert "Docs: https://hf.co/docs/my-repo" in str(card)
        # Test saving card => model card is saved and restored with custom data
        model.save_pretrained(tmp_path, model_card_kwargs=model_card_kwargs)
        card_reloaded = ModelCard.load(tmp_path / "README.md")
        assert str(card) == str(card_reloaded)

    def test_config_with_custom_coders(self, tmp_path):
        """
        Regression test for #2334. When `config` is encoded with custom coders, it should be decoded correctly.

        See https://github.com/huggingface/huggingface_hub/issues/2334.
        """
        model = DummyModelWithEncodedConfig(Namespace(a=1, b=2))
        model.save_pretrained(tmp_path)
        assert model._hub_mixin_config["a"] == 1
        assert model._hub_mixin_config["b"] == 2

        reloaded = DummyModelWithEncodedConfig.from_pretrained(tmp_path)
        assert isinstance(reloaded.config, Namespace)
        assert reloaded.config.a == 1
        assert reloaded.config.b == 2

    def test_inheritance_and_sibling_classes(self):
        """
        Test tags are not shared between sibling classes.

        Regression test for #2394.
        See https://github.com/huggingface/huggingface_hub/pull/2394.
        """
        assert DummyModelWithTag1._hub_mixin_info.model_card_data.tags == [
            "model_hub_mixin",
            "pytorch_model_hub_mixin",
            "tag1",
        ]
        assert DummyModelWithTag2._hub_mixin_info.model_card_data.tags == [
            "model_hub_mixin",
            "pytorch_model_hub_mixin",
            "tag2",
        ]
