"""Export feature snapshots + labels to Parquet with train/valid/test split."""
from __future__ import annotations

import asyncio
import datetime as dt
import json
import subprocess
import uuid
from pathlib import Path
from typing import Any

import pyarrow as pa
import pyarrow.parquet as pq
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from libs.common.logging import get_logger
from libs.common.time_utils import (
    filing_time_bucket as classify_time_bucket,
    trading_days_between,
    utc_now,
)

logger = get_logger(__name__)

MANIFEST_FILENAME = "manifest.json"


_UNIVERSE_PROFILE_MIDLARGE_LIQUID_LONG_V1 = "midlarge-liquid-long-v1"
_UNIVERSE_PROFILE_MIDPLUS_LIQUID_LONG_V1 = "midplus-liquid-long-v1"
_UNIVERSE_PROFILE_MIDWIDE_LIQUID_LONG_V1 = "midwide-liquid-long-v1"
_UNIVERSE_PROFILE_SMALLCAP_LIQUID_LONG_V1 = "smallcap-liquid-long-v1"

_UNIVERSE_PROFILES: dict[str, dict[str, Any]] = {
    _UNIVERSE_PROFILE_MIDLARGE_LIQUID_LONG_V1: {
        "market_cap_min": 2_000_000_000,
        "price_min": 15,
        "avg_dollar_volume_20d_min": 75_000_000,
        "exchange": "NYSE,NASDAQ,AMEX",
        "exclude_types": "ETF,FUND,ADR,SPAC",
    },
    _UNIVERSE_PROFILE_MIDPLUS_LIQUID_LONG_V1: {
        "market_cap_min": 1_500_000_000,
        "price_min": 12,
        "avg_dollar_volume_20d_min": 65_000_000,
        "exchange": "NYSE,NASDAQ,AMEX",
        "exclude_types": "ETF,FUND,ADR,SPAC",
    },
    _UNIVERSE_PROFILE_MIDWIDE_LIQUID_LONG_V1: {
        "market_cap_min": 1_000_000_000,
        "price_min": 10,
        "avg_dollar_volume_20d_min": 50_000_000,
        "exchange": "NYSE,NASDAQ,AMEX",
        "exclude_types": "ETF,FUND,ADR,SPAC",
    },
    _UNIVERSE_PROFILE_SMALLCAP_LIQUID_LONG_V1: {
        "market_cap_min": 500_000_000,
        "market_cap_max": 2_000_000_000,
        "price_min": 8,
        "avg_dollar_volume_20d_min": 15_000_000,
        "exchange": "NYSE,NASDAQ,AMEX",
        "exclude_types": "ETF,FUND,ADR,SPAC",
    },
}


def _get_git_commit_hash() -> str:
    """Return the current git commit hash (short), or 'unknown'."""
    try:
        result = subprocess.run(
            ["git", "rev-parse", "--short", "HEAD"],
            capture_output=True,
            text=True,
            timeout=5,
        )
        return result.stdout.strip() or "unknown"
    except Exception:
        return "unknown"


def _temporal_split(
    rows: list[dict[str, Any]],
    split_policy: str = "temporal_70_15_15",
) -> dict[str, list[dict[str, Any]]]:
    """Split rows into train/valid/test by event_date (temporal order).

    Args:
        rows: List of row dicts that must have an "event_date" field.
        split_policy: E.g. "temporal_70_15_15" → 70% train, 15% valid, 15% test.

    Returns:
        Dict with keys "train", "valid", "test".
    """
    if not rows:
        return {"train": [], "valid": [], "test": []}

    parts = split_policy.replace("temporal_", "").split("_")
    if len(parts) != 3:
        raise ValueError(f"Invalid split_policy: {split_policy}")
    train_pct, valid_pct, _ = (int(p) for p in parts)

    sorted_rows = sorted(rows, key=lambda r: r.get("event_date", ""))
    n = len(sorted_rows)
    n_train = int(n * train_pct / 100)
    n_valid = int(n * valid_pct / 100)

    return {
        "train": sorted_rows[:n_train],
        "valid": sorted_rows[n_train : n_train + n_valid],
        "test": sorted_rows[n_train + n_valid :],
    }


