import asyncio
import random
from functools import wraps
from typing import Any, Awaitable, Callable, TypeVar, cast

from requests import HTTPError, RequestException
from rich.console import Console
from rich.panel import Panel

from shadai.core.exceptions import (
    AgentConfigurationError,
    AgentExecutionError,
    AgentFunctionError,
    ConfigurationError,
    IngestionError,
    IntelligenceAPIError,
)

console = Console()

# Define TypeVars for the function return types
T = TypeVar("T")
R = TypeVar("R")


def retry_on_server_error(max_retries: int = 5, base_delay: float = 1.0) -> Callable:
    """Decorator for retrying requests on server errors."""

    def decorator(func: Callable) -> Callable:
        @wraps(func)
        async def wrapper(*args: Any, **kwargs: Any) -> Any:
            for attempt in range(max_retries):
                try:
                    return await func(*args, **kwargs)
                except RequestException as e:
                    if e.response.status_code == 402:
                        console.print(
                            "[red]⚠️  Insufficient balance. Please top up your account.[/]"
                        )
                        raise
                    else:
                        if (
                            isinstance(e, HTTPError)
                            and 400 <= e.response.status_code < 500
                        ):
                            raise
                        if attempt == max_retries - 1:
                            raise IntelligenceAPIError("Max retries reached") from e
                        wait_time = min(base_delay * (2**attempt), 8) * (
                            1 + random.random() * 0.1
                        )
                        console.print(
                            f"[yellow]⚠️  Retrying request ({attempt + 1}/{max_retries}) in {wait_time:.1f}s...[/]"
                        )
                        await asyncio.sleep(wait_time)

        return wrapper

    return decorator


def handle_errors(func: Callable[..., Awaitable[T]]) -> Callable[..., Awaitable[T]]:
    """
    Decorator for handling session errors with rich output.

    This decorator preserves the return type of the decorated function.

    Args:
        func: The async function to decorate

    Returns:
        A decorated function with the same return type as the original
    """

    ERROR_MESSAGES = {
        IntelligenceAPIError: {
            "title": "API Error",
            "suggestions": [
                "Check your API call parameters",
                "Check your API key",
                "Verify your available balance",
                "Ensure you have network connectivity",
            ],
        },
        ConfigurationError: {
            "title": "Configuration Error",
            "suggestions": [
                "Set INTELLIGENCE_API_KEY environment variable",
                "Check your configuration settings",
            ],
        },
        IngestionError: {
            "title": "Ingestion Error",
            "suggestions": [
                "Verify your input files exist",
                "Check file permissions",
                "Ensure files are in supported format",
            ],
        },
        AgentConfigurationError: {
            "title": "Agent Configuration Error",
            "suggestions": [
                "Check your agent configuration",
            ],
        },
        AgentExecutionError: {
            "title": "Agent Execution Error",
            "suggestions": [
                "Check your agent configuration",
            ],
        },
        AgentFunctionError: {
            "title": "Agent Function Error",
            "suggestions": [
                "Check your agent configuration",
            ],
        },
    }

    async def cleanup_session(instance: Any) -> None:
        """Helper function to cleanup session if needed."""
        if (
            instance
            and hasattr(instance, "_delete")
            and instance._delete
            and func.__name__ != "__aexit__"
            and not getattr(instance, "_session_cleaned_up", False)
        ):
            try:
                console.print("[yellow]⚙️  Cleaning up session...[/]")
                await instance.delete()
                setattr(instance, "_session_cleaned_up", True)
                console.print("[green]✓[/] Session cleaned up")
            except Exception as cleanup_error:
                console.print(
                    f"[red]✗[/] Failed to cleanup session: {str(cleanup_error)}"
                )

    @wraps(func)
    async def wrapper(*args: Any, **kwargs: Any) -> T:
        try:
            return await func(*args, **kwargs)

        except KeyboardInterrupt:
            console.print("\n[yellow]⚠️  Process interrupted by user[/]")
            if args:
                await cleanup_session(args[0])
            raise SystemExit(130)

        except tuple(ERROR_MESSAGES.keys()) as e:
            error_config = ERROR_MESSAGES[type(e)]
            console.print(
                Panel(
                    f"[bold red]{error_config['title']}[/]\n\n"
                    f"[yellow]Details:[/] {str(e)}\n\n"
                    f"[blue]Suggestions:[/]\n"
                    + "".join(f"• {s}\n" for s in error_config["suggestions"]),
                    title="[red]Error",
                    border_style="red",
                )
            )
            if args:
                await cleanup_session(args[0])
            raise SystemExit(1)

        except Exception as e:
            console.print(
                Panel(
                    f"[bold red]Unexpected Error[/]\n\n"
                    f"[yellow]Error Type:[/] {type(e).__name__}\n"
                    f"[yellow]Details:[/] {str(e)}",
                    title="[red]System Error",
                    border_style="red",
                )
            )
            if args:
                await cleanup_session(args[0])
            raise SystemExit(1)

    return cast(Callable[..., Awaitable[T]], wrapper)
