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

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


async def batch_llm_api_async(
    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 = 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(
                        bs.llm_api_async(
                            content=prefix + text + suffix,
                            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