import time
from stark_qa.tools.api_lib import anthropic_client
from typing import Any, Union


def complete_text_claude(message: Union[str, list], 
                         model: str = "claude-2.1",
                         json_object: bool = False,
                         max_tokens: int = 2048, 
                         temperature: float = 1.0, 
                         max_retry: int = 1,
                         sleep_time: int = 0,
                         tools: list = [],
                         **kwargs: Any) -> str:
    """
    Call the Claude API to complete a prompt.

    Args:
        message (Union[str, list]): The input message or a list of message dicts.
        model (str): The model to use for completion. Default is "claude-2.1".
        json_object (bool): Whether to output in JSON format. Default is False.
        max_tokens (int): Maximum number of tokens to generate. Default is 2048.
        temperature (float): Sampling temperature. Default is 1.0.
        max_retry (int): Maximum number of retries in case of an error. Default is 1.
        sleep_time (int): Sleep time between retries in seconds. Default is 0.
        tools (list): List of tools to use for the completion. Default is an empty list.
        **kwargs (Any): Additional keyword arguments to pass to the API.

    Returns:
        str: The completed text generated by the Claude model.

    Raises:
        Exception: If the completion fails after the maximum number of retries.
    """
    if isinstance(message, str):
        if json_object:
            message = "You are a helpful assistant designed to output in JSON format." + message
        messages = [{"role": "user", "content": message}] 
    else:
        messages = message

    for cnt in range(max_retry):
        try:
            response = anthropic_client.beta.tools.messages.create(
                messages=messages,
                model=model,
                temperature=temperature,
                max_tokens=max_tokens,
                tools=tools,
                **kwargs
            )
            completion = response.to_dict()
            return completion["content"][0]['text']
        except Exception as e:
            print(f"Attempt {cnt} failed: {e}. Retrying after {sleep_time} seconds...")
            time.sleep(sleep_time)
    
    raise Exception("Failed to complete text after maximum retries")
