"""Data Management — DataManager, Universe, HiveStorage, and incremental updates."""

import logging
import shutil
from datetime import datetime
from pathlib import Path

import plotly.graph_objects as go
import polars as pl
import structlog

# ml4t-data emits structured debug logs on every fetch/store; route them
# through stdlib logging at WARNING so the notebook output stays focused on
# the demonstration.
structlog.configure(
    wrapper_class=structlog.make_filtering_bound_logger(logging.WARNING),
)

# ml4t-data core imports
from ml4t.data import DataManager
from ml4t.data.storage import HiveStorage
from ml4t.data.storage.backend import StorageConfig
from ml4t.data.universe import Universe

from utils.paths import REPO_ROOT, get_output_dir
from utils.style import COLORS


def _rel(path):
    """Repo-relative display path (keeps absolute machine paths out of outputs)."""
    try:
        return path.relative_to(REPO_ROOT)
    except ValueError:
        return path


# Working directory for this notebook's storage examples. Wipe any artifacts
# from a previous run so the demo is fully reproducible.
STORAGE_DIR = get_output_dir(2, "data_management")
if STORAGE_DIR.exists():
    shutil.rmtree(STORAGE_DIR)
STORAGE_DIR.mkdir(parents=True, exist_ok=True)

print(f"Storage directory: {_rel(STORAGE_DIR)}")

# Production defaults — Papermill injects overrides for CI

# DataManager without storage — pure fetch mode
dm = DataManager()

# Fetch a single symbol (defaults to Yahoo Finance for equities)
aapl = dm.fetch("AAPL", "2024-01-01", "2024-12-31", provider="yahoo")

print(f"AAPL: {aapl.shape[0]} rows, {aapl.shape[1]} columns")
print(f"Date range: {aapl['timestamp'].min().date()} to {aapl['timestamp'].max().date()}")
print(f"Columns: {aapl.columns}")
aapl.head(3)

# Fetch 5 ETFs in parallel
etf_symbols = ["SPY", "QQQ", "IWM", "TLT", "GLD"]
etf_data = dm.batch_load(
    symbols=etf_symbols,
    start="2024-01-01",
    end="2024-12-31",
    provider="yahoo",
    max_workers=4,
)

print(f"Combined: {etf_data.shape[0]:,} rows across {etf_data['symbol'].n_unique()} symbols")

# The stacked frame is easier to read as a picture than as a row count. Rebase
# each ETF's close to 100 at the first 2024 session and color by asset class:
# the batch is one call, but the panel spans equities (SPY/QQQ/IWM), long bonds
# (TLT), and gold (GLD). Color carries the asset class, not the ticker — the
# three equity lines move as a bundle while bonds and gold pull away, which is
# exactly the cross-asset dispersion a multi-asset loader exists to capture.
etf_rebased = etf_data.sort("timestamp").with_columns(
    (pl.col("close") / pl.col("close").first().over("symbol") * 100).alias("rebased")
)

etf_class = {
    "SPY": "Equities",
    "QQQ": "Equities",
    "IWM": "Equities",
    "TLT": "Long bonds",
    "GLD": "Gold",
}
class_color = {"Equities": COLORS["blue"], "Long bonds": COLORS["copper"], "Gold": COLORS["amber"]}

fig = go.Figure()
seen: set[str] = set()
for sym in etf_symbols:
    s = etf_rebased.filter(pl.col("symbol") == sym)
    cls = etf_class[sym]
    fig.add_trace(
        go.Scatter(
            x=s["timestamp"].to_list(),
            y=s["rebased"].to_list(),
            mode="lines",
            line=dict(color=class_color[cls], width=1.5),
            name=cls,
            legendgroup=cls,
            showlegend=cls not in seen,
            text=sym,
            hovertemplate="%{text}: %{y:.1f}<extra></extra>",
        )
    )
    seen.add(cls)
fig.add_hline(y=100, line=dict(color=COLORS["neutral"], width=1, dash="dot"))
fig.update_layout(
    title="One batch_load call returns three asset classes on one axis",
    xaxis_title="Date",
    yaxis_title="Rebased close (Jan 2 2024 = 100)",
    height=420,
    legend_title="Asset class",
)
fig.show()

# List available universes
print("Available universes:")
for name in Universe.list_universes():
    symbols = Universe.get(name)
    print(f"  {name}: {len(symbols)} symbols")

# Access a universe directly
sp500 = Universe.SP500
print(f"\nS&P 500: {len(sp500)} symbols")
print(f"First 10: {sp500[:10]}")
print(f"Last 10:  {sp500[-10:]}")

# Use with DataManager.batch_load_universe for one-line loading
# (fetches all 503 S&P 500 symbols — use a smaller slice for demo)
sp500_sample = dm.batch_load(
    symbols=sp500[:5],
    start="2024-06-01",
    end="2024-12-31",
    provider="yahoo",
)
print(
    f"S&P 500 sample: {sp500_sample.shape[0]:,} rows, {sp500_sample['symbol'].n_unique()} symbols"
)

# Custom universes for your strategy
Universe.add_custom("etf_momentum", ["SPY", "QQQ", "IWM", "EFA", "EEM", "TLT", "GLD"])
Universe.add_custom("crypto_arb", ["BTCUSDT", "ETHUSDT", "SOLUSDT", "BNBUSDT"])

print("\nCustom universes registered:")
for name in ["etf_momentum", "crypto_arb"]:
    print(f"  {name}: {Universe.get(name)}")

# Initialize storage
storage_config = StorageConfig(
    base_path=STORAGE_DIR / "hive_demo",
    compression="zstd",
    partition_granularity="month",
)
storage = HiveStorage(config=storage_config)

# DataManager with storage — enables load/update/metadata operations
dm_stored = DataManager(storage=storage)

