"""
Async functions for performing API calls with synchronous wrappers
"""

import asyncio
import json
import random
from typing import Optional, Type, Union, List
import pydantic 

import anthropic
from google import genai
from openai import AsyncOpenAI
from groq import AsyncGroq

from pydantic import BaseModel
from google.genai.types import GenerateContentConfig

import bilbystats as bs

_openai_client = None
_openrouter_client = None
_deepseek_client = None
_groq_client = None

# ============================================================================
# ASYNC IMPLEMENTATIONS (Core functionality)
# ============================================================================


async def llm_api_async(
    content: str,
    instructions: str = "",
    model_name: str = "gpt-4o",
    response_format: Optional[Type[BaseModel]] = None,
    temperature: float = 1.0,
    max_tokens: int = None,
    mymessages = None,
    timeout = 600,
    web_search = False
) -> Union[str, BaseModel]:
    """
    Async unified interface to call different LLM APIs with optional structured output.

    ---------------------------------------------------------------------------
    ARGUMENTS:
    content : str
        User-provided prompt or content for which a response is requested.
    instructions : str
        System-level instructions to guide the model's behavior and response generation.
    model_name : str, optional
        The identifier for the target model to be used.
        Defaults to "gpt-4o".
    response_format : Optional[Type[BaseModel]], optional
        Pydantic model class to enforce structured JSON output.
        If None, returns regular string response.

    ---------------------------------------------------------------------------
    OUTPUT:
    output : Union[str, BaseModel]
        The response - either a string (regular) or Pydantic model instance (structured).
    ---------------------------------------------------------------------------
    AUTHORS: Samuel Davenport
    ---------------------------------------------------------------------------
    """
    available_models = bs.read_data("api_model_dict.json")
    #available_open_router_models = bs.read_data('openrouter_models.json')
    
    if (
        (model_name in available_models["openai"])
        or model_name in ("deepseek")
        or model_name in available_models["deepseek"]
    ):
        output = await openai_api_async(
            content=content,
            instructions=instructions,
            model_name=model_name,
            response_format=response_format,
            temperature=temperature,
            max_tokens=max_tokens,
            mymessages=mymessages,
            timeout = timeout,
            web_search = web_search
        )

    elif model_name in available_models["claude"] or model_name in ("claude"):
        if not max_tokens:
            max_tokens = 1000
        output = await claude_api_async(
            content=content,
            instructions=instructions,
            model_name=model_name,
            response_format=response_format,
            temperature=temperature,
            max_tokens=max_tokens
        )

    elif model_name in available_models["gemini"] or model_name in ("gemini"):
        output = await gemini_api_async(
            content=content,
            instructions=instructions,
            model_name=model_name,
            response_format=response_format,
            temperature=temperature,
            max_tokens=max_tokens
        )
    elif model_name[:7] == 'ollama/':
            output = await bs.ollama_async(
                content=content,
                instructions=instructions,
                model_name=model_name,
            )
    else:
        output = await openai_api_async(content, instructions, model_name, 
                                        temperature=temperature, max_tokens=max_tokens, mymessages=mymessages, timeout = timeout, response_format=response_format, web_search = web_search)
    
    return output


async def get_openai_client(model_name: str, timeout):
    """Get or create a reusable async OpenAI client"""
    global _openai_client, _deepseek_client, _openrouter_client, _groq_client
    
    available_models = bs.read_data("api_model_dict.json")
    
    if model_name[len(model_name) - 5:] == ':groq':
        _groq_client = AsyncGroq(api_key=bs.read_api_key("groq"), timeout=timeout)
        return _groq_client
        
    if model_name in available_models["openai"]:
        if _openai_client is None:
            _openai_client = AsyncOpenAI(api_key=bs.read_api_key("openai"), timeout=timeout)
        return _openai_client
    
    elif model_name in available_models["deepseek"]:
        if _deepseek_client is None:
            _deepseek_client = AsyncOpenAI(
                api_key=bs.read_api_key("deepseek"),
                base_url="https://api.deepseek.com",
            )
        return _deepseek_client
    
    else:
        if _openrouter_client is None:
            _openrouter_client = AsyncOpenAI(
                base_url="https://openrouter.ai/api/v1",
                api_key=bs.read_api_key('openrouter'),
                timeout = timeout
            )
        return _openrouter_client