def _rows_to_table(rows: list[dict[str, Any]]) -> pa.Table:
    """Convert list of dicts to a PyArrow Table."""
    if not rows:
        return pa.table({})
    # Collect ALL keys from ALL rows (not just the first)
    keys: list[str] = []
    seen: set[str] = set()
    for row in rows:
        for k in row:
            if k not in seen:
                keys.append(k)
                seen.add(k)
    arrays: dict[str, list[Any]] = {k: [] for k in keys}
    for row in rows:
        for k in keys:
            arrays[k].append(row.get(k))
    return pa.table({k: pa.array(v) for k, v in arrays.items()})


async def _resolve_universe_profile(
    universe_profile: str | None,
) -> dict[str, dict[str, Any]]:
    if not universe_profile:
        return {}
    profile = _UNIVERSE_PROFILES.get(universe_profile)
    if profile is None:
        raise ValueError(f"Unknown universe_profile: {universe_profile}")

    from libs.oracle_client import ScreenerService, make_oracle_client

    async with make_oracle_client() as client:
        svc = ScreenerService(client)
        stocks = await svc.search_all_stocks(
            market_cap_min=profile["market_cap_min"],
            market_cap_max=profile.get("market_cap_max"),
            exchange=profile["exchange"],
            exclude_types=profile["exclude_types"],
            price_min=profile["price_min"],
        )
    return {
        stock.symbol.upper(): {
            "exchange_proxy": stock.exchange,
            "market_cap_proxy": stock.market_cap,
        }
        for stock in stocks
        if stock.symbol
    }


async def _resolve_symbol_metadata(
    symbols: list[str] | set[str] | None,
    *,
    concurrency: int = 16,
) -> dict[str, dict[str, Any]]:
    """Resolve exchange/market-cap metadata for explicit symbol lists.

    This keeps symbol-based exports aligned with universe-profile exports so
    backtests can still apply market-cap proxy filters.
    """
    if not symbols:
        return {}

    from libs.oracle_client import CompanyService, ScreenerService, make_oracle_client

    normalized = sorted({str(symbol).upper() for symbol in symbols if symbol})
    semaphore = asyncio.Semaphore(concurrency)
    result: dict[str, dict[str, Any]] = {}

    async with make_oracle_client() as client:
        svc = CompanyService(client)

        async def _fetch(sym: str) -> tuple[str, dict[str, Any]]:
            async with semaphore:
                try:
                    company = await svc.get_company(sym)
                except Exception as exc:
                    logger.warning(
                        "snapshot_export_symbol_metadata_failed",
                        symbol=sym,
                        error=str(exc),
                    )
                    return sym, {}
                meta = {
                    "exchange_proxy": company.exchange,
                    "market_cap_proxy": company.market_cap,
                }
                if not meta["exchange_proxy"] and not meta["market_cap_proxy"]:
                    return sym, {}
                return sym, meta

        rows = await asyncio.gather(*(_fetch(sym) for sym in normalized))
        for sym, meta in rows:
            if meta:
                result[sym] = meta

        unresolved = [sym for sym in normalized if sym not in result]
        if unresolved:
            screener = ScreenerService(client)
            try:
                stocks = await screener.search_all_stocks(
                    exchange="NYSE,NASDAQ,AMEX",
                    exclude_types="ETF,FUND,ADR,SPAC",
                )
                lookup = {
                    stock.symbol.upper(): {
                        "exchange_proxy": stock.exchange,
                        "market_cap_proxy": stock.market_cap,
                    }
                    for stock in stocks
                    if stock.symbol
                }
                for sym in unresolved:
                    meta = lookup.get(sym)
                    if meta:
                        result[sym] = meta
            except Exception as exc:
                logger.warning(
                    "snapshot_export_symbol_metadata_screener_failed",
                    error=str(exc),
                    unresolved=len(unresolved),
                )
    return result


