# -*- coding: utf-8 -*-

import json
import os
import tempfile
import unittest

from Tokenizer.morphbpe import MorphBPETrainer
from Tokenizer.unified.bundle import TokenizerBundle


class TokenizerBundleTest(unittest.TestCase):
    def _train_tiny_morphbpe(self, tmp: str) -> str:
        trainer = MorphBPETrainer(vocab_size=200, min_pair_freq=1)
        tokenizer = trainer.train(["ᠮᠣᠩᠭᠣᠯ ᠪᠢᠴᠢᠭ", "ᠮᠣᠩᠭᠣᠯ text"])
        path = os.path.join(tmp, "tiny_morphbpe.json")
        tokenizer.save(path)
        return path

    def test_bundle_save_load_encode_and_validate(self):
        with tempfile.TemporaryDirectory() as tmp:
            morphbpe_path = self._train_tiny_morphbpe(tmp)
            bundle = TokenizerBundle.from_files(
                morphbpe_path,
                zh_source="smoke-zh",
                en_source="smoke-en",
                use_smoke_hf=True,
            )
            out_dir = os.path.join(tmp, "bundle")
            bundle.save_dir(out_dir)

            self.assertTrue(os.path.exists(os.path.join(out_dir, "config.json")))
            self.assertTrue(os.path.exists(os.path.join(out_dir, "morphbpe.json")))
            self.assertTrue(os.path.exists(os.path.join(out_dir, "vocab.json")))
            self.assertTrue(os.path.exists(os.path.join(out_dir, "manifest.json")))
            with open(os.path.join(out_dir, "manifest.json"), "r", encoding="utf-8") as f:
                manifest = json.load(f)
            self.assertIn("vocab.json", manifest["files"])

            loaded = TokenizerBundle.from_dir(out_dir)
            self.assertEqual(loaded.validate(), [])

            encoded = loaded.encode_with_spans("ᠮᠣᠩᠭᠣᠯ 文字 test", add_bos=True, add_eos=True)
            self.assertEqual(len(encoded.input_ids), len(encoded.tokens))
            self.assertGreater(len(encoded.input_ids), 4)
            self.assertEqual(encoded.tokens[0].token, "<bos>")
            self.assertEqual(encoded.tokens[-1].token, "<eos>")

            mm = loaded.encode_multimodal(
                "文字 <image> test",
                images=["img"],
                image_sizes=[(14, 14)],
            )
            self.assertEqual(len(mm.image_token_spans), 1)
            start, end = mm.image_token_spans[0]
            self.assertEqual([tok.token for tok in mm.tokens[start:end]], [
                "<image_start>",
                "<image_patch>",
                "<image_end>",
            ])
            self.assertEqual(len(mm.attention_mask), len(mm.input_ids))

            with open(os.path.join(out_dir, "vocab.json"), "a", encoding="utf-8") as f:
                f.write("\n")
            issues = loaded.validate()
            self.assertTrue(
                any("manifest hash mismatch for vocab.json" in issue for issue in issues)
            )


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