async def openai_api_async(
    content: str,
    instructions: str = "",
    model_name: str = "gpt-4o",
    response_format: Optional[Type[BaseModel]] = None,
    temperature: float = 1.0,
    max_tokens: int = None,
    mymessages = None,
    timeout = 5,
    web_search = False
) -> Union[str, BaseModel]:
    
    if model_name == "deepseek":
        model_name = "deepseek-chat"


    # Get reusable client instead of creating new one
    client = await get_openai_client(model_name, timeout = timeout)
    
    if model_name.endswith(':groq'):
        model_name = model_name[:-5] 
    
    available_models = bs.read_data("api_model_dict.json")

    messages = [] 
    if len(instructions) > 0:
        messages.append({"role": "system", "content": instructions})
        
    if mymessages:
        messages = messages + mymessages
        
    messages.append({"role": "user","content": content})

    completion_args = {
        "model": model_name,
        "messages": messages,
        "temperature": temperature,
        "max_tokens": max_tokens,
        "timeout": timeout 
    }

    if response_format:
       completion_args["response_format"] = response_format
        
    if web_search:
       completion_args["tools"] = [{"type": "web_search"}]
    
    if type(response_format) == pydantic._internal._model_construction.ModelMetaclass:
        if model_name in available_models["deepseek"]:
            response_format = {
                "type": "json_object",
            }
        completion_args["response_format"] = response_format

        completion = await client.beta.chat.completions.parse(**completion_args)

        if model_name in available_models["openai"]:
            output = completion.choices[0].message.parsed
        elif model_name in available_models["deepseek"]:
            output = completion.choices[0].message.content
    else:
        completion = await client.chat.completions.create(**completion_args)
        output = completion.choices[0].message.content

    return output


def _pydantic_to_tool_schema(
    model_class: Type[BaseModel],
) -> dict:
    """Convert a Pydantic model to Claude tool schema format."""
    schema = model_class.model_json_schema()

    return {
        "name": "structured_response",
        "description": f"Provide a structured response in {model_class.__name__} format",
        "input_schema": {
            "type": "object",
            "properties": schema.get("properties", {}),
            "required": schema.get("required", []),
        },
    }


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 = _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

async def gemini_api_async(
    content: str,
    instructions: str = "",
    model_name: str = "gemini-2.0-flash",
    response_format: Optional[Type[BaseModel]] = None,
    temperature: float = 1.0,
    max_tokens: int = 1000
) -> str:
    """
    Async generate content using the Gemini API.

    This function sends an async request to the Gemini API to generate content based on
    the provided input. If the default model name "gemini" is used, it is replaced
    with "gemini-2.0-flash" to ensure compatibility. The function returns the
    generated text content.

    ---------------------------------------------------------------------------
    ARGUMENTS:
    - content : str
        The input prompt or content to be sent to the Gemini API for generation.

    - model_name : str, optional (default="gemini-2.0-flash")
        The name of the Gemini model to use for content generation. If "gemini"
        is specified, it defaults to "gemini-2.0-flash". Other options are e.g.
        gemini-2.5-flash-preview-04-17 and gemini-2.0-flash-lite and
        gemini-2.5-pro-preview-05-06. See
        https://ai.google.dev/gemini-api/docs/models for a list of the different
        possible models.

    ---------------------------------------------------------------------------
    OUTPUT:
    - output : str
        The generated text returned by the Gemini API.

    ---------------------------------------------------------------------------
    AUTHORS: Samuel Davenport and Stephen Enwright-Ward
    ---------------------------------------------------------------------------
    """
    if model_name == "gemini":
        model_name = "gemini-2.0-flash"

    # Note: The genai client doesn't have native async support
    # We'll use asyncio.run_in_executor to make it async
    if not response_format:
        def _generate_sync():
            client = genai.Client(api_key=bs.read_api_key("gemini"))
            response = client.models.generate_content(
                model=model_name,
                contents=instructions + content,
                config=GenerateContentConfig(
                    temperature=temperature,
                    max_output_tokens=max_tokens,
                ),
            )
            return response.text
    else:
        def _generate_sync():
            client = genai.Client(
                api_key=bs.read_api_key("gemini"),
                )

            response = client.models.generate_content(
                model=model_name,
                contents=content,
                config={
                    "response_mime_type": "application/json",
                    "system_instruction": instructions,
                    "response_schema": response_format,
                    "temperature": temperature,
                    "max_output_tokens": max_tokens,
                    },
                )

            # return response.text
            return response.parsed
    

    loop = asyncio.get_event_loop()
    output = await loop.run_in_executor(None, _generate_sync)
    return output


