from itertools import tee
from typing import Optional, Dict, List

from .docker_module import DockerModule
from .llm_engine import LlmEngine
from .context_template import generate_context
from .util import get_valid_root_elements, split_response

class AgentCore:
    """
    Core orchestration class for SIA that manages the interaction between
    different modules and runs the main agent loop.
    """
    
    def __init__(
        self, 
        system_prompt: str,
        action_schema: str,
        docker_module: DockerModule,
        llm_engine: LlmEngine
    ):
        """
        Initialize the AgentCore with required components.
        
        Args:
            system_prompt: System prompt to use for the LLM
            action_schema_path: Path to the XML schema defining valid actions
            docker_module: DockerModule instance
            llm_engine: LLmEngine instance
        """
        self.system_prompt = system_prompt
        self.action_schema = action_schema
        self.docker_module = docker_module
        self.llm_engine = llm_engine
        
        self.valid_elements = get_valid_root_elements(self.action_schema)
    
    def run_iteration(self) -> None:
        """Run a single iteration of the main agent loop."""
        containers = self.docker_module.get_all_container_statuses()
        context = generate_context(containers)
        tokens = self.llm_engine.infer(
            self.system_prompt,
            context
        )
        print_tokens, response_tokens = tee(tokens)
        for token in print_tokens:
            print(token, end="", flush=True)
        response_tokens = ''.join(response_tokens)
        result = split_response(response_tokens, self.valid_elements)
        print(f"result: {result}")