from typing import AsyncGenerator

try:
    from xai_sdk import Client
    from xai_sdk.chat import system, user
except ImportError:  # pragma: no cover - optional dependency
    Client = None
    system = None
    user = None
from analogai.deepthink.infrastructure.gateways.llm import config
from analogai.deepthink.infrastructure.gateways.llm.gateway import LlmGateway


class XaiLlmGatewayImpl(LlmGateway):
    def __init__(
        self,
        system_prompt: str | None = None,
        max_tokens: int = 500,
        model: str | None = None,
    ):
        if Client is None or system is None or user is None:
            raise ImportError(
                "xai-sdk is required for the XAI gateway. Install it with 'pip install xai-sdk'."
            )
        api_key = config.XAI_API_KEY
        if not api_key:
            raise ValueError("XAI_API_KEY must be set for XAI gateway")
        if not model:
            raise ValueError("Model name must be provided for XAI gateway")
        self.model = model
        self.default_system_prompt = system_prompt or ""
        self.default_max_tokens = min(max_tokens, 500)
        self.client = Client(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
        chat = self.client.chat.create(
            model=self.model,
            messages=[
                system(final_system_prompt),
                user(user_message),
            ],
        )
        response = chat.sample()
        content = getattr(response, "content", None)
        if content is None:
            return ""
        return str(content).strip()

    async def generate_streaming_completion(
        self,
        user_message: str,
    ) -> AsyncGenerator:
        output = self.generate_completion(user_message)
        if output:
            yield output

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


if __name__ == "__main__":
    gateway = XaiLlmGatewayImpl(model="grok-3")
    result = gateway.generate_completion("Hello, how are you?")
    print(result)