import hashlib
import threading
from collections import defaultdict
from dataclasses import dataclass, field
from typing import Any, Callable, Dict, Optional

import anthropic

from groknroll.clients.base_lm import BaseLM
from groknroll.core.types import ModelUsageSummary, UsageSummary


@dataclass
class ToolUseResult:
    """Result from a tool use completion call."""

    response_text: str = ""
    tool_calls: list[dict[str, Any]] = field(default_factory=list)
    stop_reason: str = ""
    input_tokens: int = 0
    output_tokens: int = 0

    @property
    def has_tool_calls(self) -> bool:
        """Check if the response contains tool calls."""
        return len(self.tool_calls) > 0


class AnthropicClient(BaseLM):
    """
    LM Client for running models with the Anthropic API.

    PERF-003: Implements connection pooling via get_instance() classmethod.
    """

    # PERF-003: Class-level instance cache for connection pooling
    _instances: Dict[str, "AnthropicClient"] = {}
    _lock = threading.Lock()

    @classmethod
    def get_instance(
        cls,
        api_key: str,
        model_name: str | None = None,
        max_tokens: int = 32768,
        **kwargs,
    ) -> "AnthropicClient":
        """
        Get a cached client instance or create a new one.

        PERF-003: Connection pooling - reuses existing SDK instances.

        Args:
            api_key: API key for authentication
            model_name: Model name to use
            max_tokens: Maximum tokens in response
            **kwargs: Additional arguments

        Returns:
            Cached or new AnthropicClient instance
        """
        # Create cache key from configuration
        key_parts = [
            api_key,
            model_name or "",
            str(max_tokens),
        ]
        cache_key = hashlib.md5(":".join(key_parts).encode()).hexdigest()

        with cls._lock:
            if cache_key not in cls._instances:
                cls._instances[cache_key] = cls(
                    api_key=api_key,
                    model_name=model_name,
                    max_tokens=max_tokens,
                    **kwargs
                )
            return cls._instances[cache_key]

    @classmethod
    def clear_instances(cls) -> None:
        """Clear the instance cache. Useful for testing."""
        with cls._lock:
            cls._instances.clear()

    def __init__(
        self,
        api_key: str,
        model_name: str | None = None,
        max_tokens: int = 32768,
        **kwargs,
    ):
        super().__init__(model_name=model_name, **kwargs)
        self.client = anthropic.Anthropic(api_key=api_key)
        self.async_client = anthropic.AsyncAnthropic(api_key=api_key)
        self.model_name = model_name
        self.max_tokens = max_tokens

        # Per-model usage tracking
        self.model_call_counts: dict[str, int] = defaultdict(int)
        self.model_input_tokens: dict[str, int] = defaultdict(int)
        self.model_output_tokens: dict[str, int] = defaultdict(int)
        self.model_total_tokens: dict[str, int] = defaultdict(int)

    def completion(self, prompt: str | list[dict[str, Any]], model: str | None = None) -> str:
        messages, system = self._prepare_messages(prompt)

        model = model or self.model_name
        if not model:
            raise ValueError("Model name is required for Anthropic client.")

        kwargs = {"model": model, "max_tokens": self.max_tokens, "messages": messages}
        if system:
            kwargs["system"] = system

        response = self.client.messages.create(**kwargs)
        self._track_cost(response, model)
        return response.content[0].text

    async def acompletion(
        self, prompt: str | list[dict[str, Any]], model: str | None = None
    ) -> str:
        messages, system = self._prepare_messages(prompt)

        model = model or self.model_name
        if not model:
            raise ValueError("Model name is required for Anthropic client.")

        kwargs = {"model": model, "max_tokens": self.max_tokens, "messages": messages}
        if system:
            kwargs["system"] = system

        response = await self.async_client.messages.create(**kwargs)
        self._track_cost(response, model)
        return response.content[0].text

    def _prepare_messages(
        self, prompt: str | list[dict[str, Any]]
    ) -> tuple[list[dict[str, Any]], str | None]:
        """Prepare messages and extract system prompt for Anthropic API."""
        system = None

        if isinstance(prompt, str):
            messages = [{"role": "user", "content": prompt}]
        elif isinstance(prompt, list) and all(isinstance(item, dict) for item in prompt):
            # Extract system message if present (Anthropic handles system separately)
            messages = []
            for msg in prompt:
                if msg.get("role") == "system":
                    system = msg.get("content")
                else:
                    messages.append(msg)
        else:
            raise ValueError(f"Invalid prompt type: {type(prompt)}")

        return messages, system

    def _track_cost(self, response: anthropic.types.Message, model: str):
        self.model_call_counts[model] += 1
        self.model_input_tokens[model] += response.usage.input_tokens
        self.model_output_tokens[model] += response.usage.output_tokens
        self.model_total_tokens[model] += response.usage.input_tokens + response.usage.output_tokens

        # Track last call for handler to read
        self.last_prompt_tokens = response.usage.input_tokens
        self.last_completion_tokens = response.usage.output_tokens

    def get_usage_summary(self) -> UsageSummary:
        model_summaries = {}
        for model in self.model_call_counts:
            model_summaries[model] = ModelUsageSummary(
                total_calls=self.model_call_counts[model],
                total_input_tokens=self.model_input_tokens[model],
                total_output_tokens=self.model_output_tokens[model],
            )
        return UsageSummary(model_usage_summaries=model_summaries)

    def get_last_usage(self) -> ModelUsageSummary:
        return ModelUsageSummary(
            total_calls=1,
            total_input_tokens=self.last_prompt_tokens,
            total_output_tokens=self.last_completion_tokens,
        )

    def completion_stream(
        self,
        prompt: str | list[dict[str, Any]],
        model: str | None = None,
        chunk_size: int = 1,  # OPT-004: Smaller default chunk size for faster first token
    ):
        """
        Stream completion response token by token.

        OPT-004: Optimized streaming with:
        - Reduced buffer sizes for earlier first token
        - Immediate yielding of chunks
        - Backpressure handling via generator pattern

        Args:
            prompt: Text prompt or message list
            model: Model name (optional, uses default)
            chunk_size: Minimum characters to buffer before yielding (default: 1)

        Yields:
            str: Response chunks as they arrive

        Example:
            for chunk in client.completion_stream("Hello"):
                print(chunk, end="", flush=True)
        """
        messages, system = self._prepare_messages(prompt)

        model = model or self.model_name
        if not model:
            raise ValueError("Model name is required for Anthropic client.")

        kwargs = {"model": model, "max_tokens": self.max_tokens, "messages": messages}
        if system:
            kwargs["system"] = system

        # OPT-004: Stream response with message streaming
        buffer = ""
        with self.client.messages.stream(**kwargs) as stream:
            for text in stream.text_stream:
                buffer += text

                # OPT-004: Yield immediately when buffer exceeds chunk_size
                if len(buffer) >= chunk_size:
                    yield buffer
                    buffer = ""

        # Yield any remaining content
        if buffer:
            yield buffer

        # Track usage
        self.model_call_counts[model] += 1

    async def acompletion_stream(
        self,
        prompt: str | list[dict[str, Any]],
        model: str | None = None,
        chunk_size: int = 1,  # OPT-004: Smaller default chunk size
    ):
        """
        Async stream completion response token by token.

        OPT-004: Optimized async streaming with:
        - Reduced buffer sizes for earlier first token
        - Immediate yielding of chunks
        - Backpressure handling via async generator

        Args:
            prompt: Text prompt or message list
            model: Model name (optional, uses default)
            chunk_size: Minimum characters to buffer before yielding (default: 1)

        Yields:
            str: Response chunks as they arrive

        Example:
            async for chunk in client.acompletion_stream("Hello"):
                print(chunk, end="", flush=True)
        """
        messages, system = self._prepare_messages(prompt)

        model = model or self.model_name
        if not model:
            raise ValueError("Model name is required for Anthropic client.")

        kwargs = {"model": model, "max_tokens": self.max_tokens, "messages": messages}
        if system:
            kwargs["system"] = system

        # OPT-004: Async stream response
        buffer = ""
        async with self.async_client.messages.stream(**kwargs) as stream:
            async for text in stream.text_stream:
                buffer += text

                # OPT-004: Yield immediately when buffer exceeds chunk_size
                if len(buffer) >= chunk_size:
                    yield buffer
                    buffer = ""

        # Yield any remaining content
        if buffer:
            yield buffer

        # Track usage
        self.model_call_counts[model] += 1

    def completion_with_tools(
        self,
        prompt: str | list[dict[str, Any]],
        tools: list[dict[str, Any]],
        model: str | None = None,
        tool_choice: Optional[dict[str, Any]] = None,
    ) -> ToolUseResult:
        """
        Completion with native Anthropic tool use API.

        This method enables Claude to use tools/functions during the conversation.
        Tools are defined using JSON Schema format.

        Args:
            prompt: Text prompt or message list
            tools: List of tool definitions in Anthropic format:
                   [{"name": "...", "description": "...", "input_schema": {...}}, ...]
            model: Model name (optional, uses default)
            tool_choice: Optional tool choice constraint:
                        {"type": "auto"} - Claude decides (default)
                        {"type": "any"} - Claude must use a tool
                        {"type": "tool", "name": "tool_name"} - Use specific tool

        Returns:
            ToolUseResult with response text, tool calls, and usage info

        Example:
            tools = [
                {
                    "name": "read",
                    "description": "Read file contents",
                    "input_schema": {
                        "type": "object",
                        "properties": {
                            "path": {"type": "string", "description": "File path"}
                        },
                        "required": ["path"]
                    }
                }
            ]
            result = client.completion_with_tools(
                prompt="Read the README.md file",
                tools=tools
            )
            if result.has_tool_calls:
                for call in result.tool_calls:
                    print(f"Tool: {call['name']}, Input: {call['input']}")
        """
        messages, system = self._prepare_messages(prompt)

        model = model or self.model_name
        if not model:
            raise ValueError("Model name is required for Anthropic client.")

        kwargs: dict[str, Any] = {
            "model": model,
            "max_tokens": self.max_tokens,
            "messages": messages,
            "tools": tools,
        }
        if system:
            kwargs["system"] = system
        if tool_choice:
            kwargs["tool_choice"] = tool_choice

        response = self.client.messages.create(**kwargs)
        self._track_cost(response, model)

        # Parse response content
        result = ToolUseResult(
            stop_reason=response.stop_reason or "",
            input_tokens=response.usage.input_tokens,
            output_tokens=response.usage.output_tokens,
        )

        for block in response.content:
            if block.type == "text":
                result.response_text += block.text
            elif block.type == "tool_use":
                result.tool_calls.append(
                    {
                        "id": block.id,
                        "name": block.name,
                        "input": block.input,
                    }
                )

        return result

    async def acompletion_with_tools(
        self,
        prompt: str | list[dict[str, Any]],
        tools: list[dict[str, Any]],
        model: str | None = None,
        tool_choice: Optional[dict[str, Any]] = None,
    ) -> ToolUseResult:
        """
        Async completion with native Anthropic tool use API.

        Args:
            prompt: Text prompt or message list
            tools: List of tool definitions in Anthropic format
            model: Model name (optional, uses default)
            tool_choice: Optional tool choice constraint

        Returns:
            ToolUseResult with response text, tool calls, and usage info
        """
        messages, system = self._prepare_messages(prompt)

        model = model or self.model_name
        if not model:
            raise ValueError("Model name is required for Anthropic client.")

        kwargs: dict[str, Any] = {
            "model": model,
            "max_tokens": self.max_tokens,
            "messages": messages,
            "tools": tools,
        }
        if system:
            kwargs["system"] = system
        if tool_choice:
            kwargs["tool_choice"] = tool_choice

        response = await self.async_client.messages.create(**kwargs)
        self._track_cost(response, model)

        # Parse response content
        result = ToolUseResult(
            stop_reason=response.stop_reason or "",
            input_tokens=response.usage.input_tokens,
            output_tokens=response.usage.output_tokens,
        )

        for block in response.content:
            if block.type == "text":
                result.response_text += block.text
            elif block.type == "tool_use":
                result.tool_calls.append(
                    {
                        "id": block.id,
                        "name": block.name,
                        "input": block.input,
                    }
                )

        return result

    def submit_tool_results(
        self,
        messages: list[dict[str, Any]],
        tool_results: list[dict[str, Any]],
        tools: list[dict[str, Any]],
        model: str | None = None,
    ) -> ToolUseResult:
        """
        Submit tool results back to continue the conversation.

        After executing tools, use this method to send results back to Claude.

        Args:
            messages: Conversation history including the assistant's tool use message
            tool_results: List of tool results:
                         [{"tool_use_id": "...", "content": "..."}, ...]
            tools: List of available tools (same as original call)
            model: Model name (optional, uses default)

        Returns:
            ToolUseResult with Claude's next response

        Example:
            # After getting tool calls from completion_with_tools
            tool_results = []
            for call in result.tool_calls:
                output = await execute_tool(call['name'], call['input'])
                tool_results.append({
                    "tool_use_id": call['id'],
                    "content": str(output)
                })

            # Continue conversation
            next_result = client.submit_tool_results(
                messages=[...],  # Include previous messages
                tool_results=tool_results,
                tools=tools
            )
        """
        model = model or self.model_name
        if not model:
            raise ValueError("Model name is required for Anthropic client.")

        # Build tool result message
        tool_result_content = []
        for tr in tool_results:
            tool_result_content.append(
                {
                    "type": "tool_result",
                    "tool_use_id": tr["tool_use_id"],
                    "content": tr["content"],
                }
            )

        # Add tool results as user message
        conversation = messages + [{"role": "user", "content": tool_result_content}]

        # Extract system message if present
        system = None
        filtered_messages = []
        for msg in conversation:
            if msg.get("role") == "system":
                system = msg.get("content")
            else:
                filtered_messages.append(msg)

        kwargs: dict[str, Any] = {
            "model": model,
            "max_tokens": self.max_tokens,
            "messages": filtered_messages,
            "tools": tools,
        }
        if system:
            kwargs["system"] = system

        response = self.client.messages.create(**kwargs)
        self._track_cost(response, model)

        # Parse response
        result = ToolUseResult(
            stop_reason=response.stop_reason or "",
            input_tokens=response.usage.input_tokens,
            output_tokens=response.usage.output_tokens,
        )

        for block in response.content:
            if block.type == "text":
                result.response_text += block.text
            elif block.type == "tool_use":
                result.tool_calls.append(
                    {
                        "id": block.id,
                        "name": block.name,
                        "input": block.input,
                    }
                )

        return result
