from typing import AsyncGenerator
from openai import OpenAI
from analogai.foundation.settings import config as infra_config
from analogai.deepthink.infrastructure.gateways.llm.gateway import LlmGateway


class OpenaiLlmGatewayImpl(LlmGateway):
    def __init__(
        self,
        system_prompt: str | None = None,
        max_tokens: int = 300,
        model: str | None = None,
    ):
        api_key = infra_config.get("OPENAI_API_KEY", "")
        if not api_key:
            raise ValueError("OPENAI_API_KEY must be set for OpenAI gateway")
        if not model:
            raise ValueError("Model name must be provided for OpenAI gateway")
        self.model = model
        self.default_system_prompt = system_prompt or ""
        self.default_max_tokens = min(max_tokens, 300)
        self.client = OpenAI(api_key=api_key)

    def set_model(self, model: str) -> None:
        self.model = model

    def set_system_prompt(self, system_prompt: str) -> None:
        self.default_system_prompt = system_prompt

    def generate_completion(
        self,
        user_message: str,
    ) -> str:
        final_system_prompt = self.default_system_prompt
        max_tokens = self.default_max_tokens
        response = self.client.chat.completions.create(
            messages=[
                {"role": "system", "content": final_system_prompt},
                {"role": "user", "content": user_message},
            ],
            max_tokens=max_tokens,
            temperature=1.0,
            top_p=1.0,
            model=self.model,
        )
        return response.choices[0].message.content.strip()

    async def generate_streaming_completion(
        self,
        user_message: str,
    ) -> AsyncGenerator:
        final_system_prompt = self.default_system_prompt
        max_tokens = self.default_max_tokens
        stream = self.client.chat.completions.create(
            stream=True,
            messages=[
                {"role": "system", "content": final_system_prompt},
                {"role": "user", "content": user_message},
            ],
            max_tokens=max_tokens,
            temperature=1.0,
            top_p=1.0,
            model=self.model,
        )
        for chunk in stream:
            yield chunk

    def close(self) -> None:
        try:
            self.client.close()
        except AttributeError:
            pass


if __name__ == "__main__":
    gateway = OpenaiLlmGatewayImpl(model="gpt-4o-mini")
    result = gateway.generate_completion("Hello, how are you?")
    print(result)