def _parse_iso_date(raw: Any) -> dt.date | None:
    if isinstance(raw, dt.date):
        return raw
    if isinstance(raw, str):
        try:
            return dt.date.fromisoformat(raw)
        except ValueError:
            return None
    return None


async def _enrich_macro_features(rows: list[dict[str, Any]]) -> None:
    """Add VIX and HY spread from FRED as macro regime features.

    Uses forward-fill for weekends/holidays. These features enable
    macro-aware scoring: high VIX + wide HY spread = favorable PEAD regime.
    """
    if not rows:
        return
    try:
        from libs.oracle_client import FredService, make_oracle_client
        import pandas as pd

        dates = sorted({r.get("event_date", "") for r in rows if r.get("event_date")})
        if not dates:
            return
        start = dates[0]

        async with make_oracle_client() as client:
            svc = FredService(client)
            vix_resp = await svc.get_observations("VIXCLS", start=start)
            hy_resp = await svc.get_observations("BAMLH0A0HYM2", start=start)

        vix_series = pd.Series(
            {o.date: o.value for o in vix_resp.observations if o.value is not None}
        ).sort_index()
        hy_series = pd.Series(
            {o.date: o.value for o in hy_resp.observations if o.value is not None}
        ).sort_index()

        for row in rows:
            ed = row.get("event_date", "")
            if not ed:
                continue
            prior_vix = vix_series[vix_series.index <= ed]
            row["macro_vix"] = float(prior_vix.iloc[-1]) if len(prior_vix) > 0 else None
            prior_hy = hy_series[hy_series.index <= ed]
            row["macro_hy_spread"] = float(prior_hy.iloc[-1]) if len(prior_hy) > 0 else None

        enriched = sum(1 for r in rows if r.get("macro_vix") is not None)
        logger.info("macro_features_enriched", count=enriched, total=len(rows))
    except Exception as exc:
        logger.warning("macro_features_failed", error=str(exc))


def _enrich_prior_event_drift(rows: list[dict[str, Any]]) -> None:
    """Add a PIT-safe prior_event_fwd5d from the same ticker's most recent prior event.

    The feature is only populated when the prior event's 5-trading-day forward
    window is fully realized strictly before the current event date. If the most
    recent prior event has not fully realized yet, the current row gets null.
    """
    sorted_rows = sorted(
        rows,
        key=lambda r: (
            str(r.get("ticker", "")),
            str(r.get("event_date", "")),
            str(r.get("entry_date", "")),
        ),
    )
    prev_by_ticker: dict[str, tuple[float | None, dt.date | None]] = {}
    realized_on_cache: dict[dt.date, dt.date | None] = {}

    def _fwd5_realized_on(entry_date: dt.date | None) -> dt.date | None:
        if entry_date is None:
            return None
        cached = realized_on_cache.get(entry_date)
        if cached is not None or entry_date in realized_on_cache:
            return cached
        trading_days = trading_days_between(entry_date, entry_date + dt.timedelta(days=14))
        realized_on = trading_days[5] if len(trading_days) > 5 else None
        realized_on_cache[entry_date] = realized_on
        return realized_on

    for row in sorted_rows:
        ticker = str(row.get("ticker", ""))
        current_event_date = _parse_iso_date(row.get("event_date"))
        prior_value: float | None = None
        prior = prev_by_ticker.get(ticker)
        if prior is not None and current_event_date is not None:
            candidate_value, realized_on = prior
            if realized_on is not None and current_event_date > realized_on:
                prior_value = candidate_value
        row["prior_event_fwd5d"] = prior_value

        fwd5 = row.get("fwd_return_5d")
        entry_date = _parse_iso_date(row.get("entry_date"))
        prev_by_ticker[ticker] = (
            float(fwd5) if fwd5 is not None else None,
            _fwd5_realized_on(entry_date),
        )


