from __future__ import annotations

import dataclasses as dc
import datetime as dt
import operator
from functools import total_ordering
from typing import Any
from typing import Generic
from typing import Literal
from typing import NoReturn
from typing import TypeVar
from typing import Union

import numpy as np
import pandas as pd
import pyarrow as pa
from typing_extensions import TypeGuard

from .misc import is_int
from .misc import is_num


TimeValue = Union[int, float, dt.date, dt.datetime, np.datetime64, pd.Timestamp]
TimeDelta = Union[int, float, dt.timedelta, np.timedelta64, pd.Timedelta]

TimeValueType = TypeVar(
    'TimeValueType', int, float, dt.date, dt.datetime, np.datetime64, pd.Timestamp,
)
TimeDeltaType = TypeVar(
    'TimeDeltaType', int, float, dt.timedelta, np.timedelta64, pd.Timedelta,
)


@total_ordering
@dc.dataclass
class Length(Generic[TimeDeltaType]):
    value: TimeDeltaType
    in_rows: bool = False

    def __post_init__(self) -> None:
        # When deserializing on the server we need to convert str values to timedelta's
        if isinstance(self.value, str):
            self.value = pd.Timedelta(self.value)
        self.validate()

    def __str__(self) -> str:
        desc = str(self.value).replace(' ', '_')
        if self.in_rows:
            return f'{desc}_rows'
        else:
            return desc

    def __lt__(self, other: Any) -> bool | NoReturn:
        # total_ordering decorator will auto-generate all 6 comparison operators
        # https://docs.python.org/3.3/library/functools.html#functools.total_ordering
        if isinstance(other, Length):
            self._check_compatible_types(other)
            return self.value < other.value
        else:
            return self.value < other

    def __add__(self, other: Any) -> Length | NoReturn:
        if isinstance(other, Length):
            self._check_compatible_types(other)
            return Length(value=self.value + other.value, in_rows=self.in_rows)
        else:
            return Length(value=self.value + other, in_rows=self.in_rows)

    def __sub__(self, other: Any) -> Length | NoReturn:
        if isinstance(other, Length):
            self._check_compatible_types(other)
            return Length(value=self.value - other.value, in_rows=self.in_rows)
        else:
            return Length(value=self.value - other, in_rows=self.in_rows)

    def _check_compatible_types(self, other: Any) -> None | NoReturn:
        if not isinstance(other, Length):
            raise ValueError('Both objects must be of type Length')
        else:
            if is_num(self.value) != is_num(other.value):
                raise ValueError(
                    f'Incompatible types for "value" attribute: '
                    f'{type(self.value)} and {type(other.value)}',
                )
            elif not isinstance(other.value, type(self.value)):
                raise ValueError(
                    f'Incompatible types for "value" attribute: '
                    f'{type(self.value)} and {type(other.value)}',
                )
            elif self.in_rows != other.in_rows:
                raise ValueError('Incompatible "in_rows" attributes')

    @property
    def is_null(self) -> bool:
        # Checks for NaN / None / NaT
        return pd.isna(self.value)

    @property
    def is_inf(self) -> bool:
        return is_num(self.value) and np.isinf(self.value)

    def validate(self, time_type: pa.DataType | None = None) -> None | NoReturn:
        if self.in_rows:
            if not is_int(self.value):
                raise ValueError('Value should be an int if "in_rows=True"')
        elif time_type is not None:
            compatible_types = True
            if isinstance(time_type, pa.TimestampType) or is_datetype(time_type):
                if not isinstance(
                    self.value, (dt.timedelta, np.timedelta64, pd.Timedelta),
                ):
                    compatible_types = False
            else:
                if not is_num(self.value):
                    compatible_types = False

            if not compatible_types:
                msg = (
                    f'Value of type {type(self.value)} is incompatible '
                    f'with time id column values of type {time_type}'
                )
                if is_int(self.value):
                    msg += '\nDid you forget to pass "in_rows=True"?'
                raise ValueError(msg)

        if not self.is_null:
            if is_num(self.value):
                if self.value <= 0:
                    raise ValueError('Length value should be positive')
            else:
                if self.value <= dt.timedelta(0):  # type: ignore[operator]
                    raise ValueError('Length value should be positive')

    def to_sql(self) -> str:
        # For timedeltas we need to convert to a SQL interval string,
        # for all other types we just convert the numeric value to str
        # NB: As of DuckDB 0.9.2 interval time conversion only works w/ integer values:
        # https://github.com/duckdb/duckdb/issues/9762
        # FIXME: remove int() below when #9773 is in the next DuckDB release
        # https://github.com/duckdb/duckdb/pull/9773
        if isinstance(self.value, (dt.timedelta, pd.Timedelta)):
            length_in_secs = self.value.total_seconds()
            sql = f"'{int(length_in_secs)} seconds'::interval"
        elif isinstance(self.value, np.timedelta64):
            length_in_secs = self.value / np.timedelta64(1, 's')
            sql = f"'{int(length_in_secs)} seconds'::interval"
        else:
            sql = str(self.value)

        return sql


def is_datetime(x: Any) -> TypeGuard[
    dt.date | dt.datetime | np.datetime64
    | pd.Timestamp | pa.TimestampType | pa.Date32Scalar | pa.Date64Scalar
]:
    return (
        isinstance(x, (dt.datetime, np.datetime64, pd.Timestamp, pa.TimestampType))
        or is_datetype(type(x))
    )


def is_timedelta(x: Any) -> TypeGuard[
    dt.timedelta | np.timedelta64 | pd.Timedelta | pa.DurationType
]:
    return isinstance(x, (dt.timedelta, np.timedelta64, pd.Timedelta, pa.DurationType))


def is_datetype(x: Any) -> TypeGuard[pa.Date32Scalar | pa.Date64Scalar]:
    return x == dt.date or x == pa.date32() or x == pa.date64()


def shift_time_value(
    time_value: TimeValueType, time_shift: TimeDeltaType,
    direction: Literal['left', 'right'],
) -> TimeValueType | NoReturn:
    if direction == 'left':
        op = operator.sub
    elif direction == 'right':
        op = operator.add
    else:
        raise ValueError(f'Invalid direction: "{direction}"')

    if (
        # Check for int/float compatibility
        (is_num(time_value) and is_num(time_shift))
        # Check for python datetime - timedelta compatibility
        or (isinstance(time_value, dt.datetime) and isinstance(time_shift, dt.timedelta))  # noqa: E501
        # Check for numpy datetime64 - timedelta64 compatibility
        or (isinstance(time_value, np.datetime64) and isinstance(time_shift, np.timedelta64))  # noqa: E501
        # Check for pandas Timestamp - Timedelta compatibility
        or (isinstance(time_value, pd.Timestamp) and isinstance(time_shift, pd.Timedelta))  # noqa: E501
    ):
        return op(time_value, time_shift)
    else:
        raise TypeError(
            f'time value of type {type(time_value)} is incompatible '
            f'with a time shift of type {type(time_shift)}',
        )
