from typing import Optional, Type
import anthropic
from pydantic import BaseModel
import bilbystats as bs


async def claude_api_async(
    content: str,
    instructions: str = "",
    model_name: str = "claude-3-7-sonnet-20250219",
    response_format: Optional[Type[BaseModel]] = None,
    temperature: float = 1.0,
    max_tokens: int = None,
) -> str:
    """
    Async call to the Claude API to generate a response based on user input and system instructions.

    This function initializes the Claude API client, constructs a message payload with the given
    instructions and content, and retrieves the model's response text asynchronously.

    ---------------------------------------------------------------------------
    ARGUMENTS:
    instructions : str
        System-level instructions guiding the model's behavior and response style.
    content : str
        User-provided content or prompt for which a response is requested.
    model_name : str, optional
        The name of the Claude model to be used for generating the response.
        Defaults to "claude-3-7-sonnet-20250219".
    temperature : float, optional
        Temperature parameter for controlling randomness in generation.
        Defaults to 1.0.
    response_format : Optional[Type[BaseModel]], optional
        Pydantic model class to enforce structured JSON output.
    ---------------------------------------------------------------------------
    OUTPUT:
    output : str
        The text content of the model's response generated by Claude.

    ---------------------------------------------------------------------------
    AUTHORS: Samuel Davenport and Stephen Enwright-Ward
    ---------------------------------------------------------------------------
    """
    if model_name in ("claude"):
        model_name = "claude-3-7-sonnet-20250219"

    key = bs.read_api_key("claude")
    client = anthropic.AsyncAnthropic(api_key=key)

    api_call_args = {
        "model": model_name,
        "max_tokens": max_tokens,
        "temperature": temperature,
        "system": instructions,
    }

    if response_format:
        tool_schema = bs._pydantic_to_tool_schema(response_format)
        api_call_args["tools"] = [tool_schema]
        api_call_args["tool_choice"] = {
            "type": "tool",
            "name": "structured_response",
        }
        content = f"{content}\n\nPlease provide your response using the structured format tool."

    messages = [
        {
            "role": "user",
            "content": [
                {
                    "type": "text",
                    "text": content,
                }
            ],
        }
    ]

    api_call_args["messages"] = messages

    message = await client.messages.create(
        **api_call_args,
    )

    if response_format:
        # Extract structured response
        tool_use = message.content[0]

        if tool_use.type == "tool_use":
            output = response_format(**tool_use.input)

        else:
            raise ValueError("Expected tool use in response")

    else:
        output = message.content[0].text

    await client.close()
    return output