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


class AzureLlmGatewayImpl(LlmGateway):
    def __init__(
        self,
        system_prompt: str | None = None,
        max_tokens: int = 300,
        model: str | None = None,
        azure_client: AzureOpenAI | None = None,
    ):
        api_key = infra_config.get("AZURE_AI_FOUNDRY_API_KEY", "")
        if not api_key:
            raise ValueError("AZURE_AI_FOUNDRY_API_KEY must be set for Azure gateway")
        if not model:
            raise ValueError("Model name must be provided for Azure gateway")
        timeout_seconds = float(infra_config.get("LLM_TIMEOUT_SECONDS", 60))
        self.model = model
        self.default_system_prompt = system_prompt or ""
        self.default_max_tokens = min(max_tokens, 300)
        endpoint = infra_config.get(
            "AZURE_AI_FOUNDRY_ENDPOINT",
            "https://analogai-llms.cognitiveservices.azure.com/",
        )
        self.client = azure_client or AzureOpenAI(
            api_version="2024-12-01-preview",
            azure_endpoint=endpoint,
            api_key=api_key,
            timeout=timeout_seconds,
        )

    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 = AzureLlmGatewayImpl(model="gpt-4o")
    result = gateway.generate_completion("Hello, how are you?")
    print(result)