async def _backfill_market_fields(rows: list[dict[str, Any]]) -> None:
    pending_rows = [
        row for row in rows
        if row.get("ticker")
        and (
            row.get("avg_dollar_volume_20d") is None
            or row.get("reaction_day_low") is None
            or row.get("reaction_day_high") is None
        )
    ]
    if not pending_rows:
        return

    from libs.backtest.snapshot_store import SnapshotStore
    from libs.common.config import get_settings

    symbols = sorted({str(row["ticker"]).upper() for row in pending_rows if row.get("ticker")})
    dates = [
        parsed
        for row in pending_rows
        for parsed in [_parse_iso_date(row.get("reaction_date") or row.get("event_date"))]
        if parsed is not None
    ]
    if not symbols or not dates:
        return

    settings = get_settings()
    date_range = (min(dates) - dt.timedelta(days=45), max(dates) + dt.timedelta(days=5))
    bars_by_symbol, _ = await SnapshotStore._fetch_price_data(
        symbols,
        date_range,
        settings.stock_oracle_url,
    )

    for row in pending_rows:
        ticker = str(row.get("ticker") or "").upper()
        reaction_date = _parse_iso_date(row.get("reaction_date") or row.get("event_date"))
        if not ticker or reaction_date is None:
            continue
        date_bars = bars_by_symbol.get(ticker, {})
        if reaction_date not in date_bars:
            continue
        sorted_dates = sorted(date_bars)
        idx = sorted_dates.index(reaction_date)
        event_bar = date_bars[reaction_date]
        if row.get("reaction_day_low") is None:
            row["reaction_day_low"] = event_bar.get("low")
        if row.get("reaction_day_high") is None:
            row["reaction_day_high"] = event_bar.get("high")
        if row.get("avg_dollar_volume_20d") is None and idx > 0:
            prior = sorted_dates[max(0, idx - 20) : idx]
            if prior:
                row["avg_dollar_volume_20d"] = sum(
                    float(date_bars[d]["close"]) * float(date_bars[d]["volume"])
                    for d in prior
                ) / len(prior)


