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


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