from .base_ai_config import BaseAIConfig
from .base_chat import BaseChat
from .base_chat_history import BaseChatHistory
from .client_config import ClientConfig
from .message import SingleChatMessage, ChatMessage

class BaseClient:
    def __init__(self, base_chat, base_chat_history, client_config):
        self.base_chat = base_chat
        self.base_chat_history = base_chat_history
        self.client_config = client_config
        self.session_map = {}
    
    def chat(self, session_id, user_input):
        if isinstance(user_input, str):
            user_message = SingleChatMessage.from_text(user_input, 'user')
        else:
            user_message = user_input
        
        if session_id not in self.session_map:
            self.session_map[session_id] = BaseChatHistory()
        
        chat_history = self.session_map[session_id]
        chat_history.add_message(user_message)
        
        # Get current messages and add system prompt
        current_messages = chat_history.get_messages()
        system_prompt = self.client_config.get_system_prompt()
        system_message = SingleChatMessage('system', system_prompt)
        
        # Create chat message with system prompt at the beginning
        messages_with_system = [system_message] + current_messages
        chat_message = ChatMessage(messages_with_system)
        
        # Get AI response
        ai_response = self.base_chat.chat(chat_message)
        
        # Add AI response to history
        chat_history.add_message(ai_response)
        
        return ai_response
    
    def chat_stream(self, session_id, user_input):
        if isinstance(user_input, str):
            user_message = SingleChatMessage.from_text(user_input, 'user')
        else:
            user_message = user_input
        
        if session_id not in self.session_map:
            self.session_map[session_id] = BaseChatHistory()
        
        chat_history = self.session_map[session_id]
        chat_history.add_message(user_message)
        
        # Get current messages and add system prompt
        current_messages = chat_history.get_messages()
        system_prompt = self.client_config.get_system_prompt()
        system_message = SingleChatMessage('system', system_prompt)
        
        # Create chat message with system prompt at the beginning
        messages_with_system = [system_message] + current_messages
        chat_message = ChatMessage(messages_with_system)
        
        # Get AI response stream
        ai_response = None
        for chunk in self.base_chat.chat_stream(chat_message):
            if isinstance(chunk, SingleChatMessage):
                # This is the final message object
                ai_response = chunk
            else:
                # This is a content chunk
                yield chunk
        
        # Add AI response to history
        if ai_response:
            chat_history.add_message(ai_response)
            yield ai_response
    
    @classmethod
    def create(cls):
        ai_config = BaseAIConfig()
        base_chat = BaseChat(ai_config)
        base_chat_history = BaseChatHistory()
        client_config = ClientConfig()
        return cls(base_chat, base_chat_history, client_config)