import unittest

from itertools import tee

from . import test_data

from sia.llm_engine import LlmEngine
from sia.llm_engine.local_llm_engine import LocalLlmEngine

class LocalLlmEngineTest(unittest.TestCase):
    def setUp(self):
        self.model_path = "/root/models/current"

    @unittest.skip("Takes too long")
    def test_initialization(self):
        llm_engine = LocalLlmEngine(self.model_path, 0.2, 1024)
        self.assertIsInstance(llm_engine, LlmEngine)

    @unittest.skip("Takes too long")
    def test_infer(self):
        main_context = "This is a test"
        llm_engine = LocalLlmEngine(self.model_path, 0.2, 1024)
        tokens = llm_engine.infer(test_data.echo_system_prompt, main_context, "")
        print_tokens, result_tokens = tee(tokens)
        for token in print_tokens:
            print(token, end="", flush=True)
        result = ''.join(result_tokens)
        self.assertEqual(result, f"{main_context}<test_tag>{main_context}</test_tag>")

    @unittest.skip("Takes too long")
    def test_token_count(self):
        """Test token counting returns reasonable numbers"""
        llm_engine = LocalLlmEngine(self.model_path, 0.2, 1024)
        
        # Test with short inputs
        count = llm_engine.token_count("Short prompt", "Brief context")
        self.assertGreater(count, 5)
        self.assertLess(count, 50)
        
        # Test with longer inputs
        long_prompt = "A detailed system prompt with multiple sentences. " * 5
        long_context = "An extensive context containing various details. " * 10
        count = llm_engine.token_count(long_prompt, long_context)
        self.assertGreater(count, 100)
        self.assertLess(count, 1000)

    @unittest.skip("Takes too long")
    def test_token_limit(self):
        """Test token limit is within expected range for modern LLMs"""
        llm_engine = LocalLlmEngine(self.model_path, 0.2, 1024)
        limit = llm_engine.token_limit()
        self.assertGreaterEqual(limit, 2048)
        self.assertLessEqual(limit, 1048576)