async def export_dataset_snapshot(
    session: AsyncSession,
    snapshot_id: str | None,
    split_policy: str,
    output_dir: str | Path,
    feature_version: str = "market_v1",
    feature_versions: list[str] | None = None,
    label_version: str = "label-1.0.0",
    parser_version: str = "rule-1.0.0",
    symbols: list[str] | None = None,
    universe_profile: str | None = None,
    start_date: dt.date | None = None,
    end_date: dt.date | None = None,
) -> dict[str, Any]:
    """Join FeatureSnapshot + EventLabel and export to Parquet.

    Args:
        session: Async DB session.
        snapshot_id: Unique ID for this snapshot (generated if None).
        split_policy: Temporal split policy string (e.g. "temporal_70_15_15").
        output_dir: Root directory for output files.
        feature_version: Snapshot name filter (used when feature_versions is None).
        feature_versions: Merge multiple feature types (e.g. ["market_v1", "event_v1"]).
        label_version: Label version filter for EventLabel.
        parser_version: Parser version filter for Event.
        universe_profile: Optional named live screener profile to filter symbols.
        start_date: Optional inclusive lower bound on Event.event_date.
        end_date: Optional inclusive upper bound on Event.event_date.

    Returns:
        Manifest dict with metadata and row counts.
    """
    from libs.db.models import Event, EventLabel, FeatureSnapshot, SymbolMaster

    if snapshot_id is None:
        snapshot_id = str(uuid.uuid4())

    if feature_versions:
        versions = list(dict.fromkeys(feature_versions))
    else:
        versions = [feature_version]
        # The live research snapshots merge event parser features into the market rows.
        # Keep export defaults aligned so OOT snapshots exercise the same strategy fields.
        if feature_version == "market_v1":
            versions.append("event_v1")

    out_path = Path(output_dir) / snapshot_id
    out_path.mkdir(parents=True, exist_ok=True)

    profile_symbol_meta = await _resolve_universe_profile(universe_profile)
    explicit_symbols = {s.upper() for s in symbols} if symbols else None
    if profile_symbol_meta:
        profile_symbols = set(profile_symbol_meta.keys())
        explicit_symbols = profile_symbols if explicit_symbols is None else (explicit_symbols & profile_symbols)
    explicit_symbol_meta: dict[str, dict[str, Any]] = {}
    if explicit_symbols:
        missing_meta_symbols = (
            explicit_symbols - set(profile_symbol_meta.keys())
            if profile_symbol_meta
            else explicit_symbols
        )
        explicit_symbol_meta = await _resolve_symbol_metadata(missing_meta_symbols)

    # Query: JOIN feature_snapshots + event_labels via event_id
    # When merging multiple versions, query all and group by event_id
    stmt = (
        select(FeatureSnapshot, EventLabel, Event, SymbolMaster)
        .join(EventLabel, FeatureSnapshot.event_id == EventLabel.event_id)
        .join(Event, FeatureSnapshot.event_id == Event.event_id)
        .outerjoin(SymbolMaster, Event.symbol_id == SymbolMaster.symbol_id)
        .where(FeatureSnapshot.snapshot_name.in_(versions))
        .where(EventLabel.label_version == label_version)
        .where(EventLabel.label_status.in_(["ok", "truncated"]))
        .where(EventLabel.invalid_event_for_labeling.is_(False))
    )

    if explicit_symbols:
        stmt = stmt.where(
            Event.symbol_id.in_(
                select(SymbolMaster.symbol_id).where(
                    SymbolMaster.ticker.in_(sorted(explicit_symbols))
                )
            )
        )

    result = await session.execute(stmt)
    pairs = result.all()

    # Group features by event_id, merge feature_json from all versions
    event_features: dict[str, dict[str, Any]] = {}
    event_labels: dict[str, Any] = {}
    event_dates: dict[str, dt.date] = {}
    event_filed_at: dict[str, dt.datetime | None] = {}
    event_symbols: dict[str, str] = {}
    event_symbol_meta: dict[str, dict[str, Any]] = {}
    for fs, lbl, evt, sym in pairs:
        eid = fs.event_id
        if eid not in event_features:
            event_features[eid] = {}
            event_labels[eid] = lbl
            event_dates[eid] = evt.event_date
            event_filed_at[eid] = evt.filed_at_utc
            ticker = (sym.ticker if sym and sym.ticker else "").upper()
            event_symbols[eid] = ticker
            event_symbol_meta[eid] = {
                "asset_type_proxy": sym.asset_type if sym else None,
                "exchange_proxy": (sym.venue if sym else None),
            }
        event_features[eid].update(fs.feature_json)

    preliminary_rows: list[dict[str, Any]] = []
    for eid, features in event_features.items():
        ticker = event_symbols.get(eid, "")
        if explicit_symbols is not None and ticker not in explicit_symbols:
            continue
        lbl = event_labels[eid]
        row: dict[str, Any] = {
            "event_id": eid,
            "ticker": ticker,
            "snapshot_name": "+".join(versions),
            "snapshot_version": "1.0.0",
            **features,
            "entry_convention": lbl.entry_convention,
            "reaction_date": lbl.reaction_date.isoformat() if lbl.reaction_date else None,
            "entry_date": lbl.entry_date.isoformat() if lbl.entry_date else None,
            "entry_price": float(lbl.entry_price) if lbl.entry_price else None,
            "fwd_return_1d": float(lbl.fwd_return_1d) if lbl.fwd_return_1d else None,
            "fwd_return_3d": float(lbl.fwd_return_3d) if lbl.fwd_return_3d else None,
            "fwd_return_5d": float(lbl.fwd_return_5d) if lbl.fwd_return_5d else None,
            "hit_pos_1r_within_3d": lbl.hit_pos_1r_within_3d,
            "hit_neg_1r_within_3d": lbl.hit_neg_1r_within_3d,
            "close_up_after_3d": lbl.close_up_after_3d,
            "close_up_after_5d": lbl.close_up_after_5d,
            "mfe_3d": float(lbl.mfe_3d) if lbl.mfe_3d else None,
            "mae_3d": float(lbl.mae_3d) if lbl.mae_3d else None,
            "mfe_5d": float(lbl.mfe_5d) if lbl.mfe_5d else None,
            "mae_5d": float(lbl.mae_5d) if lbl.mae_5d else None,
            "fwd_return_10d": float(lbl.fwd_return_10d) if lbl.fwd_return_10d else None,
            "fwd_return_20d": float(lbl.fwd_return_20d) if lbl.fwd_return_20d else None,
            "mfe_10d": float(lbl.mfe_10d) if lbl.mfe_10d else None,
            "mae_10d": float(lbl.mae_10d) if lbl.mae_10d else None,
            "mfe_20d": float(lbl.mfe_20d) if lbl.mfe_20d else None,
            "mae_20d": float(lbl.mae_20d) if lbl.mae_20d else None,
            "label_status": lbl.label_status,
            "label_version": lbl.label_version,
        }
        # Use DB Event.event_date as authoritative source
        ed = event_dates.get(eid)
        row["event_date"] = ed.isoformat() if ed else row.get("event_date", "")
        filed_at = event_filed_at.get(eid)
        if (
            isinstance(filed_at, dt.datetime)
            and str(row.get("filing_time_bucket", "unknown")).lower() == "unknown"
        ):
            row["filing_time_bucket"] = classify_time_bucket(filed_at)

        if start_date is not None and ed is not None and ed < start_date:
            continue
        if end_date is not None and ed is not None and ed > end_date:
            continue

        symbol_meta = dict(event_symbol_meta.get(eid, {}))
        if ticker and ticker in explicit_symbol_meta:
            symbol_meta.update(explicit_symbol_meta[ticker])
        if ticker and ticker in profile_symbol_meta:
            symbol_meta.update(profile_symbol_meta[ticker])
        row.update(symbol_meta)
        preliminary_rows.append(row)

    await _backfill_market_fields(preliminary_rows)

    # Compute cross-event momentum: prior same-ticker event's 5d forward return
    _enrich_prior_event_drift(preliminary_rows)

    # Enrich with macro regime features (VIX, HY spread)
    await _enrich_macro_features(preliminary_rows)

    rows: list[dict[str, Any]] = []
    for row in preliminary_rows:
        market_cap_proxy = row.get("market_cap_proxy")
        avg_dollar_volume_20d = row.get("avg_dollar_volume_20d")
        profile = _UNIVERSE_PROFILES.get(universe_profile) if universe_profile else None
        if profile is not None:
            if market_cap_proxy is None or float(market_cap_proxy) < float(profile["market_cap_min"]):
                continue
            if "market_cap_max" in profile and float(market_cap_proxy) > float(profile["market_cap_max"]):
                continue
            if (
                avg_dollar_volume_20d is None
                or float(avg_dollar_volume_20d) < float(profile["avg_dollar_volume_20d_min"])
            ):
                continue
        rows.append(row)

    logger.info("snapshot_export_rows", snapshot_id=snapshot_id, total=len(rows))

    splits = _temporal_split(rows, split_policy)
    row_counts: dict[str, int] = {}

    for split_name, split_rows in splits.items():
        parquet_path = out_path / f"{split_name}.parquet"
        table = _rows_to_table(split_rows)
        pq.write_table(table, str(parquet_path))
        row_counts[split_name] = len(split_rows)
        logger.info("split_written", split=split_name, rows=len(split_rows), path=str(parquet_path))

    manifest: dict[str, Any] = {
        "snapshot_id": snapshot_id,
        "created_at_utc": utc_now().isoformat(),
        "code_commit_hash": _get_git_commit_hash(),
        "feature_version": "+".join(versions),
        "parser_version": parser_version,
        "label_version": label_version,
        "split_policy": split_policy,
        "universe_profile": universe_profile,
        "start_date": start_date.isoformat() if start_date else None,
        "end_date": end_date.isoformat() if end_date else None,
        "row_counts": row_counts,
        "total_rows": len(rows),
        "output_dir": str(out_path),
    }

    manifest_path = out_path / MANIFEST_FILENAME
    manifest_path.write_text(json.dumps(manifest, indent=2))
    logger.info("manifest_written", path=str(manifest_path), snapshot_id=snapshot_id)

    return manifest