"""
Unit tests for LawSLM PyTorch Vision Encoder, OCR layout analyzer, and VLM Pipeline.
"""

import unittest
import torch
from PIL import Image
import io
import base64

from slm.vision.encoder import VisionEncoder, VisionPatchEmbedding
from slm.vision.ocr import DocumentVisualAnalyzer
from slm.vision.vlm import LawSLMVisionPipeline
from slm.model.transformer_lm import SLMForCausalLM
from slm.config.model_config import ModelConfig
from slm.tokenizer.bpe import BPETokenizer


class TestVisionPipeline(unittest.TestCase):

    def setUp(self):
        self.config = ModelConfig(vocab_size=200, d_model=64, n_heads=2, n_layers=1)
        self.tokenizer = BPETokenizer()
        self.tokenizer.train_on_texts(["Hello world! Vision test."], vocab_size=200)
        self.config.vocab_size = len(self.tokenizer.vocab)
        self.model = SLMForCausalLM(self.config)

    def test_vision_patch_embedding(self):
        patch_embed = VisionPatchEmbedding(image_size=224, patch_size=16, embed_dim=64)
        dummy_img = torch.randn(2, 3, 224, 224)
        out = patch_embed(dummy_img)
        self.assertEqual(out.shape[0], 2)
        self.assertEqual(out.shape[2], 64)

    def test_vision_encoder_forward(self):
        encoder = VisionEncoder(image_size=224, patch_size=16, d_model=64, n_layers=1, n_heads=2)
        dummy_img = torch.randn(1, 3, 224, 224)
        tokens = encoder(dummy_img)
        self.assertEqual(tokens.shape[0], 1)
        self.assertEqual(tokens.shape[2], 64)

    def test_ocr_and_vlm_pipeline(self):
        analyzer = DocumentVisualAnalyzer()
        img = Image.new("RGB", (300, 400), color="white")
        buf = io.BytesIO()
        img.save(buf, format="PNG")
        b64 = base64.b64encode(buf.getvalue()).decode("utf-8")

        pipeline = LawSLMVisionPipeline(self.model, self.tokenizer)
        res = pipeline.analyze_image_question(image_base64=b64, question="What legal document is this?")

        self.assertIn("question", res)
        self.assertIn("answer", res)
        self.assertIn("ocr", res)
        self.assertEqual(res["ocr"]["dimensions"], "300x400")


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