"""XNYS trading calendar helpers (America/New_York sessions, UTC storage)."""
from __future__ import annotations

import datetime as dt
from functools import lru_cache
from zoneinfo import ZoneInfo

import numpy as np
import pandas as pd

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


@lru_cache(maxsize=1)
def _xnys():
    import exchange_calendars as xc
    return xc.get_calendar("XNYS", start="2000-01-01", end="2027-12-31")


@lru_cache(maxsize=1)
def sessions() -> pd.DatetimeIndex:
    """All XNYS sessions 2000→2027 as tz-naive dates."""
    return _xnys().sessions.tz_localize(None) if _xnys().sessions.tz is not None else _xnys().sessions


def session_array() -> np.ndarray:
    return sessions().to_numpy().astype("datetime64[D]")


def is_session(day: dt.date) -> bool:
    return np.datetime64(day) in set(session_array().tolist())


def shift_trading_days(day: dt.date | np.datetime64 | pd.Timestamp, n: int) -> dt.date:
    """Shift by n trading days; if `day` is not a session, start from the previous session (n<=0) or next (n>0)."""
    arr = session_array()
    d = np.datetime64(pd.Timestamp(day).date())
    pos = int(np.searchsorted(arr, d))
    if pos < len(arr) and arr[pos] == d:
        base = pos
    else:
        base = pos - 1 if n <= 0 else pos
    idx = base + n
    if idx < 0 or idx >= len(arr):
        raise ValueError("date out of calendar range")
    return arr[idx].astype("datetime64[D]").astype(dt.date)


def trading_days_between(a: dt.date, b: dt.date) -> int:
    """Number of sessions strictly after a up to and including b (negative if b<a)."""
    arr = session_array()
    ia, ib = np.searchsorted(arr, np.datetime64(a), side="right"), np.searchsorted(arr, np.datetime64(b), side="right")
    return int(ib - ia)


def month_end_sessions(start: dt.date, end: dt.date) -> list[dt.date]:
    s = pd.Series(sessions())
    s = s[(s >= pd.Timestamp(start)) & (s <= pd.Timestamp(end))]
    return [d.date() for d in s.groupby(s.dt.to_period("M")).max()]


def close_time_utc(session: dt.date) -> pd.Timestamp:
    """Open instant (UTC) of the closing-print bar (16:00 ET, or 13:00 ET on early-close days)."""
    cal = _xnys()
    ts = pd.Timestamp(session)
    close = cal.session_close(ts)
    return pd.Timestamp(close).tz_convert("UTC") if close.tz is not None else pd.Timestamp(close).tz_localize("UTC")


def regular_session_bounds_utc(session: dt.date) -> tuple[pd.Timestamp, pd.Timestamp]:
    """[open 09:30 ET, close bar open] in UTC for the session (inclusive bar-open bounds)."""
    cal = _xnys()
    ts = pd.Timestamp(session)
    o, c = cal.session_open(ts), cal.session_close(ts)
    o = pd.Timestamp(o).tz_convert("UTC") if o.tz is not None else pd.Timestamp(o).tz_localize("UTC")
    c = pd.Timestamp(c).tz_convert("UTC") if c.tz is not None else pd.Timestamp(c).tz_localize("UTC")
    return o, c