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


class AnthropicLlmGatewayImpl(LlmGateway):
    def __init__(
        self,
        system_prompt: str | None = None,
        max_tokens: int = 300,
        model: str | None = None,
    ):
        api_key = infra_config.get("ANTHROPIC_API_KEY", "")
        if not api_key:
            raise ValueError("ANTHROPIC_API_KEY must be set for Anthropic gateway")
        if not model:
            raise ValueError("Model name must be provided for Anthropic gateway")
        self.model = model
        self.default_system_prompt = system_prompt or ""
        self.default_max_tokens = min(max_tokens, 300)
        self.client = anthropic.Anthropic(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:
        max_tokens = self.default_max_tokens
        final_system_prompt = self.default_system_prompt
        response = self.client.messages.create(
            model=self.model,
            max_tokens=max_tokens,
            system=final_system_prompt,
            messages=[{"role": "user", "content": user_message}],
        )
        parts = []
        for block in response.content or []:
            try:
                text = block.text
            except AttributeError:
                text = None
            if text:
                parts.append(text)
        return "".join(parts).strip()

    async def generate_streaming_completion(
        self,
        user_message: str,
    ) -> AsyncGenerator[str, None]:
        max_tokens = self.default_max_tokens
        final_system_prompt = self.default_system_prompt
        stream = self.client.messages.stream(
            model=self.model,
            max_tokens=max_tokens,
            system=final_system_prompt,
            messages=[{"role": "user", "content": user_message}],
        )

        with stream as response_stream:
            for event in response_stream:
                try:
                    event_type = event.type
                except AttributeError:
                    event_type = None
                if event_type != "content_block_delta":
                    continue
                try:
                    delta = event.delta
                except AttributeError:
                    delta = None
                if delta is None:
                    continue
                try:
                    content = delta.text
                except AttributeError:
                    content = None
                if content:
                    yield content


if __name__ == "__main__":
    gateway = AnthropicLlmGatewayImpl(model="claude-sonnet-4-6")
    result = gateway.generate_completion("Hello, how are you?")
    print(result)