"""Trading window utilities for the autonomous trading engine.

Pure computation module that determines whether a given timestamp falls
within the allowed trading window (9:45 AM – 3:45 PM ET on weekdays),
whether the US market is open, and when the next trading window opens.

Uses ``zoneinfo.ZoneInfo("America/New_York")`` for Eastern Time handling.
Checks major US market holidays for 2026.
"""

from __future__ import annotations

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

# US Eastern timezone
ET = ZoneInfo("America/New_York")

# Trading window boundaries (excludes first/last 15 min of market hours)
WINDOW_OPEN = time(9, 45)
WINDOW_CLOSE = time(15, 45)

# Full US market hours
MARKET_OPEN = time(9, 30)
MARKET_CLOSE = time(16, 0)

# Weekday range: Monday=0 … Friday=4
_WEEKDAYS = range(0, 5)


def _us_market_holidays_2026() -> set[date]:
    """Return a set of US market holiday dates for 2026.

    Major holidays observed by NYSE/NASDAQ:
    - New Year's Day (Jan 1)
    - MLK Day (3rd Monday of January)
    - Presidents' Day (3rd Monday of February)
    - Good Friday (April 3)
    - Memorial Day (last Monday of May)
    - Juneteenth (June 19)
    - Independence Day (July 3 observed — July 4 is Saturday)
    - Labor Day (1st Monday of September)
    - Thanksgiving (4th Thursday of November)
    - Christmas (Dec 25)
    """
    return {
        date(2026, 1, 1),    # New Year's Day
        date(2026, 1, 19),   # MLK Day (3rd Monday)
        date(2026, 2, 16),   # Presidents' Day (3rd Monday)
        date(2026, 4, 3),    # Good Friday
        date(2026, 5, 25),   # Memorial Day (last Monday)
        date(2026, 6, 19),   # Juneteenth
        date(2026, 7, 3),    # Independence Day (observed)
        date(2026, 9, 7),    # Labor Day (1st Monday)
        date(2026, 11, 26),  # Thanksgiving (4th Thursday)
        date(2026, 12, 25),  # Christmas
    }


_HOLIDAYS_2026 = _us_market_holidays_2026()


def is_within_trading_window(dt: datetime) -> bool:
    """Return True if *dt* is between 9:45 AM ET and 3:45 PM ET on a weekday.

    The timestamp is first converted to US/Eastern time.  Weekends and
    US market holidays (2026) are always outside the window.
    """
    et_dt = dt.astimezone(ET)
    if et_dt.weekday() not in _WEEKDAYS:
        return False
    if et_dt.date() in _HOLIDAYS_2026:
        return False
    t = et_dt.time()
    return WINDOW_OPEN <= t < WINDOW_CLOSE


def next_window_open(dt: datetime) -> datetime:
    """Return the next datetime when the trading window opens (9:45 AM ET).

    If *dt* is before 9:45 AM ET on a weekday the same day's open is
    returned.  Otherwise the next weekday's 9:45 AM ET is returned.
    """
    et_dt = dt.astimezone(ET)
    today_open = et_dt.replace(
        hour=WINDOW_OPEN.hour,
        minute=WINDOW_OPEN.minute,
        second=0,
        microsecond=0,
    )

    # If we haven't reached today's open yet and it's a weekday, return today
    if et_dt < today_open and et_dt.weekday() in _WEEKDAYS:
        return today_open

    # Otherwise advance to the next weekday
    candidate = et_dt + timedelta(days=1)
    candidate = candidate.replace(
        hour=WINDOW_OPEN.hour,
        minute=WINDOW_OPEN.minute,
        second=0,
        microsecond=0,
    )
    while candidate.weekday() not in _WEEKDAYS:
        candidate += timedelta(days=1)
    return candidate


def is_market_open(dt: datetime) -> bool:
    """Return True if *dt* is during US market hours (9:30 AM – 4:00 PM ET) on a weekday.

    Returns False on weekends and US market holidays (2026).
    """
    et_dt = dt.astimezone(ET)
    if et_dt.weekday() not in _WEEKDAYS:
        return False
    if et_dt.date() in _HOLIDAYS_2026:
        return False
    t = et_dt.time()
    return MARKET_OPEN <= t < MARKET_CLOSE