"""Deterministic bar aggregation (1m → 5m/30m/1h). Buckets are aligned on the ET wall clock so that the
09:30 bucket covers 09:30:00–09:34:59 regardless of DST. Never fills empty buckets."""
from __future__ import annotations

import polars as pl
import pyarrow as pa
from findata_schema import BARS_SCHEMA, INTERVAL_SECONDS, Interval


def aggregate_bars(table: pa.Table, target: Interval, source: Interval = Interval.m1) -> pa.Table:
    if INTERVAL_SECONDS[target] % INTERVAL_SECONDS[source] != 0:
        raise ValueError(f"{target} is not a multiple of {source}")
    df = pl.from_arrow(table)
    if df.height == 0:
        return table.cast(BARS_SCHEMA)
    every = f"{INTERVAL_SECONDS[target]}s"
    # Truncate in local exchange time, then convert back: bucket boundaries follow the wall clock.
    local = pl.col("timestamp_utc").dt.convert_time_zone("America/New_York")
    df = df.with_columns(local.dt.truncate(every).dt.convert_time_zone("UTC").alias("bucket"))
    df = df.sort(["symbol_as_reported", "timestamp_utc"])
    agg = (df.group_by(["instrument_id", "symbol_as_reported", "venue", "asset_class", "session_date", "currency",
                        "adjustment_type", "source_id", "bucket"], maintain_order=True)
           .agg(pl.col("open").first(), pl.col("high").max(), pl.col("low").min(), pl.col("close").last(),
                pl.col("volume").sum(), pl.len().cast(pl.Int32).alias("trade_count"),
                ((pl.col("close") * pl.col("volume")).sum() / pl.col("volume").sum()).alias("vwap"),
                pl.col("ingested_at").max(), pl.col("quality_flags").list.explode(keep_nulls=False, empty_as_null=False).drop_nulls().unique().alias("quality_flags"),
                pl.col("revision_id").max()))
    agg = agg.rename({"bucket": "timestamp_utc"}).with_columns(pl.lit(target.value).alias("interval"))
    agg = agg.with_columns(pl.when(pl.col("volume") > 0).then(pl.col("vwap")).otherwise(None).alias("vwap"))
    out = agg.select(BARS_SCHEMA.names).to_arrow()
    return out.cast(BARS_SCHEMA)