from openai import OpenAI
from .message import SingleChatMessage

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['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)
    
    def chat_stream(self, chat_message):
        response = self.client.chat.completions.create(
            model=self.config.get_model(),
            messages=chat_message.to_ai_messages(),
            stream=True
        )
        
        full_content = ""
        role = None
        metadata = {
            'id': None,
            'created': None,
            'model': None
        }
        
        for chunk in response:
            if chunk.choices:
                choice = chunk.choices[0]
                if choice.delta:
                    if choice.delta.role:
                        role = choice.delta.role
                    if choice.delta.content:
                        content_chunk = choice.delta.content
                        full_content += content_chunk
                        yield content_chunk
            
            # Capture metadata from the first chunk
            if not metadata['id']:
                metadata['id'] = chunk.id
                metadata['created'] = chunk.created
                metadata['model'] = chunk.model
        
        # After streaming is complete, yield the final message object
        if role and full_content:
            yield SingleChatMessage(role, full_content, metadata)