import unittest
import os
import tempfile
import time
from xml.etree import ElementTree as ET

from sia.llm_engine import LlmEngine

class TestLlmEngine(unittest.TestCase):
    def setUp(self):
        """Set up test environment before each test."""
        # Create a temporary bash script to simulate the LLM engine
        self.temp_script = self._create_mock_llm_script()
        os.chmod(self.temp_script, 0o755)  # Make it executable
        
        # Initialize the LlmEngine with our mock script
        self.engine = LlmEngine(self.temp_script, "/root/sia/action_schema.xsd")
    
    def tearDown(self):
        """Clean up after the test."""
        if hasattr(self, 'engine') and self.engine:
            self.engine._terminate_process()
    
    def _create_mock_llm_script(self):
        """Create a bash script that simulates an LLM engine."""
        fd, path = tempfile.mkstemp(suffix='.sh')
        with os.fdopen(fd, 'w') as f:
            # Paste the content from the mock_llm_script_debug artifact
            # This is a placeholder - you'll replace this with the actual script content
            f.write('''#!/bin/bash

# Mock LLM engine for testing
# This script simulates an LLM engine subprocess that responds to the three commands:
# <token_limit/>, <token_count>...</token_count>, and <infer_xml>...</infer_xml>

# Function to read XML input until a complete closing tag is found
read_xml_input() {
    local input=""
    local line
    local char_count=0
    
    # Read until we get a complete XML command
    while IFS= read -r line; do
        input="${input}${line}"
        char_count=$((char_count + ${#line}))
        
        # Debug the actual content (first 30 chars)
        
        if [[ "$input" == *"<token_limit/>"* ]]; then
            printf "1024" 
            printf "\004"  # EOT character (hex 04)
            return
        fi
        
        if [[ "$input" == *"<token_count>"*"</token_count>"* ]]; then
            printf "405"
            printf "\004"  # EOT character (hex 04)
            return
        fi
        
        if [[ "$input" == *"<infer_xml>"*"</infer_xml>"* ]]; then
            generate_response
            return
        fi
        
    done
}

# Function to generate a response token by token
generate_response() {
    printf "<"
    sleep 0.01
    
    printf "reason"
    sleep 0.01
    
    printf "ing"
    sleep 0.01
    
    printf ">"
    sleep 0.01
    
    printf "This"
    sleep 0.01
    
    printf " is"
    sleep 0.01
    
    printf " a"
    sleep 0.01
    
    printf " test"
    sleep 0.01
    
    printf " response."
    sleep 0.01
    
    printf "</"
    sleep 0.01
    
    printf "reason"
    sleep 0.01
    
    printf "ing"
    sleep 0.01
    
    printf ">"
    sleep 0.01
    
    printf "\004"
}

# Main loop - keep reading input and responding
iteration=0
while true; do
    iteration=$((iteration + 1))
    read_xml_input
done''')
        return path

    def test_token_limit(self):
        """Test retrieving the token limit from the LLM engine."""
        limit = self.engine.token_limit()
        self.assertEqual(limit, 1024, "Token limit should be 1024")

    def test_token_count(self):
        """Test counting tokens in a prompt."""
        system_prompt = "You are a helpful AI assistant."
        
        # Create a simple context ElementTree
        context_et = ET.Element("context")
        context_et.set("time", "2024-10-18T12:00:00Z")
        context_et.text = "Test context"
        
        prefix = ""
        
        count = self.engine.token_count(system_prompt, context_et, prefix)
        self.assertEqual(count, 405, "Token count should be 405")

    def test_inference_token_by_token(self):
        """Test inference with the LLM engine, ensuring tokens come one by one."""
        system_prompt = "You are a helpful AI assistant."
        
        # Create a simple context ElementTree
        context_et = ET.Element("context")
        context_et.set("time", "2024-10-18T12:00:00Z")
        context_et.text = "Test context"
        
        prefix = ""
        
        # Get the generator for token-by-token output
        token_generator = self.engine.infer(system_prompt, context_et, prefix)
        
        # Collect tokens
        response = ""
        tokens_received = 0
        
        for token in token_generator:
            response += token
            tokens_received += 1
            # Each "token" should be small (in the mock script, it's up to 10 characters)
            self.assertLessEqual(len(token), 10, "Each token should be small")
        
        # Verify we received multiple tokens
        self.assertGreater(tokens_received, 10, "Should receive multiple tokens")
        
        expected = "<reasoning>This is a test response.</reasoning>"
        self.assertEqual(response, expected, "Inference response should match expected output")

    def test_multiple_inferences(self):
        """Test running multiple inference cycles one after another."""
        system_prompt = "You are a helpful AI assistant."
        
        # Create a simple context ElementTree
        context_et = ET.Element("context")
        context_et.set("time", "2024-10-18T12:00:00Z")
        context_et.text = "Test context"
        
        prefix = ""
        
        # Run first inference
        token_generator = self.engine.infer(system_prompt, context_et, prefix)
        response = ""
        for token in token_generator:
            response += token
        
        expected = "<reasoning>This is a test response.</reasoning>"
        self.assertEqual(response, expected, "First inference response should match expected output")
        
        # Add a small delay between inferences to see if it helps
        time.sleep(0.1)
        
        # Run second inference
        token_generator = self.engine.infer(system_prompt, context_et, prefix)
        response = ""
        for token in token_generator:
            response += token
        
        self.assertEqual(response, expected, "Second inference response should match expected output")
        
        # We should be able to get the token limit again after inference
        limit = self.engine.token_limit()
        self.assertEqual(limit, 1024, "Token limit should still be 1024 after multiple inferences")

    def test_restart(self):
        """Test restarting the LLM engine."""
        # First get a token limit
        limit = self.engine.token_limit()
        self.assertEqual(limit, 1024, "Token limit should be 1024")
        
        # Restart the engine
        self.engine.restart()
        
        # Get token limit again
        limit = self.engine.token_limit()
        self.assertEqual(limit, 1024, "Token limit should still be 1024 after restart")
        
        # Engine should still work for inference after restart
        system_prompt = "You are a helpful AI assistant."
        
        # Create a simple context ElementTree
        context_et = ET.Element("context")
        context_et.set("time", "2024-10-18T12:00:00Z")
        context_et.text = "Test context"
        
        prefix = ""
        
        token_generator = self.engine.infer(system_prompt, context_et, prefix)
        response = ""
        for token in token_generator:
            response += token
        
        expected = "<reasoning>This is a test response.</reasoning>"
        self.assertEqual(response, expected, "Inference should work after restart")

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