import os
import yaml
from dotenv import load_dotenv
from openai import OpenAI

class BaseAIConfig:
    def __init__(self):
        load_dotenv()
        self.config = self._load_config()
        self.base_url = self._get_config_value('diona.ai.modelset.base-url')
        self.api_key = self._get_api_key()
        self.model = self._get_config_value('diona.ai.modelset.model')
        self.proxy = self._get_config_value('diona.ai.clientset.proxy', False)
        self.use_proxy = self._get_config_value('diona.ai.clientset.use-proxy', False)
    
    def _load_config(self):
        config_path = os.path.join(os.getcwd(), 'config.yml')
        if os.path.exists(config_path):
            with open(config_path, 'r', encoding='utf-8') as f:
                return yaml.safe_load(f)
        return {}
    
    def _get_config_value(self, key, default=None):
        keys = key.split('.')
        value = self.config
        for k in keys:
            if isinstance(value, dict) and k in value:
                value = value[k]
            else:
                return default
        return value
    
    def _get_api_key(self):
        # Priority: 1. Environment variable 2. .env file 3. config.yml
        api_key = os.environ.get('OPENAI_API_KEY')
        if api_key:
            return api_key
        api_key = self._get_config_value('diona.ai.modelset.api-key')
        return api_key
    
    def get_base_url(self):
        return self.base_url
    
    def get_api_key(self):
        return self.api_key
    
    def get_model(self):
        return self.model
    
    def get_proxy(self):
        return self.proxy
    
    def is_use_proxy(self):
        return self.use_proxy

class BaseChat:
    def __init__(self, config):
        self.config = config
        self.client = self._create_client()
    
    def _create_client(self):
        kwargs = {
            'api_key': self.config.get_api_key(),
            'base_url': self.config.get_base_url()
        }
        if self.config.is_use_proxy() and self.config.get_proxy():
            kwargs['http_client'] = {
                'proxy': self.config.get_proxy()
            }
        return OpenAI(**kwargs)
    
    def chat(self, chat_message):
        response = self.client.chat.completions.create(
            model=self.config.get_model(),
            messages=chat_message.to_ai_messages()
        )
        ai_response = response.choices[0].message
        metadata = {
            'id': response.id,
            'created': response.created,
            'model': response.model,
            'usage': {
                'prompt_tokens': response.usage.prompt_tokens,
                'completion_tokens': response.usage.completion_tokens,
                'total_tokens': response.usage.total_tokens
            }
        }
        return SingleChatMessage(ai_response.role, ai_response.content, metadata)

class SingleChatMessage:
    def __init__(self, role, content, metadata=None):
        self.role = role
        self.content = content
        self.metadata = metadata or {}
    
    @classmethod
    def from_text(cls, text, role='user'):
        return cls(role, text)
    
    def to_ai_message(self):
        return {
            'role': self.role,
            'content': self.content
        }

class BaseChatHistory:
    def __init__(self, max_messages=6):
        self.max_messages = max_messages
        self.messages = []
    
    def add_message(self, message):
        self.messages.append(message)
        if len(self.messages) > self.max_messages:
            self.messages.pop(0)
    
    def get_messages(self):
        return self.messages
    
    def get_chat_message(self):
        return ChatMessage(self.messages)
    
    def persist(self):
        # To be implemented by subclass
        pass

class ClientConfig:
    def __init__(self):
        self.config = self._load_config()
        self.prompt_path = self._get_config_value('diona.ai.clientset.prompt')
        self.system_prompt = self._load_system_prompt()
    
    def _load_config(self):
        config_path = os.path.join(os.getcwd(), 'config.yml')
        if os.path.exists(config_path):
            with open(config_path, 'r', encoding='utf-8') as f:
                return yaml.safe_load(f)
        return {}
    
    def _get_config_value(self, key, default=None):
        keys = key.split('.')
        value = self.config
        for k in keys:
            if isinstance(value, dict) and k in value:
                value = value[k]
            else:
                return default
        return value
    
    def _load_system_prompt(self):
        if self.prompt_path and os.path.exists(self.prompt_path):
            with open(self.prompt_path, 'r', encoding='utf-8') as f:
                return f.read().strip()
        return "You are a helpful assistant"
    
    def get_system_prompt(self):
        return self.system_prompt

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
    
    @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)

class ChatMessage:
    def __init__(self, messages):
        self.messages = messages
    
    def to_ai_messages(self):
        return [msg.to_ai_message() for msg in self.messages]

def interactive_chat():
    client = BaseClient.create()
    session_id = "default"
    
    print("Welcome to Diona AI Chat!")
    print("Type 'exit' to quit.")
    
    while True:
        user_input = input("<diona@user> ")
        if user_input.strip().lower() == "exit":
            print("Goodbye!")
            break
        
        response = client.chat(session_id, user_input)
        print(f"<@assistant> {response.content}")
        print()
