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

from app.bars.models import Timeframe

UTC = ZoneInfo("UTC")
EASTERN = ZoneInfo("America/New_York")
SESSION_OPEN = time(18, 0)


def _session_open_local(current: datetime) -> datetime:
    session_date = current.date() if current.timetz().replace(tzinfo=None) >= SESSION_OPEN else current.date() - timedelta(days=1)
    return datetime.combine(session_date, SESSION_OPEN, EASTERN)


def bucket_start(t: int, tf: Timeframe) -> int:
    if tf not in (Timeframe.H4, Timeframe.D1):
        return (t // tf.seconds) * tf.seconds

    current = datetime.fromtimestamp(t, UTC).astimezone(EASTERN)
    session_open = _session_open_local(current)
    if tf is Timeframe.D1:
        return int(session_open.timestamp())

    # CME's 4h anchors are wall-clock ET anchors. This intentionally makes the
    # DST-transition bucket three or five elapsed hours instead of shifting it.
    elapsed_wall = current.replace(tzinfo=None) - session_open.replace(tzinfo=None)
    bucket_hours = int(elapsed_wall.total_seconds() // 14400) * 4
    local_start = session_open.replace(tzinfo=None) + timedelta(hours=bucket_hours)
    return int(local_start.replace(tzinfo=EASTERN).timestamp())