import asyncio
import time
from collections.abc import Awaitable, Callable
from datetime import UTC, datetime
from typing import ParamSpec, TypeVar

from boulderopalscaleupsdk.common.dtypes import Duration
from boulderopalscaleupsdk.device.common import Component
from boulderopalscaleupsdk.device.processor.common import ComponentParameter
from pydantic import TypeAdapter
from pydantic.dataclasses import dataclass
from rich.columns import Columns
from rich.console import Console
from rich.table import Table

_T = TypeVar("_T")
_P = ParamSpec("_P")


def _split_table_sizes(length: int) -> list[int]:
    # We split the parameters into sets of up to preferred_width columns,
    # each one will be displayed in a separate table.
    # (If there'd be a spare column, we add it to the first table.)
    preferred_width = 4
    if length <= preferred_width + 1:
        return [length]
    set_sizes = [preferred_width] * (length // preferred_width)
    remainder = length - sum(set_sizes)
    if remainder == 1:
        set_sizes[0] += 1
    elif remainder != 0:
        set_sizes.append(remainder)
    return set_sizes


@dataclass
class _DisplayInfo:
    label: str
    unit: str
    scale: float

    def _get_value(self, value: float | Duration | None) -> float | None:
        if isinstance(value, float | int):
            return value * self.scale
        if isinstance(value, Duration):
            return value.to_seconds() * self.scale
        return None

    @property
    def table_header(self) -> str:
        return f"{self.label}\n({self.unit})"

    def get_table_info(self, parameter: ComponentParameter | None) -> tuple[str, str, str]:
        if parameter is None:
            return "None", "None", "None"
        value = self._get_value(parameter.value)
        em = self._get_value(parameter.err_minus)
        ep = self._get_value(parameter.err_plus)
        std = em + ep if em is not None and ep is not None else None
        return str(value), str(std), str(parameter.calibration_status)


def display_element(reference: str, component: Component) -> None:
    console = Console()

    display_info = {
        k: TypeAdapter(_DisplayInfo).validate_python(extra["display"])
        for k, v in type(component).model_fields.items()
        if isinstance(extra := v.json_schema_extra, dict) and "display" in extra
    }
    headers = list(display_info.keys())

    tables = []
    for length in _split_table_sizes(len(display_info)):
        table = Table()
        table.add_column("", justify="right")
        values = ["value"]
        stds = ["std"]
        calibration_statuses = ["calibration status"]
        for _ in range(length):
            header = headers.pop(0)
            display = display_info[header]
            parameter = getattr(component, header)

            table.add_column(display.table_header, justify="center")
            value, std, calibration_status = display.get_table_info(parameter)
            values.append(value)
            stds.append(std)
            calibration_statuses.append(calibration_status)

        table.add_row(*values)
        table.add_row(*stds)
        table.add_row(*calibration_statuses)
        tables.append(table)

    console.print(Columns(tables, title=f"{type(component).__name__} {reference}"))


def trace(
    fn: Callable[..., _T],
    *fn_args,
    **fn_kwargs,
) -> tuple[_T, float]:
    """Trace the execution time of a function."""

    start_time = datetime.now(UTC)
    result = fn(*fn_args, **fn_kwargs)
    end_time = datetime.now(UTC)
    duration_ms = (end_time - start_time).total_seconds() * 1000
    return result, duration_ms


def retry_with(
    max_retries: int = 3,
    delay_seconds: float = 0.3,
    exponential_factor: float = 2.0,
) -> Callable[
    [Callable[_P, Awaitable[_T]]],
    Callable[_P, Awaitable[_T]],
]:
    def decorator(fn: Callable[_P, Awaitable[_T]]) -> Callable[_P, Awaitable[_T]]:
        async def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _T:
            retries = 0
            delay_base = delay_seconds
            multiplier = exponential_factor
            while True:
                try:
                    return await fn(*args, **kwargs)
                except Exception as e:
                    retries += 1
                    if retries > max_retries:
                        raise RuntimeError(f"Max retries exceeded. Last error: {e}") from e
                    await asyncio.sleep(delay_base * (multiplier ** (retries - 1)))

        return wrapper

    return decorator


def sync_retry_with(
    max_retries: int = 3,
    delay_seconds: float = 0.3,
    exponential_factor: float = 2.0,
) -> Callable[
    [Callable[_P, _T]],
    Callable[_P, _T],
]:
    def decorator(fn: Callable[_P, _T]) -> Callable[_P, _T]:
        def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _T:
            retries = 0
            delay_base = delay_seconds
            multiplier = exponential_factor
            while True:
                try:
                    return fn(*args, **kwargs)
                except Exception as e:
                    retries += 1
                    if retries > max_retries:
                        raise RuntimeError(f"Max retries exceeded. Last error: {e}") from e
                    time.sleep(delay_base * (multiplier ** (retries - 1)))

        return wrapper

    return decorator
