from __future__ import annotations

from datetime import UTC, date, datetime
from pathlib import Path

from findata_market_firstrate.archive import parse_bar_file
from findata_quality import aggregate_bars, validate_bars
from findata_quality.calendar import xnys
from findata_quality.validators import session_coverage
from findata_schema import Adjustment, AssetClass, Interval, QualityFlag

FIX = Path(__file__).parent / "fixtures" / "firstrate"
NOW = datetime(2026, 11, 1, 12, 0, tzinfo=UTC)


def _parse(text: str, symbol="TEST"):
    return parse_bar_file(text.encode(), symbol, AssetClass.stock, Interval.m1, Adjustment.raw, "firstrate", NOW)[0]


def _flags(t):
    return [set(f) for f in t.column("quality_flags").to_pylist()]


def test_clean_real_session_has_no_error_flags():
    t = parse_bar_file((FIX / "AAPL_day_1min_UNADJUSTED.txt").read_bytes(), "AAPL", AssetClass.stock, Interval.m1, Adjustment.raw, "firstrate", NOW)[0]
    out, anomalies = validate_bars(t, Interval.m1)
    errors = [a for a in anomalies if a["severity"] == "error"]
    assert errors == []
    assert out.num_rows == 900


def test_row_level_flags():
    rows = "\n".join([
        "2026-10-30 09:30:00,10,9,8,9.5,100",      # high < open → ohlc_inconsistent
        "2026-10-30 09:31:00,10,11,9,10,-5",       # negative volume
        "2026-10-30 09:32:30,10,11,9,10,5",        # off interval grid
        "2026-10-30 09:33:00,10,11,9,10,5",
        "2026-10-30 09:33:00,10,11,9,10,5",        # duplicate
        "2026-10-30 03:59:00,10,11,9,10,5",        # before 04:00 ET → outside_session
        "2026-10-30 09:34:00,10,20,9,15,5",        # +50 % vs previous close → price_jump
        "2026-10-30 09:35:00,0,0,0,0,5",           # nonpositive price
        "2026-11-01 09:30:00,10,11,9,10,5",        # Sunday → non_trading_day
    ])
    out, anomalies = validate_bars(_parse(rows), Interval.m1)
    by_ts = {r["timestamp_utc"]: set(r["quality_flags"]) for r in out.to_pylist()}
    f = lambda h, m, d=30: by_ts[datetime(2026, 10, d, h, m, tzinfo=UTC)]
    assert QualityFlag.ohlc_inconsistent in f(13, 30)
    assert QualityFlag.negative_volume in f(13, 31)
    assert any(QualityFlag.off_interval in v for v in by_ts.values())
    assert QualityFlag.duplicate in f(13, 33)
    assert QualityFlag.outside_session in f(7, 59)
    assert QualityFlag.price_jump in f(13, 34)
    assert QualityFlag.nonpositive_price in f(13, 35)
    assert QualityFlag.non_trading_day in by_ts[datetime(2026, 11, 1, 14, 30, tzinfo=UTC)]
    codes = {a["code"] for a in anomalies}
    assert {"ohlc_inconsistent", "negative_volume", "duplicate", "non_trading_day"} <= codes
    # prices are never altered by validation
    assert out.column("open").to_pylist()[0] == 10.0


def test_session_calendar_expected_bars_and_early_close():
    cal = xnys()
    assert cal.is_session(date(2026, 10, 30)) and not cal.is_session(date(2026, 11, 1))
    assert not cal.is_session(date(2026, 11, 26))                      # Thanksgiving
    assert cal.expected_regular_bars(date(2026, 10, 30)) == 391          # 09:30..16:00 inclusive
    assert cal.expected_regular_bars(date(2026, 11, 27)) == 211          # early close 13:00
    assert cal.expected_regular_bars(date(2026, 11, 1)) == 0


def test_session_coverage_rows():
    t = parse_bar_file((FIX / "A_day_1min_UNADJUSTED.txt").read_bytes(), "A", AssetClass.stock, Interval.m1, Adjustment.raw, "firstrate", NOW)[0]
    cov = session_coverage(t, Interval.m1)
    assert len(cov) == 1 and cov[0]["symbol"] == "A" and cov[0]["bar_count"] == 390 and cov[0]["expected_bars"] == 391


def test_aggregate_5m_matches_manual_expectation():
    t = parse_bar_file((FIX / "AAPL_day_1min_UNADJUSTED.txt").read_bytes(), "AAPL", AssetClass.stock, Interval.m1, Adjustment.raw, "firstrate", NOW)[0]
    agg = aggregate_bars(t, Interval.m5)
    import polars as pl
    one = pl.from_arrow(t).sort("timestamp_utc")
    five = pl.from_arrow(agg).sort("timestamp_utc")
    # 09:30 ET bucket = 13:30 UTC covers 13:30..13:34
    win = one.filter((pl.col("timestamp_utc") >= datetime(2026, 10, 30, 13, 30, tzinfo=UTC)) & (pl.col("timestamp_utc") < datetime(2026, 10, 30, 13, 35, tzinfo=UTC)))
    b = five.filter(pl.col("timestamp_utc") == datetime(2026, 10, 30, 13, 30, tzinfo=UTC)).row(0, named=True)
    assert b["open"] == win["open"][0] and b["close"] == win["close"][-1]
    assert b["high"] == win["high"].max() and b["low"] == win["low"].min()
    assert abs(b["volume"] - win["volume"].sum()) < 1e-6 and b["trade_count"] == 5
    assert b["interval"] == "5m" and b["session_date"] == date(2026, 10, 30)
    assert abs(five["volume"].sum() - one["volume"].sum()) < 1e-6          # volume conserved
    assert five.height <= 16 * 12                                          # 04:00–20:00 → ≤ 192 buckets, none manufactured
    assert all(ts.minute % 5 == 0 for ts in five["timestamp_utc"].to_list())


def test_aggregate_buckets_follow_wall_clock_across_dst():
    # 09:30 EST (14:30 UTC) and 09:30 EDT (13:30 UTC) must each start a bucket
    t = _parse("2026-03-06 09:31:00,1,1,1,1,1\n2026-03-06 09:34:00,2,2,2,2,1\n2026-03-09 09:31:00,3,3,3,3,1\n2026-03-09 09:34:00,4,4,4,4,1\n")
    agg = aggregate_bars(t, Interval.m5)
    ts = agg.column("timestamp_utc").to_pylist()
    assert ts == [datetime(2026, 3, 6, 14, 30, tzinfo=UTC), datetime(2026, 3, 9, 13, 30, tzinfo=UTC)]
    assert agg.column("open").to_pylist() == [1.0, 3.0] and agg.column("close").to_pylist() == [2.0, 4.0]


def test_aggregate_zero_volume_bucket_has_null_vwap():
    t = _parse("2026-10-30 09:30:00,1,1,1,1,0\n")
    agg = aggregate_bars(t, Interval.m5)
    assert agg.column("vwap")[0].as_py() is None