import asyncio
import hashlib
import json
import logging
import sqlite3
from datetime import datetime, timedelta
from typing import Any, Dict, Optional, Union, List, Tuple

from inspect_ai.model import Model, GenerateConfig


logger = logging.getLogger(__name__)


async def generate_with_retry(
    model: Model,
    prompt_or_messages: Union[str, List],
    config: GenerateConfig,
    max_retries: int = 20,
    base_delay: float = 2.0,
    desired_schema=None,
) -> Any:
    """
    Generate response with retry logic and optional JSON parsing.

    Args:
        model: The model to use for generation
        prompt_or_messages: Prompt string or list of messages
        config: Generation configuration
        max_retries: Maximum number of retry attempts (default: 20)
        base_delay: Base delay in seconds between retries (default: 2.0)

    Returns:
        The model response or parsed JSON object

    Raises:
        Exception: If all retry attempts fail
    """
    last_exception = None

    for attempt in range(max_retries + 1):
        try:
            # Generate response
            result = await model.generate(
                prompt_or_messages,
                config=config,
            )

            # Try to parse JSON from the completion using .find and .rfind
            completion = result.completion
            json_start = completion.find("{")
            json_end = completion.rfind("}") + 1

            if json_start != -1 and json_end > json_start:
                json_str = completion[json_start:json_end]
                parsed = json.loads(json_str)
                if desired_schema:
                    parsed = desired_schema(**parsed)
                    parsed = parsed.model_dump(mode="json")

                return result, parsed, attempt
            else:
                raise json.JSONDecodeError("No JSON found in completion", completion, 0)

        except Exception as e:
            last_exception = e
            logger.warning(f"Attempt {attempt + 1}/{max_retries + 1} failed: {str(e)}")

            if attempt < max_retries:
                # Wait before retrying with exponential backoff
                delay = base_delay * (1.5**attempt)
                logger.info(f"Waiting {delay:.2f} seconds before retry...")
                await asyncio.sleep(delay)
            else:
                logger.error(
                    f"All {max_retries + 1} attempts failed. Last error: {str(e)}"
                )

    # If we get here, all retries failed
    raise last_exception
