import aiohttp
import asyncio
import os
from dotenv import load_dotenv
import logging

logger = logging.getLogger("deepsearch-mcp")

# Load environment variables from .env file
load_dotenv()

# Get API key from environment variables
OPENROUTER_API_KEY = os.getenv('OPENROUTER_API_KEY')

# Add these constants at the top with your other imports
API_URL = "https://openrouter.ai/api/v1/chat/completions"


async def make_api_call(prompt, model="openai/o3-mini", temperature=0.1, session=None, max_retries=5, base_delay=1):
    """Make API call with exponential backoff retry logic."""
    logger.info(f"Making API call to OpenRouter with model: {model}")
    headers = {
        "Authorization": f"Bearer {OPENROUTER_API_KEY}",
        "HTTP-Referer": "http://localhost:8000",
    }
    data = {
        "messages": [{"role": "user", "content": prompt}],
        "model": model,
        "max_tokens": 100000,
        "temperature": temperature
    }

    should_close_session = False
    if session is None:
        session = aiohttp.ClientSession()
        should_close_session = True

    try:
        for attempt in range(max_retries):
            try:
                async with session.post(API_URL, json=data, headers=headers) as response:
                    if response.status == 200:
                        resp_json = await response.json()
                        if 'choices' in resp_json:
                            logger.info(
                                "Successfully received response from OpenRouter")
                            return resp_json['choices'][0]['message']['content']
                    elif response.status == 429:  # Rate limit exceeded
                        delay = base_delay * (2 ** attempt)
                        logger.warning(
                            f"Rate limit exceeded. Retrying in {delay} seconds...")
                        await asyncio.sleep(delay)
                        continue
                    else:
                        error_text = await response.text()
                        logger.error(
                            f"Error response (status {response.status}): {error_text}")

                    delay = base_delay * (2 ** attempt)
                    logger.warning(
                        f"Attempt {attempt+1} failed. Retrying in {delay} seconds...")
                    await asyncio.sleep(delay)

            except Exception as e:
                if attempt == max_retries - 1:
                    logger.error(
                        f"All retry attempts failed. Final error: {str(e)}")
                    raise
                delay = base_delay * (2 ** attempt)
                logger.warning(
                    f"Error occurred: {str(e)}. Retrying in {delay} seconds...")
                await asyncio.sleep(delay)
    finally:
        if should_close_session:
            await session.close()

    raise Exception("Max retries exceeded")
