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


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 bs.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