import unittest

from transformers import AutoTokenizer

import taoverse.model.storage.disk.utils as utils
from taoverse.model.data import Model, ModelId
from taoverse.model.storage.disk.disk_model_store import DiskModelStore
from tests.utils import get_test_model


class TestDiskModelStore(unittest.TestCase):
    def setUp(self):
        self.base_dir = "test-models"
        self.disk_store = DiskModelStore(self.base_dir)
        self.tiny_model = get_test_model()

    def tearDown(self):
        # Clean up the model directory.
        self.disk_store.delete_unreferenced_models(dict(), 0)

    def test_get_path(self):
        hotkey = "hotkey0"

        expected_path = utils.get_local_miner_dir("test-models", hotkey)
        actual_path = self.disk_store.get_path(hotkey)

        self.assertEqual(expected_path, actual_path)

    def test_store_and_retrieve_model(self):
        hotkey = "hotkey0"
        model_id = ModelId(
            namespace="TestPath",
            name="TestModel",
            competition_id=1,
            secure_hash="TestHash1",
            commit="TestCommit",
        )

        pt_model = self.tiny_model

        model = Model(id=model_id, pt_model=pt_model)

        # Store the model locally.
        self.disk_store.store_model(hotkey, model)

        # Retrieve the model locally.
        retrieved_model = self.disk_store.retrieve_model(hotkey, model_id, kwargs={})

        # Check that they match.
        self.assertEqual(str(model), str(retrieved_model))

    def test_store_and_retrieve_model_with_tokenizer(self):
        hotkey = "hotkeyTok"
        model_id = ModelId(
            namespace="TestPathTok",
            name="TestModelTok",
            competition_id=1,
            secure_hash="TestHash1Tok",
            commit="TestCommitTok",
        )

        pt_model = self.tiny_model
        tokenizer = AutoTokenizer.from_pretrained("Xenova/gpt-4")

        model = Model(id=model_id, pt_model=pt_model, tokenizer=tokenizer)

        # Store the model locally.
        self.disk_store.store_model(hotkey, model)

        # Retrieve the model locally.
        retrieved_model = self.disk_store.retrieve_model(hotkey, model_id, kwargs={})

        # Overwrite the name of the tokenizer to avoid it using the local path.
        retrieved_model.tokenizer.name_or_path = "Xenova/gpt-4"

        # Check that they match.
        self.assertEqual(str(model), str(retrieved_model))

    def test_delete_unreferenced_models(self):
        hotkey = "hotkey0"

        # Make 2 model ids with different hashes / commits.
        model_id_1 = ModelId(
            namespace="TestPath",
            name="TestModel",
            competition_id=1,
            secure_hash="TestHash1",
            commit="TestCommit",
        )
        model_id_2 = ModelId(
            namespace="TestPath",
            name="TestModel",
            competition_id=1,
            secure_hash="TestHash2",
            commit="TestCommit2",
        )

        pt_model = self.tiny_model

        model_1 = Model(id=model_id_1, pt_model=pt_model)
        model_2 = Model(id=model_id_2, pt_model=pt_model)

        # Store both models locally.
        self.disk_store.store_model(hotkey, model_1)
        self.disk_store.store_model(hotkey, model_2)

        # Create the mapping of hotkey to model_id with only the 2nd model.
        valid_models_by_hotkey = dict()
        valid_models_by_hotkey[hotkey] = model_id_2

        # Clear the unreferenced models
        self.disk_store.delete_unreferenced_models(valid_models_by_hotkey, 0)

        # Confirm that model 1 is deleted
        with self.assertRaises(Exception):
            self.disk_store.retrieve_model(hotkey, model_id_1, kwargs={})

        # Confirm that model 2 is still there
        model_2_retrieved = self.disk_store.retrieve_model(
            hotkey, model_id_2, kwargs={}
        )
        self.assertEqual(str(model_2), str(model_2_retrieved))

    def test_delete_unreferenced_models_removed_hotkey(self):
        hotkey_1 = "hotkey1"
        hotkey_2 = "hotkey2"

        # Make 2 model ids with different hashes / commits.
        model_id_1 = ModelId(
            namespace="TestPath",
            name="TestModel",
            competition_id=1,
            secure_hash="TestHash1",
            commit="TestCommit",
        )
        model_id_2 = ModelId(
            namespace="TestPath",
            name="TestModel",
            competition_id=1,
            secure_hash="TestHash2",
            commit="TestCommit2",
        )

        pt_model = self.tiny_model

        model_1 = Model(id=model_id_1, pt_model=pt_model)
        model_2 = Model(id=model_id_2, pt_model=pt_model)

        # Store both models locally.
        self.disk_store.store_model(hotkey_1, model_1)
        self.disk_store.store_model(hotkey_2, model_2)

        # Create the mapping of hotkey to model_id with only the 2nd model.
        valid_models_by_hotkey = dict()
        valid_models_by_hotkey[hotkey_2] = model_id_2

        # Clear the unreferenced models
        self.disk_store.delete_unreferenced_models(valid_models_by_hotkey, 0)

        # Confirm that model 1 is deleted
        with self.assertRaises(Exception):
            self.disk_store.retrieve_model(hotkey_1, model_id_1, kwargs={})

        # Confirm that model 2 is still there
        model_2_retrieved = self.disk_store.retrieve_model(
            hotkey_2, model_id_2, kwargs={}
        )
        self.assertEqual(str(model_2), str(model_2_retrieved))

    def test_delete_unreferenced_models_in_grace(self):
        hotkey = "hotkey0"

        # Make 2 model ids with different hashes / commits.
        model_id_1 = ModelId(
            namespace="TestPath",
            name="TestModel",
            competition_id=1,
            secure_hash="TestHash1",
            commit="TestCommit",
        )
        model_id_2 = ModelId(
            namespace="TestPath",
            name="TestModel",
            competition_id=1,
            secure_hash="TestHash2",
            commit="TestCommit2",
        )

        pt_model = self.tiny_model

        model_1 = Model(id=model_id_1, pt_model=pt_model)
        model_2 = Model(id=model_id_2, pt_model=pt_model)

        # Store both models locally.
        self.disk_store.store_model(hotkey, model_1)
        self.disk_store.store_model(hotkey, model_2)

        # Create the mapping of hotkey to model_id with only the 2nd model.
        valid_models_by_hotkey = dict()
        valid_models_by_hotkey[hotkey] = model_id_2

        # Clear the unreferenced models
        self.disk_store.delete_unreferenced_models(valid_models_by_hotkey, 60)

        # Confirm that model 1 is still there.
        model_1_retrieved = self.disk_store.retrieve_model(
            hotkey, model_id_1, kwargs={}
        )
        self.assertEqual(str(model_1), str(model_1_retrieved))

        # Confirm that model 2 is still there
        model_2_retrieved = self.disk_store.retrieve_model(
            hotkey, model_id_2, kwargs={}
        )
        self.assertEqual(str(model_2), str(model_2_retrieved))

    def test_delete_unreferenced_models_removed_hotkey_in_grace(self):
        hotkey_1 = "hotkey1"
        hotkey_2 = "hotkey2"

        # Make 2 model ids with different hashes / commits.
        model_id_1 = ModelId(
            namespace="TestPath",
            name="TestModel",
            competition_id=1,
            secure_hash="TestHash1",
            commit="TestCommit",
        )
        model_id_2 = ModelId(
            namespace="TestPath",
            name="TestModel",
            competition_id=1,
            secure_hash="TestHash2",
            commit="TestCommit2",
        )

        pt_model = self.tiny_model

        model_1 = Model(id=model_id_1, pt_model=pt_model)
        model_2 = Model(id=model_id_2, pt_model=pt_model)

        # Store both models locally.
        self.disk_store.store_model(hotkey_1, model_1)
        self.disk_store.store_model(hotkey_2, model_2)

        # Create the mapping of hotkey to model_id with only the 2nd model.
        valid_models_by_hotkey = dict()
        valid_models_by_hotkey[hotkey_2] = model_id_2

        # Clear the unreferenced models
        self.disk_store.delete_unreferenced_models(valid_models_by_hotkey, 60)

        # Confirm that model 1 is still there
        model_1_retrieved = self.disk_store.retrieve_model(
            hotkey_1, model_id_1, kwargs={}
        )
        self.assertEqual(str(model_1), str(model_1_retrieved))

        # Confirm that model 2 is still there
        model_2_retrieved = self.disk_store.retrieve_model(
            hotkey_2, model_id_2, kwargs={}
        )
        self.assertEqual(str(model_2), str(model_2_retrieved))


if __name__ == "__main__":
    unittest.main()
