"""Market session calendar for US equities (XNYS) with extended hours and DST handled via zoneinfo.

Policy (documented in docs/DATA_DICTIONARY.md):
- Regular session 09:30–16:00 ET; vendor extended hours 04:00–20:00 ET (last 1-minute bar opens 19:59).
- Early closes (e.g. day after Thanksgiving, Christmas Eve) come from exchange_calendars; regular-hours
  expectations use the calendar's actual close for that date.
- `expected_regular_bars(date)` = minutes between open and close (inclusive of the 16:00 bar which the vendor
  emits as the closing auction print) — used for coverage %, never to manufacture bars.
"""
from __future__ import annotations

from datetime import date, datetime, time, timedelta
from functools import lru_cache
from zoneinfo import ZoneInfo

import exchange_calendars as xcals
import pandas as pd

ET = ZoneInfo("America/New_York")
UTC = ZoneInfo("UTC")

EXT_OPEN = time(4, 0)
EXT_CLOSE = time(20, 0)   # exclusive


class SessionCalendar:
    def __init__(self, name: str = "XNYS", start: str = "1999-01-01", end: str = "2030-12-31"):
        self.name = name
        self.cal = xcals.get_calendar(name, start=start, end=end)
        self._sessions = set(d.date() for d in self.cal.sessions)

    def is_session(self, d: date) -> bool:
        return d in self._sessions

    def regular_open_close(self, d: date) -> tuple[datetime, datetime]:
        """Open/close as timezone-aware ET datetimes for session `d` (handles early closes and DST)."""
        ts = pd.Timestamp(d)
        o = self.cal.session_open(ts).tz_convert("America/New_York").to_pydatetime()
        c = self.cal.session_close(ts).tz_convert("America/New_York").to_pydatetime()
        return o, c

    def expected_regular_bars(self, d: date, interval_seconds: int = 60) -> int:
        if not self.is_session(d):
            return 0
        o, c = self.regular_open_close(d)
        return int((c - o).total_seconds() // interval_seconds) + 1  # + closing print bar

    def extended_bounds_utc(self, d: date) -> tuple[datetime, datetime]:
        start = datetime.combine(d, EXT_OPEN, tzinfo=ET).astimezone(UTC)
        end = datetime.combine(d, EXT_CLOSE, tzinfo=ET).astimezone(UTC)
        return start, end

    def sessions_between(self, start: date, end: date) -> list[date]:
        return [d for d in sorted(self._sessions) if start <= d <= end]

    def previous_session(self, d: date) -> date:
        cur = d - timedelta(days=1)
        while cur not in self._sessions:
            cur -= timedelta(days=1)
        return cur


@lru_cache
def xnys() -> SessionCalendar:
    return SessionCalendar("XNYS")