from datetime import datetime
from zoneinfo import ZoneInfo

import pytest

from app.bars.models import Timeframe
from app.bars.session import bucket_start

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


def epoch(value: str, zone=UTC) -> int:
    return int(datetime.fromisoformat(value).replace(tzinfo=zone).timestamp())


@pytest.mark.parametrize(
    ("value", "tf", "expected"),
    [
        ("2026-08-09T22:00:00", Timeframe.D1, "2026-08-09T22:00:00"),
        ("2026-08-10T16:37:00", Timeframe.D1, "2026-08-09T22:00:00"),
        ("2026-08-14T20:59:00", Timeframe.D1, "2026-08-13T22:00:00"),
        ("2026-08-10T21:30:00", Timeframe.D1, "2026-08-09T22:00:00"),
        ("2026-08-10T22:00:00", Timeframe.D1, "2026-08-10T22:00:00"),
        ("2026-08-10T01:59:00", Timeframe.H4, "2026-08-09T22:00:00"),
        ("2026-08-10T02:00:00", Timeframe.H4, "2026-08-10T02:00:00"),
        ("2026-08-10T17:59:00", Timeframe.H4, "2026-08-10T14:00:00"),
    ],
)
def test_session_boundaries(value, tf, expected):
    assert bucket_start(epoch(value), tf) == epoch(expected)


def test_intraday_buckets_use_utc_boundaries():
    assert bucket_start(epoch("2026-08-10T12:37:45"), Timeframe.M15) == epoch(
        "2026-08-10T12:30:00"
    )


@pytest.mark.parametrize(
    ("value", "expected"),
    [
        # Spring forward: the 22:00 ET bucket ends at 02:00 EDT after three real hours.
        ("2026-03-08T06:59:00", "2026-03-08T03:00:00"),
        ("2026-03-08T07:00:00", "2026-03-08T07:00:00"),
        # Fall back: the 22:00 ET bucket lasts five real hours and ends at 02:00 EST.
        ("2026-11-01T06:59:00", "2026-11-01T02:00:00"),
        ("2026-11-01T07:00:00", "2026-11-01T07:00:00"),
    ],
)
def test_four_hour_wall_clock_anchor_across_dst(value, expected):
    assert bucket_start(epoch(value), Timeframe.H4) == epoch(expected)


@pytest.mark.parametrize(
    ("local_value", "expected_local"),
    [
        ("2026-03-08T18:00:00", "2026-03-08T18:00:00"),
        ("2026-11-01T18:00:00", "2026-11-01T18:00:00"),
    ],
)
def test_sunday_open_across_dst(local_value, expected_local):
    assert bucket_start(epoch(local_value, ET), Timeframe.D1) == epoch(expected_local, ET)