"""Core Artip class for managing temporal period identifiers."""

from datetime import date, datetime

import polars as pl

from .exceptions import (
    InvalidGranularityError,
    InvalidIntervalError,
    PeriodTooLargeError,
)
from .period_id import PeriodId

TimeType = date | datetime

VALID_GRANULARITIES = {"year", "month", "day", "hour", "minute", "second"}


class Artip:
    """
    Manager for generating and interpreting temporal period identifiers in Power BI.

    This class encodes temporal periods as scalar integers optimized for Vertipaq
    storage, enabling efficient filtering in Power BI reports. Each Artip instance
    maintains a calendar of all generated periods for export to Power BI.

    Attributes:
        min_time: Minimum time boundary (inclusive).
        max_time: Maximum time boundary (inclusive).
        granularity: Temporal granularity ('year', 'month', 'day', 'hour', 'minute', 'second').
        period_max_size: Maximum allowed period size in granularity units (None = unlimited).

    Examples:
        >>> from datetime import date
        >>> artip = Artip(
        ...     min_time=date(2024, 1, 1),
        ...     max_time=date(2024, 12, 31),
        ...     granularity='day'
        ... )
        >>> period_id = artip.make_period_id(date(2024, 1, 1), date(2024, 1, 31))
        >>> calendar = artip.make_calendar_table()
    """

    _MICROS_PER_DAY = 86_400_000_000
    _INTERVAL_BY_GRANULARITY = {
        "year": "1y",
        "month": "1mo",
        "day": "1d",
        "hour": "1h",
        "minute": "1m",
        "second": "1s",
    }

    _TRUNCATE_BY_GRANULARITY = {
        "year":   lambda dt: dt.replace(month=1, day=1, hour=0, minute=0, second=0, microsecond=0),
        "month":  lambda dt: dt.replace(day=1, hour=0, minute=0, second=0, microsecond=0),
        "day":    lambda dt: dt.replace(hour=0, minute=0, second=0, microsecond=0),
        "hour":   lambda dt: dt.replace(minute=0, second=0, microsecond=0),
        "minute": lambda dt: dt.replace(second=0, microsecond=0),
        "second": lambda dt: dt.replace(microsecond=0),
    }

    _COUNT_BY_GRANULARITY = {
        "second": lambda s, e: int((e - s).total_seconds()) + 1,
        "minute": lambda s, e: int((e - s).total_seconds() / 60) + 1,
        "hour":   lambda s, e: int((e - s).total_seconds() / 3600) + 1,
        "day":    lambda s, e: (e - s).days + 1,
        "month":  lambda s, e: (e.year - s.year) * 12 + (e.month - s.month) + 1,
        "year":   lambda s, e: e.year - s.year + 1,
    }

    def __init__(
        self,
        min_time: TimeType,
        max_time: TimeType,
        granularity: str = "day",
        period_max_size: int | None = None,
    ) -> None:
        """
        Initialize an Artip instance.

        Args:
            min_time: Minimum time boundary.
            max_time: Maximum time boundary.
            granularity: Temporal granularity. Defaults to 'day'.
            period_max_size: Maximum period size in units. None for unlimited.

        Raises:
            InvalidIntervalError: If min_time >= max_time or times are invalid.
            InvalidGranularityError: If granularity is not supported.
            ValueError: If period_max_size is not None and not a positive integer.

        Examples:
            >>> from datetime import date
            >>> artip = Artip(date(2024, 1, 1), date(2024, 12, 31))
        """
        if granularity not in VALID_GRANULARITIES:
            raise InvalidGranularityError(
                f"Granularity must be one of {VALID_GRANULARITIES}, got {granularity}"
            )

        self._truncate = self._TRUNCATE_BY_GRANULARITY[granularity]
        self._count_units = self._COUNT_BY_GRANULARITY[granularity]

        min_dt = self._truncate(self.to_datetime(min_time))
        max_dt = self._truncate(self.to_datetime(max_time))

        if min_dt >= max_dt:
            raise InvalidIntervalError(
                f"min_time must be less than max_time, got {min_dt} >= {max_dt}"
            )

        if period_max_size is not None:
            if (
                isinstance(period_max_size, bool)
                or not isinstance(period_max_size, int)
                or period_max_size <= 0
            ):
                raise ValueError(
                    f"period_max_size must be a positive integer or None, got {period_max_size!r}"
                )

        self.min_time = min_dt
        self.max_time = max_dt
        self.granularity = granularity
        self.period_max_size = period_max_size

        # Calculate internal bounds
        self.min_id = 1
        total_units = self._count_units(self.min_time, self.max_time)
        self.max_id = total_units

        # Calculate time_id_size (number of digits required for max_id)
        self.time_id_size = self.calculate_digits(self.max_id)

        # Calculate period_size if not provided
        if self.period_max_size is None:
            self.period_size = self.time_id_size
        else:
            self.period_size = self.calculate_digits(self.period_max_size)

        self._multiplier = 10 ** self.period_size
        # Store generated periods in compact form: packed start/end datetimes.
        self._periods: dict[int, tuple[int, int]] = {}  # period_id -> (packed_start, packed_end)
        self._period_ids: dict[int, PeriodId] = {}  # period_value -> cached PeriodId

    @staticmethod
    def _pack_datetime(dt: datetime) -> int:
        """Pack a naive datetime into a single integer for compact in-memory storage."""
        day_ordinal = dt.toordinal()
        micros_of_day = (
            ((dt.hour * 60 + dt.minute) * 60 + dt.second) * 1_000_000 + dt.microsecond
        )
        return day_ordinal * Artip._MICROS_PER_DAY + micros_of_day

    @staticmethod
    def _unpack_datetime(packed: int) -> datetime:
        """Unpack a compact datetime integer back into a datetime object."""
        day_ordinal, micros_of_day = divmod(packed, Artip._MICROS_PER_DAY)
        hour, rem = divmod(micros_of_day, 3_600_000_000)
        minute, rem = divmod(rem, 60_000_000)
        second, microsecond = divmod(rem, 1_000_000)
        base = datetime.fromordinal(day_ordinal)
        return base.replace(
            hour=hour,
            minute=minute,
            second=second,
            microsecond=microsecond,
        )

    @staticmethod
    def to_datetime(time_value: TimeType) -> datetime:
        """Convert date or datetime to datetime."""
        if isinstance(time_value, datetime):
            return time_value
        return datetime.combine(time_value, datetime.min.time())

    @staticmethod
    def calculate_digits(value: int) -> int:
        """Calculate the number of digits needed to represent a value."""
        if value <= 0:
            return 1
        return len(str(value))

    @staticmethod
    def count_time_units(start: datetime, end: datetime, granularity: str) -> int:
        """Count the number of time units between two datetimes."""
        if granularity == "second":
            delta = (end - start).total_seconds()
            return int(delta) + 1
        if granularity == "minute":
            delta = (end - start).total_seconds() / 60
            return int(delta) + 1
        if granularity == "hour":
            delta = (end - start).total_seconds() / 3600
            return int(delta) + 1
        if granularity == "day":
            delta = (end - start).days
            return delta + 1
        if granularity == "month":
            months = (end.year - start.year) * 12 + (end.month - start.month)
            return months + 1
        if granularity == "year":
            years = end.year - start.year
            return years + 1
        raise InvalidGranularityError(f"Unsupported granularity: {granularity}")
    def make_period_id(self, start: TimeType, end: TimeType) -> PeriodId:
        """
        Generate a unique scalar identifier for a temporal period.

        The period must be fully contained within [min_time, max_time] and not
        exceed period_max_size if specified.

        Args:
            start: Period start time (inclusive).
            end: Period end time (inclusive).

        Returns:
            A PeriodId encoding the period.

        Raises:
            InvalidIntervalError: If start > end or period is outside bounds.
            PeriodTooLargeError: If period exceeds period_max_size.

        Examples:
            >>> from datetime import date
            >>> artip = Artip(date(2024, 1, 1), date(2024, 12, 31))
            >>> pid = artip.make_period_id(date(2024, 1, 1), date(2024, 1, 31))
            >>> print(pid)
            101
        """
        start_dt = self._truncate(self.to_datetime(start))
        end_dt = self._truncate(self.to_datetime(end))

        if start_dt > end_dt:
            raise InvalidIntervalError(
                f"start must be <= end, got {start_dt} > {end_dt}"
            )

        # If period is entirely outside bounds, it is not representable: return PeriodId(0)
        if end_dt < self.min_time or start_dt > self.max_time:
            return PeriodId._create(0)

        # Clamp to [min_time, max_time] bounds for ID calculation
        start_clamped = max(start_dt, self.min_time)
        end_clamped = min(end_dt, self.max_time)

        # Calculate start_id and end_id
        if start_dt < self.min_time:
            start_id = 0
        else:
            start_id = self._count_units(self.min_time, start_clamped)

        if end_dt > self.max_time:
            end_id = self.max_id + 1
        else:
            end_id = self._count_units(self.min_time, end_clamped)

        # Calculate period length
        period_length = end_id - start_id

        # Check period size constraint
        if self.period_max_size is not None and period_length > self.period_max_size:
            raise PeriodTooLargeError(
                f"Period length {period_length} exceeds maximum {self.period_max_size}"
            )

        # Construct period_id: start_id * 10^period_size + period_length
        period_value = start_id * self._multiplier + period_length

        # Return cached PeriodId if this period was already generated
        if period_value in self._period_ids:
            return self._period_ids[period_value]
        else:
        # Generate PeriodId and cache it
            period_id = PeriodId._create(period_value)
            self._period_ids[period_value] = period_id
            self._periods[period_value] = (
                self._pack_datetime(start_dt),
                self._pack_datetime(end_dt),
            )
        return period_id

    def get_period_range(self, period_id: PeriodId) -> tuple[datetime, datetime]:
        """
        Retrieve the time range for a generated period identifier.

        Args:
            period_id: A PeriodId previously generated by this instance.

        Returns:
            Tuple of (start_time, end_time) as datetime objects.

        Raises:
            InvalidIntervalError: If the period_id was not generated by this instance.

        Examples:
            >>> from datetime import date
            >>> artip = Artip(date(2024, 1, 1), date(2024, 12, 31))
            >>> pid = artip.make_period_id(date(2024, 1, 1), date(2024, 1, 31))
            >>> start, end = artip.get_period_range(pid)
            >>> (start.day, end.day)
            (1, 31)
        """
        if period_id.value not in self._periods:
            raise InvalidIntervalError(
                f"PeriodId {period_id} was not generated by this Artip instance"
            )

        packed_start, packed_end = self._periods[period_id.value]
        start_dt = self._unpack_datetime(packed_start)
        end_dt = self._unpack_datetime(packed_end)
        return (start_dt, end_dt)

    def make_calendar_table(self) -> pl.LazyFrame:
        """
        Generate a periodic calendar table for Power BI filtering.

        Creates a table with columns: date, period_id, and label. The table contains
        one row per date per period, with labels indicating the date's relationship
        to the period ('starts', 'ends', 'in').

        Returns:
            A Polars LazyFrame with schema:
            - date: datetime or date based on granularity
            - period_id: Int64 (the PeriodId value)
            - label: Categorical ('starts', 'ends', 'in')

        Examples:
            >>> from datetime import date
            >>> artip = Artip(date(2024, 1, 1), date(2024, 1, 10), granularity='day')
            >>> artip.make_period_id(date(2024, 1, 1), date(2024, 1, 5))
            >>> cal = artip.make_calendar_table()
            >>> len(cal)
            7
        """
        if not self._periods:
            # Return empty lazy frame with correct schema
            return pl.DataFrame(
                {
                    "date": pl.Series(dtype=pl.Datetime),
                    "period_id": pl.Series(dtype=pl.Int64),
                    "label": pl.Categorical(),
                },
            ).lazy()

        period_data = [
            {
                "period_id": pid,
                "start_date": self._unpack_datetime(packed_start),
                "end_date": self._unpack_datetime(packed_end),
            }
            for pid, (packed_start, packed_end) in self._periods.items()
        ]
        periods_lf = pl.DataFrame(period_data).lazy()

        # Build all row groups with LazyFrame operations to maximize Polars execution.
        in_lf = periods_lf.select(
            pl.col("period_id"),
            pl.date_ranges(
                pl.col("start_date"),
                pl.col("end_date"),
                interval=self._INTERVAL_BY_GRANULARITY[self.granularity],
            ).alias("date"),
        ).explode("date").with_columns(
            pl.col("date").cast(pl.Datetime),
            pl.col("period_id").cast(pl.Int64),
            label=pl.lit("in", dtype=pl.Categorical),
        ).select(["date", "period_id", "label"])

        starts_lf = periods_lf.select(
            pl.col("start_date").alias("date"),
            pl.col("period_id").cast(pl.Int64),
            label=pl.lit("starts", dtype=pl.Categorical),
        ).with_columns(pl.col("date").cast(pl.Datetime)).select(["date", "period_id", "label"])

        ends_lf = periods_lf.select(
            pl.col("end_date").alias("date"),
            pl.col("period_id").cast(pl.Int64),
            label=pl.lit("ends", dtype=pl.Categorical),
        ).with_columns(pl.col("date").cast(pl.Datetime)).select(["date", "period_id", "label"])

        calendar_lf = pl.concat([in_lf, starts_lf, ends_lf]).sort(["date", "period_id"])
        return calendar_lf
