"""Data loading for the crypto-quant-eval dataset.

Loads curated parquet files from the dataset repo or a local copy,
asserts schema and no-lookahead invariants, and exposes simple accessors.
"""

from __future__ import annotations

import hashlib
import json
from dataclasses import dataclass
from pathlib import Path

import pandas as pd

REQUIRED_BARS = {
    "datetime", "open", "high", "low", "close", "volume",
}
REQUIRED_FUNDING = {"timestamp", "funding_rate"}


def sha256_file(path: str | Path) -> str:
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 20), b""):
            h.update(chunk)
    return h.hexdigest()


@dataclass
class EvalData:
    """Curated bar data plus aligned funding ticks."""

    bars: pd.DataFrame        # index: UTC datetime, OHLCV columns
    funding: pd.DataFrame     # index: UTC datetime (8h ticks), funding_rate
    symbol: str
    freq: str

    def assert_sane(self) -> None:
        # bar start time may be carried either as the DatetimeIndex or as
        # a 'datetime' column; the index form is what load_local produces.
        bar_cols = set(self.bars.columns)
        if isinstance(self.bars.index, pd.DatetimeIndex):
            bar_cols = bar_cols | {"datetime"}
        missing = REQUIRED_BARS - bar_cols
        if missing:
            raise ValueError(f"bars missing columns: {sorted(missing)}")
        fund_cols = set(self.funding.columns)
        if isinstance(self.funding.index, pd.DatetimeIndex):
            fund_cols = fund_cols | {"timestamp"}
        missing = REQUIRED_FUNDING - fund_cols
        if missing:
            raise ValueError(f"funding missing columns: {sorted(missing)}")
        if not self.bars.index.is_monotonic_increasing:
            raise ValueError("bars index is not sorted")
        if self.bars.index.has_duplicates:
            raise ValueError("bars index has duplicate timestamps")
        # no-lookahead invariant: bars are closed bars only.
        # The curated dataset stores bar START time; a bar starting at t is
        # only usable at t+freq. load() shifts nothing here — consumers must
        # use shift-based accessors below, never bars.loc[t] inside a
        # signal computed at t.
        nulls = self.bars[["open", "high", "low", "close"]].isna().sum().sum()
        if nulls:
            raise ValueError(f"bars contain {nulls} null OHLC values")
        bad_ohlc = (
            (self.bars["high"] < self.bars["low"])
            | (self.bars["close"] > self.bars["high"])
            | (self.bars["open"] > self.bars["high"])
            | (self.bars["close"] < self.bars["low"])
            | (self.bars["open"] < self.bars["low"])
        ).sum()
        if bad_ohlc:
            raise ValueError(f"{bad_ohlc} bars violate OHLC ordering")

    def check_funding_coverage(self, freq: str) -> int:
        """Return count of funding ticks falling inside the bar range."""
        mask = (self.funding.index >= self.bars.index[0]) & (
            self.funding.index <= self.bars.index[-1]
        )
        return int(mask.sum())


def load_local(root: str | Path, symbol: str = "BTCUSDT", freq: str = "1h") -> EvalData:
    """Load bars_<symbol>_<freq>.parquet and funding_<symbol>.parquet from root."""
    root = Path(root)
    bars = pd.read_parquet(root / f"bars_{symbol}_{freq}.parquet")
    funding = pd.read_parquet(root / f"funding_{symbol}.parquet")
    # curated parquet stores bar start time as a 'datetime' column; the
    # synthetic/local layout keeps a DatetimeIndex. Accept both.
    for df in (bars, funding):
        if "datetime" in df.columns and not isinstance(df.index, pd.DatetimeIndex):
            ts = pd.to_datetime(df["datetime"], utc=True)
            df.index = ts
        elif "timestamp" in df.columns and not isinstance(df.index, pd.DatetimeIndex):
            df.index = pd.to_datetime(df["timestamp"], utc=True)
    bars = bars.drop(columns=["datetime"], errors="ignore")
    funding = funding.drop(columns=["datetime", "timestamp"], errors="ignore")
    for df in (bars, funding):
        if not isinstance(df.index, pd.DatetimeIndex):
            raise ValueError("parquet files must carry a datetime index")
        if df.index.tz is None:
            df.index = df.index.tz_localize("UTC")
        else:
            df.index = df.index.tz_convert("UTC")
    data = EvalData(bars=bars, funding=funding, symbol=symbol, freq=freq)
    data.assert_sane()
    return data


def load_manifest(root: str | Path) -> dict:
    """Read the source-hash manifest; every eval run must record this."""
    return json.loads((Path(root) / "manifest.json").read_text())