"""Check the published SDK router and adapter without downloading weights."""
import unittest
from unittest.mock import patch

import torch
from transformers import Qwen2Config, Qwen2ForCausalLM

from core.engine import run_parallel_generation
from core.schema import StructuredSchema


class Tokenizer:
    pad_token_id = 0

    def encode(self, text, return_tensors=None, **kwargs):
        ids = [1 + ord(c) % 95 for c in text]
        return torch.tensor([ids]) if return_tensors == 'pt' else ids


class SDKTest(unittest.TestCase):
    def test_tree_is_default_and_preserves_response(self):
        torch.set_num_threads(2)
        torch.manual_seed(42)
        config = Qwen2Config(vocab_size=97, hidden_size=64, intermediate_size=128,
                             num_hidden_layers=2, num_attention_heads=4, num_key_value_heads=2)
        config._attn_implementation = 'sdpa'
        model = Qwen2ForCausalLM(config).eval()
        schema = StructuredSchema({'risk': {'type':'boolean', 'description':'Risk present'},
                                   'level': {'type':'enum', 'description':'Risk level', 'choices':['LOW','HIGH']}})
        self.assertEqual(run_parallel_generation.__module__, 'core.engine_tree')
        with patch('core.engine_tree.get_torch_engine', return_value=(model, Tokenizer(), 'cpu')):
            result = run_parallel_generation('A suspicious request.', schema)
        self.assertEqual(result['mode'], 'parallel_constrained_tree')
        self.assertEqual(set(result['parsed_json']), {'risk', 'level'})
        self.assertIsInstance(result['parsed_json']['risk']['value'], bool)
        self.assertIn(result['parsed_json']['level']['value'], ['LOW', 'HIGH'])
        self.assertFalse(result['has_calibrated_probabilities'])
        self.assertEqual(result['num_fields'], 2)


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