async def detect_sentiment_async(
    text: str,
    model_name: str = "gpt-4o",
) -> str:
    """
    Async detect the sentiment of a given text using a specified language model.

    This function uses the OpenAI API to analyze the sentiment of the provided
    text asynchronously. It classifies the sentiment as 'positive', 'neutral', or 'negative',
    and returns only the sentiment label.

    ---------------------------------------------------------------------------
    ARGUMENTS:
    text : str
        The text for which sentiment analysis is to be performed.
    model_name : str, optional (default="gpt-4o")
        The name of the OpenAI model used for sentiment analysis.

    ---------------------------------------------------------------------------
    OUTPUT:
    sentiment : str
        The sentiment classification of the input text ('positive',
        'neutral', or 'negative').

    ---------------------------------------------------------------------------
    AUTHORS: Samuel Davenport and Stephen Enwright-Ward
    ---------------------------------------------------------------------------
    """
    instructions = 'You are an expert in sentiment analysis. Given any input text, determine the overall sentiment and classify it as one of the following categories: positive, neutral, or negative. Respond with only one of these words — "positive", "neutral", or "negative" — and nothing else.'
    sentiment = await llm_api_async(
        text,
        instructions,
        model_name=model_name,
    )

    return sentiment


async def batch_llm_api_async_orig(
    texts: list[str],
    instructions: str = "",
    model_name: str = "gpt-4o",
    prefix: str = "",
    max_concurrent: int = 50,
    batch_size: int = 100,
    max_retries: int = 3,
    base_delay: float = 1.0,
    response_format: Optional[Type[BaseModel]] = None,
    temperature: float = 1.0,
    max_tokens: int = None,
    mymessages = None
) -> list[str]:
    """
    Process multiple texts through LLM API calls with batching, concurrency control, and retry logic.
    ---------------------------------------------------------------------------
    ARGUMENTS:
    texts : list[str]
        list of texts to process.
    instructions : str
        System instructions for the LLM.
    model_name : str, optional
        Name of the model to use.
    max_concurrent : int, optional
        Maximum number of concurrent requests.
    batch_size : int, optional
        Number of requests per batch.
    max_retries : int, optional
        Maximum number of retry attempts for failed requests.
    base_delay : float, optional
        Base delay in seconds for exponential backoff.
    response_format : Optional[Type[BaseModel]], optional
        Pydantic model class to enforce structured JSON output.
        If None, returns regular string response.
    ---------------------------------------------------------------------------
    OUTPUT:
    responses : list[str]
        list of responses from the LLM, in the same order as input texts.
    ---------------------------------------------------------------------------
    AUTHOR: Samuel Davenport
    ---------------------------------------------------------------------------
    """
    semaphore = asyncio.Semaphore(max_concurrent)

    async def process_single_text_with_retry(
        text: str,
        text_index: int,
        response_format: Optional[Type[BaseModel]] = None,
    ) -> str:
        async with semaphore:
            for attempt in range(max_retries + 1):
                try:
                    return await llm_api_async(
                        content=prefix + text,
                        instructions=instructions,
                        model_name=model_name,
                        response_format=response_format,
                        temperature = temperature,
                        max_tokens = max_tokens,
                        mymessages = mymessages
                    )

                except Exception as e:
                    error_str = str(e)
                    print(error_str)

                    # Check if it's a retryable error (503, rate limiting, etc.)
                    is_retryable = any(
                        code in error_str
                        for code in [
                            "503",
                            "429",
                            "502",
                            "504",
                            "UNAVAILABLE",
                            "RATE_LIMIT"
                        ]
                    )
                    print(is_retryable)

                    if is_retryable and attempt < max_retries:
                        # Exponential backoff with jitter
                        delay = base_delay * (2**attempt) + random.uniform(0, 1)
                        print(
                            f"Text {text_index}: Attempt {attempt + 1} failed ({error_str[:50]}...), retrying in {delay:.1f}s"
                        )
                        await asyncio.sleep(delay)

                    else:
                        # Final attempt failed or non-retryable error
                        print(f"Error processing text {text_index}: {e}")
                        return ""  # or raise exception if you prefer

            return ""  # Fallback if all retries exhausted

    bs.loader(1, 1, "Progress")
    responses = []
    total_batches = (len(texts) - 1) // batch_size + 1

    for batch_idx in range(total_batches):
        start_idx = batch_idx * batch_size
        end_idx = min(start_idx + batch_size, len(texts))
        batch_texts = texts[start_idx:end_idx]

        bs.loader(batch_idx, total_batches, "Progress")

        # Create tasks with text indices for better error reporting
        tasks = [
            process_single_text_with_retry(
                text=text,
                text_index=start_idx + i,
                response_format=response_format,
            )
            for i, text in enumerate(batch_texts)
        ]

        # Changed to False since we handle exceptions internally
        batch_responses = await asyncio.gather(
            *tasks,
            return_exceptions=False,
        )

        responses.extend(batch_responses)

        # Add small delay between batches to be nice to the API
        if batch_idx < total_batches - 1:
            await asyncio.sleep(0.1)

    bs.loader(total_batches, total_batches, "Progress")
    return responses

