from __future__ import annotations
import anthropic


INPUT_COST_PER_1M = 3.0   # claude-sonnet-4-6 $/1M tokens
OUTPUT_COST_PER_1M = 15.0


class ClaudeClient:
    MODEL = "claude-sonnet-4-6"

    def __init__(self, api_key: str):
        self._client = anthropic.Anthropic(api_key=api_key)

    def ask(
        self,
        prompt: str,
        system: str = "",
        max_tokens: int = 1024,
    ) -> tuple[str, dict]:
        messages = [{"role": "user", "content": prompt}]
        kwargs = {"model": self.MODEL, "max_tokens": max_tokens, "messages": messages}
        if system:
            kwargs["system"] = system

        response = self._client.messages.create(**kwargs)
        text = response.content[0].text if response.content else ""
        usage = {
            "input_tokens": response.usage.input_tokens,
            "output_tokens": response.usage.output_tokens,
            "cost_usd": self._estimate_cost(response.usage.input_tokens, response.usage.output_tokens),
        }
        return text, usage

    def stream_ask(self, prompt: str, system: str = "", max_tokens: int = 1024):
        messages = [{"role": "user", "content": prompt}]
        kwargs = {"model": self.MODEL, "max_tokens": max_tokens, "messages": messages}
        if system:
            kwargs["system"] = system

        with self._client.messages.stream(**kwargs) as stream:
            for text in stream.text_stream:
                yield text

    def _estimate_cost(self, input_tokens: int, output_tokens: int) -> float:
        return (input_tokens / 1_000_000 * INPUT_COST_PER_1M) + (output_tokens / 1_000_000 * OUTPUT_COST_PER_1M)

    @staticmethod
    def estimate_tokens(text: str) -> int:
        return max(1, len(text) // 4)