symbols = ["AAPL", "MSFT", "GOOGL"]
stored_keys = {}
for symbol in symbols:
    key = dm_stored.load(symbol, "2023-01-01", "2024-12-31", provider="yahoo")
    stored_keys[symbol] = key
    print(f"  Stored {symbol} → key: {key}")

# List what's in storage. `storage.list_keys()` walks the on-disk layout, so it
# reports every symbol regardless of metadata-file contents.
stored_symbols = sorted(storage.list_keys())
print(f"Symbols in storage: {stored_symbols}")

# Read back with a date-range filter — only the matching month=k partitions
# touch disk, so reading 2024 from a 2-year archive halves the I/O.
aapl_2024 = storage.read(
    stored_keys["AAPL"],
    start_date=datetime(2024, 1, 1),
    end_date=datetime(2024, 12, 31),
).collect()
print(f"\nAAPL 2024 only: {len(aapl_2024)} rows (partition-pruned)")
print(f"Date range: {aapl_2024['timestamp'].min().date()} to {aapl_2024['timestamp'].max().date()}")

# Check metadata for stored symbols
for symbol in symbols:
    meta = dm_stored.get_metadata(symbol)
    if meta:
        print(f"\n{symbol}:")
        for k, v in list(meta.items())[:5]:
            print(f"  {k}: {v}")

# Ask the store what it wrote. The directory names are not addressable from outside -
# the key is encoded for filesystem safety, and each write commits into a new generation
# directory so a failed write cannot leave a half-written partition visible - so
# `partitions()` is how a caller reports the layout.
for symbol in symbols:
    parts = storage.partitions(stored_keys[symbol])
    print(f"{symbol}: {len(parts)} partitions, {sum(p.size_bytes for p in parts) / 1024:.1f} KB")

print("\nAAPL partitions (first 8):")
for part in storage.partitions(stored_keys["AAPL"])[:8]:
    print(f"  {part.label}  ({part.size_bytes / 1024:.1f} KB)")

aapl_sizes = pl.DataFrame(
    [
        {"period": part.label, "size_kb": part.size_bytes / 1024}
        for part in storage.partitions(stored_keys["AAPL"])
    ]
)

fig = go.Figure(
    go.Bar(
        x=aapl_sizes["period"].to_list(),
        y=aapl_sizes["size_kb"].to_list(),
        marker_color=COLORS["blue"],
    )
)
fig.update_layout(
    title="AAPL Hive storage: one Parquet partition per month",
    xaxis_title="Partition (year-month)",
    yaxis_title="Partition size (KB)",
    height=420,
    showlegend=False,
)
fig.show()

# update() checks what's already stored and only fetches new data
for symbol in symbols:
    key = dm_stored.update(symbol, lookback_days=7, provider="yahoo")
    print(f"  Updated {symbol} → {key}")

# Verify data is current
for symbol in symbols:
    meta = dm_stored.get_metadata(symbol)
    if meta and "last_updated" in meta:
        print(f"  {symbol} last updated: {meta['last_updated']}")

from ml4t.data.update_manager import GapDetector

# Pass `exclude_weekends=True` so Saturdays and Sundays don't count as gaps.
# The cached series here is calendar-dense (each non-trading day carries the
# prior close forward), so the detector reports no gaps. For a sparse,
# trading-days-only feed it would instead flag every missing session, including
# holidays — without an exchange calendar it cannot tell a holiday from a true
# gap, so pair it with a calendar-aware check for end-of-day pipelines.
gap_detector = GapDetector(exclude_weekends=True)

for symbol, key in stored_keys.items():
    df = storage.read(key).collect()
    gaps = gap_detector.detect_gaps(df, frequency="daily")
    if gaps:
        print(f"{symbol}: {len(gaps)} gap(s) detected")
        for gap in gaps[:3]:
            print(f"  {gap['start'].date()} to {gap['end'].date()} ({gap['size_days']} days)")
    else:
        print(f"{symbol}: No gaps (complete)")

def production_pipeline(
    universe_name: str,
    start: str,
    end: str,
    storage_path: Path,
) -> pl.DataFrame:
    """Fetch, store, validate, and assemble a stacked DataFrame for a universe.

    The same pattern drives `data/download_all.py` for every asset class —
    only the universe and provider differ.
    """
    from ml4t.data.validation import OHLCVValidator

    symbols = Universe.get(universe_name)
    print(f"Universe '{universe_name}': {len(symbols)} symbols")

    config = StorageConfig(base_path=storage_path, compression="zstd")
    store = HiveStorage(config=config)
    manager = DataManager(storage=store, enable_validation=True)

    stored = {}
    for symbol in symbols:
        stored[symbol] = manager.load(symbol, start, end, provider="yahoo")
    print(f"Fetched: {len(stored)} symbols")

    validator = OHLCVValidator(max_return_threshold=0.5)
    issues = 0
    for symbol, key in stored.items():
        df = store.read(key).collect()
        result = validator.validate(df)
        if not result.passed:
            issues += result.error_count
            print(f"  {symbol}: {result.error_count} validation issues")
    print(f"Validated: {issues} total issue(s) across {len(stored)} symbols")

    frames = [
        store.read(key).collect().with_columns(pl.lit(symbol).alias("symbol"))
        for symbol, key in stored.items()
    ]
    combined = pl.concat(frames)
    print(f"Result: {combined.shape[0]:,} rows, {combined['symbol'].n_unique()} symbols")
    return combined

# Run pipeline on a small universe
pipeline_output = production_pipeline(
    universe_name="etf_momentum",
    start="2024-01-01",
    end="2024-12-31",
    storage_path=STORAGE_DIR / "pipeline_demo",
)

pipeline_output.head()