async def batch_llm_api_async_mr(
    texts: List[str], 
    instructions: str = "", 
    model_name: str = "gpt-4o", 
    max_rounds: int = 3, 
    columns=['label'], 
    prefix: str = "",
    max_concurrent: int = 50,
    max_retries: int = 3,
    base_delay: float = 1.0,
    response_format: Optional[Type[BaseModel]] = None,
    temperature: float = 1.0,
    max_tokens: int = None,
    mymessages = None,
    direct_eval = True,
    timeout = 5
) -> List[str]:
    """
    Call a batched asynchronous LLM API and retry missing labels for multiple rounds.

    This function sends a batch of text prompts to an asynchronous LLM API with a
    requirement that the response is in JSON format. It automatically retries failed
    or empty label responses up to `max_rounds` times. The response is converted to
    a DataFrame and updated after each retry round.

    ---------------------------------------------------------------------------
    ARGUMENTS:
    texts : List[str]
        A list of input text strings to be processed by the LLM.
    instructions : str, optional
        Instruction string guiding the LLM's output; must request JSON format.
    model_name : str, optional
        The name of the model to use for inference (default is "gpt-4o").
    max_rounds : int, optional
        Maximum number of retry rounds for handling failed or empty responses.
    columns : list or None, optional
        Optional list of expected JSON keys to structure the output DataFrame.
    prefix : str, optional
        Optional prefix to prepend to each input prompt.
    max_concurrent : int, optional
        Maximum number of concurrent requests in the batch.
    batch_size : int, optional
        Number of texts to send per batch request.
    max_retries : int, optional
        Number of retries per individual batch call in case of failure.
    base_delay : float, optional
        Base delay in seconds between retries for exponential backoff.
    response_format : Optional[Type[BaseModel]], optional
        Pydantic model for validating the structure of each response.

    ---------------------------------------------------------------------------
    OUTPUT:
    df : pd.DataFrame
        A DataFrame containing parsed JSON responses from the LLM, with updated 
        entries for any initially missing or empty labels.
    ---------------------------------------------------------------------------
    AUTHORS: Samuel Davenport
    ---------------------------------------------------------------------------
    """
    if 'json' not in instructions.lower():
        raise Exception("The output needs to be requested to be in JSON format")
    
    responses = await bs.batch_llm_api_async(texts, instructions, model_name, 
                                             prefix = prefix, max_concurrent = max_concurrent, 
                                             max_retries = max_retries, base_delay = base_delay, 
                                             response_format = response_format, 
                                             temperature = temperature, mymessages=mymessages, timeout = timeout)
    df = bs.json2df(responses, columns = columns, direct_eval = direct_eval)

    for i in range(max_rounds - 1):
        empty_labels = df[df[columns[0]].isna() | (df[columns[0]].str.strip() == '')]
        empty_label_index = empty_labels.index
        
        if len(empty_label_index) > 0:
            print('Re-running the ' + str(len(empty_label_index)) + ' texts that failed: Round ' + str(i + 1) + ' of corrections.\n')

            texts2update = [texts[i] for i in empty_label_index]
            
            new_responses = await bs.batch_llm_api_async(texts2update, 
                                                                instructions, 
                                                                model_name, 
                                                                prefix = prefix,
                                                                max_concurrent = max_concurrent,
                                                                max_retries = max_retries,
                                                                base_delay = base_delay,
                                                                response_format = response_format,
                                                                temperature = temperature,
                                                                max_tokens = max_tokens,
                                                                mymessages = mymessages,
                                                                timeout = timeout)
            new_df = bs.json2df(new_responses, columns = columns, direct_eval = direct_eval)
            new_df.index = empty_label_index
            df.update(new_df)
        else:
            return df

    return df


