"""Exchange-calendar wrapper: US/KRX sessions, close times and briefing schedule.

All date arithmetic goes through `exchange_calendars` (XNYS / XKRX) and `zoneinfo`; no UTC
offsets are hardcoded, so DST switches, holidays and early closes are handled automatically.
"""

from __future__ import annotations

import logging
from collections.abc import Iterable
from datetime import date, datetime, timedelta
from zoneinfo import ZoneInfo

import exchange_calendars as xc
import pandas as pd

from us2kr.config import ET, KST

log = logging.getLogger(__name__)

ET_TZ = ZoneInfo(ET)
KST_TZ = ZoneInfo(KST)

_US_CAL = "XNYS"
_KR_CAL = "XKRX"

# KRX closures announced by KRX / the Korean government but missing from the pinned exchange_calendars
# data (4.13.2). Election days and 임시공휴일 are decided after a library release, so the data lags.
# Drop an entry once a release marks it as a holiday (the override is harmless while both agree).
KNOWN_KRX_CLOSURES: frozenset[date] = frozenset(
    {
        date(2025, 10, 10),  # 임시공휴일 (추석 연휴 연장)
        date(2026, 6, 3),  # 제9회 전국동시지방선거
    }
)


class MarketClock:
    """Session / briefing-time calculator backed by XNYS and XKRX calendars.

    Calendars are built once for a window around `anchor` (default: today) and rebuilt
    transparently if a query falls outside that window (long-running daemon safety).
    """

    def __init__(
        self,
        run_offset_minutes: int = 75,
        *,
        anchor: date | None = None,
        years_back: int = 2,
        years_ahead: int = 2,
        extra_krx_holidays: Iterable[date] | None = None,
    ) -> None:
        if run_offset_minutes < 0:
            raise ValueError("run_offset_minutes must be >= 0")
        self.run_offset = timedelta(minutes=run_offset_minutes)
        self._years_back = years_back
        self._years_ahead = years_ahead
        # Ad-hoc KRX closures on top of the library data (KNOWN_KRX_CLOSURES + operator-supplied dates).
        self._krx_closures: frozenset[date] = KNOWN_KRX_CLOSURES | frozenset(extra_krx_holidays or ())
        self._build(anchor or datetime.now(tz=KST_TZ).date())

    # ------------------------------------------------------------------ calendars

    # Pad the requested range so Jan 1 .. Dec 31 of the window always lie strictly inside
    # [first_session, last_session] of both calendars (no exchange closes > ~5 days in a row).
    _PAD = timedelta(days=10)

    def _build(self, anchor: date) -> None:
        start = date(anchor.year - self._years_back, 1, 1) - self._PAD
        end = date(anchor.year + self._years_ahead, 12, 31) + self._PAD
        self._xnys = xc.get_calendar(_US_CAL, start=start, end=end)
        self._xkrx = xc.get_calendar(_KR_CAL, start=start, end=end)
        # Bounds common to both calendars (sessions are tz-naive Timestamps in v4).
        self._first = max(self._xnys.first_session, self._xkrx.first_session).date()
        self._last = min(self._xnys.last_session, self._xkrx.last_session).date()
        log.debug("calendars built: %s .. %s", self._first, self._last)

    @staticmethod
    def _as_date(d: date | datetime) -> date:
        return d.date() if isinstance(d, datetime) else d

    def _ts(self, d: date | datetime) -> pd.Timestamp:
        """Normalise to a tz-naive midnight Timestamp inside the calendar range.

        May rebuild `self._xnys` / `self._xkrx`, so callers must evaluate this BEFORE
        binding a calendar method (`ts = self._ts(d); self._xnys.f(ts)`).
        """
        d = self._as_date(d)
        if not (self._first <= d <= self._last):
            log.info("date %s outside calendar window %s..%s; rebuilding", d, self._first, self._last)
            self._build(d)
        return pd.Timestamp(d)

    @staticmethod
    def _require_aware(now: datetime) -> datetime:
        if now.tzinfo is None or now.utcoffset() is None:
            raise ValueError("`now` must be a tz-aware datetime")
        return now

    # ------------------------------------------------------------------ sessions

    def is_us_session(self, d: date) -> bool:
        """True if XNYS trades on calendar date `d`."""
        ts = self._ts(d)
        return bool(self._xnys.is_session(ts))

    def is_krx_session(self, d: date) -> bool:
        """True if KRX trades on calendar date `d` (library data + known ad-hoc closures)."""
        if self._as_date(d) in self._krx_closures:
            return False
        ts = self._ts(d)
        return bool(self._xkrx.is_session(ts))

    def next_krx_session(self, on_or_after: date) -> date:
        """Earliest KRX session on or after `on_or_after` (ad-hoc closures respected)."""
        d = self._as_date(on_or_after)
        while True:
            ts = self._ts(d)
            d = self._xkrx.date_to_session(ts, direction="next").date()
            if d not in self._krx_closures:
                return d
            d += timedelta(days=1)

    def us_session_close(self, session: date) -> datetime:
        """Regular-session close as tz-aware America/New_York datetime (early closes respected)."""
        ts = self._ts(session)
        if not self._xnys.is_session(ts):
            raise ValueError(f"{self._as_date(session)} is not an XNYS session")
        close_utc: pd.Timestamp = self._xnys.session_close(ts)
        return close_utc.to_pydatetime().astimezone(ET_TZ)

    def previous_us_session(self, before: date) -> date:
        """Latest US session strictly before `before`."""
        ts = self._ts(self._as_date(before) - timedelta(days=1))
        return self._xnys.date_to_session(ts, direction="previous").date()

    def next_us_session(self, after: date) -> date:
        """Earliest US session strictly after `after`."""
        ts = self._ts(self._as_date(after) + timedelta(days=1))
        return self._xnys.date_to_session(ts, direction="next").date()

    def latest_completed_session(self, now: datetime) -> date:
        """Latest US session whose regular close is <= `now`."""
        now = self._require_aware(now)
        today_et = now.astimezone(ET_TZ).date()
        ts = self._ts(today_et)
        candidate = self._xnys.date_to_session(ts, direction="previous").date()
        while self.us_session_close(candidate) > now:
            candidate = self.previous_us_session(candidate)
        return candidate

    # ------------------------------------------------------------------ briefing schedule

    def briefing_time_for(self, session: date) -> datetime:
        """Briefing run time = session close + run offset, expressed in Asia/Seoul."""
        return (self.us_session_close(session) + self.run_offset).astimezone(KST_TZ)

    def krx_date_for(self, session: date) -> date:
        """Korean calendar date on which the briefing for `session` is delivered."""
        return self.briefing_time_for(session).date()

    def next_briefing(self, now: datetime) -> tuple[date, datetime]:
        """(session, run_at KST) for the earliest briefing strictly after `now`."""
        now = self._require_aware(now)
        today_et = now.astimezone(ET_TZ).date()
        session = self.previous_us_session(today_et)
        while self.briefing_time_for(session) <= now:
            session = self.next_us_session(session)
        return session, self.briefing_time_for(session)

    def catchup_candidate(self, now: datetime, window_hours: int) -> date | None:
        """Session whose briefing time lies in (now - window, now]; None when nothing was missed."""
        now = self._require_aware(now)
        window = timedelta(hours=window_hours)
        session = self.latest_completed_session(now)
        while True:
            run_at = self.briefing_time_for(session)
            if run_at <= now:
                return session if run_at > now - window else None
            session = self.previous_us_session(session)