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


async def batch_mr(
    texts: list[str],
    instructions: str = "",
    model_name: str = "gpt-4o",
    prefix: str = "",
    suffix: 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,
    no_chinese: bool = False
) -> list[str]:
    """
    Run batch_llm_api_async and retry empty string outputs until all outputs are recorded
    or max_rounds is reached.

    ---------------------------------------------------------------------------
    ARGUMENTS:
    texts : list[str]
        List of input texts to process.
    instructions : str, optional (default="")
        System prompt/instructions for the LLM.
    model_name : str, optional (default="gpt-4o")
        Name of the LLM model to use.
    prefix : str, optional (default="")
        Prefix to add to each text.
    max_concurrent : int, optional (default=75)
        Maximum number of concurrent API calls.
    max_retries : int, optional (default=3)
        Maximum retries per individual API call.
    base_delay : float, optional (default=1.0)
        Base delay for exponential backoff.
    response_format : Optional[Type[BaseModel]], optional (default=None)
        Pydantic model for structured output.
    temperature : float, optional (default=1.0)
        Sampling temperature for the LLM.
    max_tokens : int, optional (default=None)
        Maximum tokens in the response.
    mymessages : optional (default=None)
        Custom messages to pass to the API.
    timeout : int, optional (default=5)
        Timeout for each API call in seconds.
    max_rounds : int, optional (default=3)
        Maximum retry rounds for failed responses.
    no_chinese : bool, optional (default=False)
        If True, retry responses that contain Chinese characters.

    ---------------------------------------------------------------------------
    OUTPUT:
    responses : list[str]
        List of LLM responses corresponding to each input text.
    ---------------------------------------------------------------------------
    AUTHORS: Samuel Davenport
    ---------------------------------------------------------------------------
    """

    def needs_retry(response: str) -> bool:
        """Check if a response needs to be retried."""
        if response == "":
            return True
        if no_chinese and bs.contains_chinese(response):
            return True
        return False

    responses = await bs.batch_llm_api_async(
        texts=texts,
        instructions=instructions,
        model_name=model_name,
        prefix=prefix,
        suffix=suffix,
        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 responses that need retry
        retry_indices = [i for i, r in enumerate(responses) if needs_retry(r)]

        if not retry_indices:
            break

        # Retry only the failed responses
        retry_texts = [texts[i] for i in retry_indices]
        retry_responses = await bs.batch_llm_api_async(
            texts=retry_texts,
            instructions=instructions,
            model_name=model_name,
            prefix=prefix,
            suffix=suffix,
            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(retry_indices, retry_responses):
            responses[idx] = retry_response
    
    return responses