from __future__ import annotations

import datetime as dt
import re
from typing import Callable, Iterable, Any
import pandas as pd

# Define a type for period functions to allow passing 'now'
PeriodFunc = Callable[[dt.datetime], tuple[dt.date, dt.date]]


class DateUtils:
    """
    Period resolution & normalization for ETL artifacts.
    Standardizes on UTC Everywhere + ISO 8601 'Z'.
    """

    # ---- Dynamic registries ----
    _PERIOD_FUNCTIONS: dict[str, PeriodFunc] = {}
    _PERIOD_PATTERNS: list[
        tuple[re.Pattern[str], Callable[[re.Match[str], dt.datetime], dict[str, Any]]]
    ] = []

    _LAST_N_DAYS_RE = re.compile(r"^last_(\d+)_days$")
    _WINDOW_RE = re.compile(r"^(\d{4}-\d{2}-\d{2})\.\.(\d{4}-\d{2}-\d{2})$")

    # ---------------- 1. New Helper for "Z" Formatting ----------------
    @staticmethod
    def to_iso_z(value: dt.datetime) -> str:
        """
        Convert datetime to ISO 8601 string with 'Z' suffix.
        Ensures UTC timezone if naive.
        """
        # Ensure it is timezone aware (assume UTC if missing)
        if value.tzinfo is None:
            value = value.replace(tzinfo=dt.timezone.utc)

        # Convert to UTC to be safe, then format
        return value.astimezone(dt.timezone.utc).isoformat().replace("+00:00", "Z")

    # ---------------- Core coercion helpers ----------------

    @staticmethod
    def _ensure_date(value: str | dt.date | dt.datetime | pd.Timestamp) -> dt.date:
        """Ensure the input is converted to a datetime.date."""
        if isinstance(value, dt.date) and not isinstance(value, dt.datetime):
            return value
        if isinstance(value, dt.datetime):
            return value.date()
        if isinstance(value, pd.Timestamp):
            return value.to_pydatetime().date()
        if isinstance(value, str):
            # Try pandas parser first (robust), then ISO date
            try:
                return pd.to_datetime(value, errors="raise").date()  # type: ignore
            except Exception:
                pass
            try:
                return dt.date.fromisoformat(value)
            except Exception:
                pass
        raise ValueError(f"Unsupported date format: {value!r}")

    # Public alias
    ensure_date = _ensure_date

    @staticmethod
    def _ensure_datetime(
        value: str | dt.date | dt.datetime | pd.Timestamp,
        tz: dt.tzinfo = dt.timezone.utc,
    ) -> dt.datetime:
        """Convert input to timezone-aware datetime (defaults to UTC)."""
        if isinstance(value, dt.datetime):
            return value if value.tzinfo else value.replace(tzinfo=tz)
        if isinstance(value, dt.date):
            return dt.datetime(value.year, value.month, value.day, tzinfo=tz)
        if isinstance(value, pd.Timestamp):
            dtt = value.to_pydatetime()
            return dtt if dtt.tzinfo else dtt.replace(tzinfo=tz)
        if isinstance(value, str):
            ts = pd.to_datetime(value, errors="raise", utc=False)
            dtt = ts.to_pydatetime()
            return dtt if getattr(dtt, "tzinfo", None) else dtt.replace(tzinfo=tz)
        raise ValueError(f"Unsupported datetime format: {value!r}")

    # ---------------- Week / Month / Quarter helpers ----------------

    @classmethod
    def calc_week_range(cls, ref: dt.date) -> tuple[dt.date, dt.date]:
        """Start (Mon) and end (Sun) for the week containing reference_date."""
        start = ref - dt.timedelta(days=ref.weekday())
        end = start + dt.timedelta(days=6)
        return start, end

    @staticmethod
    def get_year_timerange(year: int) -> tuple[dt.date, dt.date]:
        return dt.date(year, 1, 1), dt.date(year, 12, 31)

    @classmethod
    def get_first_day_of_the_quarter(cls, ref: dt.date) -> dt.date:
        quarter = (ref.month - 1) // 3 + 1
        return dt.date(ref.year, 3 * quarter - 2, 1)

    @classmethod
    def get_last_day_of_the_quarter(cls, ref: dt.date) -> dt.date:
        quarter = (ref.month - 1) // 3 + 1
        # Calculate first day of next quarter instead of relying on fixed month addition
        target_month = 3 * quarter + 1
        target_year = ref.year
        if target_month > 12:
            target_month = 1
            target_year += 1
        first_day_next_q = dt.date(target_year, target_month, 1)
        return first_day_next_q - dt.timedelta(days=1)

    @classmethod
    def get_month_range(cls, ref_date: dt.date, n: int = 0) -> tuple[dt.date, dt.date]:
        """
        Range for current month (n=0) or +/- n months relative to ref_date.
        If n == 0, end is ref_date (e.g. today). Otherwise end is calendar month end.
        """
        target_month = (ref_date.month - 1 + n) % 12 + 1
        target_year = ref_date.year + (ref_date.month - 1 + n) // 12
        start = dt.date(target_year, target_month, 1)

        if n == 0:
            return start, ref_date

        next_month = (target_month % 12) + 1
        next_year = target_year + (target_month == 12)
        end = dt.date(next_year, next_month, 1) - dt.timedelta(days=1)
        return start, end

    # ---------------- Period registration ----------------

    @classmethod
    def register_period(cls, name: str, func: PeriodFunc) -> None:
        """
        Dynamically register a new named period.
        The callable must accept 'now' (datetime) and return (start_date, end_date).
        """
        cls._PERIOD_FUNCTIONS[name] = func

    @classmethod
    def register_pattern(
        cls,
        pattern: str | re.Pattern[str],
        resolver: Callable[[re.Match[str], dt.datetime], dict[str, Any]],
    ) -> None:
        """
        Register a regex-based dynamic period.
        The resolver receives (match, now).
        """
        compiled = re.compile(pattern) if isinstance(pattern, str) else pattern
        cls._PERIOD_PATTERNS.append((compiled, resolver))

    # ---------------- Default named periods ----------------

    @classmethod
    def _get_default_periods(cls) -> dict[str, PeriodFunc]:
        """
        Returns lambdas that accept 'now' (datetime) to ensure
        everyone uses the exact same frozen clock time.
        """
        return {
            "today": lambda now: (now.date(), now.date()),
            "yesterday": lambda now: (
                now.date() - dt.timedelta(days=1),
                now.date() - dt.timedelta(days=1),
            ),
            "current_week": lambda now: cls.calc_week_range(now.date()),
            "last_week": lambda now: cls.calc_week_range(
                now.date() - dt.timedelta(days=7)
            ),
            "current_month": lambda now: cls.get_month_range(now.date(), n=0),
            "last_month": lambda now: cls.get_month_range(now.date(), n=-1),
            "current_year": lambda now: cls.get_year_timerange(now.year),
            "last_year": lambda now: cls.get_year_timerange(now.year - 1),
            "current_quarter": lambda now: (
                cls.get_first_day_of_the_quarter(now.date()),
                cls.get_last_day_of_the_quarter(now.date()),
            ),
            "ytd": lambda now: (dt.date(now.year, 1, 1), now.date()),
            "itd": lambda now: (dt.date(1900, 1, 1), now.date()),
        }

    @classmethod
    def period_keys(cls) -> Iterable[str]:
        """List available named periods (defaults + registered)."""
        d = dict(cls._get_default_periods())
        d.update(cls._PERIOD_FUNCTIONS)
        return d.keys()

    # ---------------- Flexible resolver ----------------

    @classmethod
    def resolve_period(
        cls,
        period: str | None = None,
        *,
        now: dt.datetime | None = None,
        tz: dt.tzinfo = dt.timezone.utc,
        **overrides: Any,
    ) -> tuple[str, dict[str, Any]]:
        """
        Resolve a period into (canonical_key, params).
        Uses 'now' as the single source of truth for time.
        """
        key = (period or "today").strip()

        # Single Source of Truth for Time
        # If 'now' isn't provided, generate it in UTC immediately.
        if now is None:
            now = dt.datetime.now(tz)
        else:
            now = now.astimezone(tz)

        # 1) Named periods (Pass 'now' down!)
        period_functions = cls._get_default_periods()
        period_functions.update(cls._PERIOD_FUNCTIONS)  # type: ignore

        if key in period_functions:
            start, end = period_functions[key](now)
            params = {"start_on": start.isoformat(), "end_on": end.isoformat()}
            params.update(overrides)
            return key, params

        # 2) Regex patterns
        for patt, resolver in cls._PERIOD_PATTERNS:
            m = patt.fullmatch(key)
            if m:
                out = resolver(m, now)
                canonical = out.get("canonical", "custom")
                params = {k: v for k, v in out.items() if k != "canonical"}
                params.update(overrides)
                return canonical, params

        # 2b) default 'last_N_days'
        m = cls._LAST_N_DAYS_RE.match(key)
        if m:
            days = int(m.group(1))
            end_date = now.date()
            start_date = (now - dt.timedelta(days=days)).date()
            params = {
                "start_on": start_date.isoformat(),
                "end_on": end_date.isoformat(),
            }
            params.update(overrides)
            return "custom", params

        # 3) explicit date window: YYYY-MM-DD..YYYY-MM-DD
        m2 = cls._WINDOW_RE.fullmatch(key)
        if m2:
            start_on, end_on = m2.group(1), m2.group(2)
            params = {"start_on": start_on, "end_on": end_on}
            params.update(overrides)
            return "custom", params

        # 4) fallback (unknown key)
        return key, dict(overrides)

    # ---------------- Backward-compatible API ----------------

    @classmethod
    def parse_period(cls, **kwargs: Any) -> tuple[dt.date, dt.date]:
        """
        Return (start_date, end_date) as datetime.date.
        Warning: Uses system time if 'now' is not provided in kwargs.
        """
        period = kwargs.get("period", "today")

        # We call resolve_period to handle the logic
        canonical, params = cls.resolve_period(period, **kwargs)

        if "start_on" in params and "end_on" in params:
            start = cls._ensure_date(params["start_on"])
            end = cls._ensure_date(params["end_on"])
            return start, end

        if "start_ts" in params and "end_ts" in params:
            sdt = cls._ensure_datetime(params["start_ts"]).date()
            edt = cls._ensure_datetime(params["end_ts"]).date()
            return sdt, edt

        raise ValueError(
            f"Could not derive date range from period '{period}' (canonical='{canonical}'). "
            f"Params: {params}"
        )


# ---------------- Default dynamic patterns registration ----------------


def _register_default_patterns() -> None:
    """
    Register common dynamic patterns:
      - last_{n}_hours  (ISO datetimes; useful for freshness windows)
    """

    def last_x_hours(match: re.Match[str], now: dt.datetime) -> dict[str, Any]:
        hours = int(match.group(1))
        end_ts = now
        start_ts = now - dt.timedelta(hours=hours)

        return {
            "canonical": "custom",
            # Strict ISO 8601 with Z
            "start_ts": DateUtils.to_iso_z(start_ts),
            "end_ts": DateUtils.to_iso_z(end_ts),
            "max_age_minutes": max(15, min(hours * 10, 240)),
        }

    DateUtils.register_pattern(r"last_(\d+)_hours", last_x_hours)


# Register defaults at import time
_register_default_patterns()