async def batch_llm_api_async(
    texts: list[str],
    instructions: str = "",
    model_name: str = "gpt-4o",
    prefix: str = "",
    max_concurrent: int = 75,
    max_retries: int = 3,
    base_delay: float = 1.0,
    response_format: Optional[Type[BaseModel]] = None,
    temperature: float = 1.0,
    max_tokens: int = None,
    mymessages = None,
    timeout = 10,
    web_search = False
) -> list[str]:
    
    semaphore = asyncio.Semaphore(max_concurrent)
    completed = 0
    total = len(texts)
    
    global _openai_client, _deepseek_client, _openrouter_client, _groq_client
    _openai_client = None
    _deepseek_client = None
    _openrouter_client = None
    _groq_client = None
    
    async def process_single_text_with_retry(
        text: str,
        text_index: int,
        response_format: Optional[Type[BaseModel]] = None,
    ) -> str:
        nonlocal completed
        
        async with semaphore:
            for attempt in range(max_retries + 1):
                try:
                    # Wrap the API call with asyncio.wait_for for client-side timeout control
                    result = await asyncio.wait_for(
                        llm_api_async(
                            content=prefix + text,
                            instructions=instructions,
                            model_name=model_name,
                            response_format=response_format,
                            temperature=temperature,
                            max_tokens=max_tokens,
                            mymessages=mymessages,
                            timeout = timeout, # Still pass to API as well
                            web_search = web_search
                        ),
                        timeout=timeout  # Client-side timeout enforcement
                    )
                    
                    # Update progress
                    completed += 1
                    bs.loader(completed, total, "Progress")
                    
                    return result

                except (Exception, asyncio.CancelledError, asyncio.TimeoutError) as e:
                    error_str = str(e)
                    error_type = type(e).__name__
                    
                    is_retryable = any(
                        code in error_str
                        for code in ["503", "429", "502", "504", "UNAVAILABLE", "RATE_LIMIT", "timeout", "TimeoutError", "CancelledError"]
                    ) or isinstance(e, (asyncio.CancelledError, asyncio.TimeoutError))

                    if is_retryable and attempt < max_retries:
                        delay = base_delay * (2**attempt) + random.uniform(0, 1)
                        #print(f"Text {text_index}: Attempt {attempt + 1} failed ({error_type}), retrying in {delay:.1f}s")
                        await asyncio.sleep(delay)
                    else:
                        completed += 1
                        bs.loader(completed, total, "Progress")
                        #print(f"Text {text_index}: Failed after {attempt + 1} attempts - {error_type}")
                        return ""

            return ""

    # Process all at once
    tasks = [
        process_single_text_with_retry(text, i, response_format)
        for i, text in enumerate(texts)
    ]

    responses = await asyncio.gather(*tasks, return_exceptions=False)
    return responses

async def batch_mr(
    texts: list[str],
    instructions: str = "",
    model_name: str = "gpt-4o",
    prefix: str = "",
    max_concurrent: int = 75,
    max_retries: int = 3,
    base_delay: float = 1.0,
    response_format: Optional[Type[BaseModel]] = None,
    temperature: float = 1.0,
    max_tokens: int = None,
    mymessages = None,
    timeout = 5,
    max_rounds: int = 3
) -> list[str]:
    """
    Run batch_llm_api_async and retry empty string outputs until all outputs are recorded
    or max_rounds is reached.
    """
    responses = await batch_llm_api_async(
        texts=texts,
        instructions=instructions,
        model_name=model_name,
        prefix=prefix,
        max_concurrent=max_concurrent,
        max_retries=max_retries,
        base_delay=base_delay,
        response_format=response_format,
        temperature=temperature,
        max_tokens=max_tokens,
        mymessages=mymessages,
        timeout=timeout
    )
    
    for round_num in range(1, max_rounds):
        # Find indices with empty responses
        empty_indices = [i for i, r in enumerate(responses) if r == ""]
        
        if not empty_indices:
            break
            
        # Retry only the empty responses
        retry_texts = [texts[i] for i in empty_indices]
        retry_responses = await batch_llm_api_async(
            texts=retry_texts,
            instructions=instructions,
            model_name=model_name,
            prefix=prefix,
            max_concurrent=max_concurrent,
            max_retries=max_retries,
            base_delay=base_delay,
            response_format=response_format,
            temperature=temperature,
            max_tokens=max_tokens,
            mymessages=mymessages,
            timeout=timeout
        )
        
        # Update responses with retry results
        for idx, retry_response in zip(empty_indices, retry_responses):
            responses[idx] = retry_response
    
    return responses