"""Point-in-time snapshot data loader for the backtester.

The public interface is:
  store = SnapshotStore.load(snapshot_dir, split_name, oracle_url, db_dsn)
  rows = store.get_candidates_for_date(date)
  bar  = store.get_bar(symbol, date)
  macro = store.get_macro_for_date(date)

For tests, inject data directly:
  store = SnapshotStore(candidates_by_exec_date=..., bars_by_symbol_date=..., macro_by_date=...)
"""
from __future__ import annotations

import asyncio
import datetime as dt
from pathlib import Path
from typing import Any

import pyarrow.parquet as pq

from libs.common.logging import get_logger

logger = get_logger(__name__)


class SnapshotStore:
    """In-memory cache of point-in-time backtest data."""

    def __init__(
        self,
        candidates_by_exec_date: dict[dt.date, list[dict[str, Any]]],
        bars_by_symbol_date: dict[str, dict[dt.date, dict[str, Any]]],
        macro_by_date: dict[dt.date, dict[str, Any]] | None = None,
    ) -> None:
        self._candidates = {
            date: list(rows)
            for date, rows in candidates_by_exec_date.items()
        }
        self._candidates_by_reaction_date = self._build_reaction_index(self._candidates)
        self._bars = bars_by_symbol_date
        self._macro = macro_by_date or {}

    # ------------------------------------------------------------------
    # Public query interface
    # ------------------------------------------------------------------

    def get_candidates_for_date(self, date: dt.date) -> list[dict[str, Any]]:
        """Return candidates where execution_date == date. No look-ahead."""
        return list(self._candidates.get(date, []))

    def get_candidates_for_reaction_date(self, date: dt.date) -> list[dict[str, Any]]:
        """Return candidates where reaction_date == date. No look-ahead."""
        return list(self._candidates_by_reaction_date.get(date, []))

    def get_bar(self, symbol: str, date: dt.date) -> dict[str, Any] | None:
        """Return OHLCV bar for symbol on date, or None if missing."""
        sym_bars = self._bars.get(symbol)
        if sym_bars is None:
            return None
        return sym_bars.get(date)

    def get_macro_for_date(self, date: dt.date) -> dict[str, Any]:
        """Return macro observations for date (empty dict if none)."""
        return dict(self._macro.get(date, {}))

    def all_execution_dates(self) -> list[dt.date]:
        """Sorted list of dates that have at least one candidate."""
        return sorted(self._candidates.keys())

    def all_reaction_dates(self) -> list[dt.date]:
        """Sorted list of dates that have at least one reaction-date candidate."""
        return sorted(self._candidates_by_reaction_date.keys())

    def all_trading_days(self, include_reaction_dates: bool = False) -> list[dt.date]:
        """All NYSE trading days from first to last execution date (inclusive).

        Use this to drive the simulation loop so stop/target/time exits are
        checked on every trading day, not just candidate days.
        """
        from libs.backtest.calendar import get_trading_days

        dates = set(self.all_execution_dates())
        if include_reaction_dates:
            dates.update(self.all_reaction_dates())
        if not dates:
            return []
        ordered = sorted(dates)
        return get_trading_days(ordered[0], ordered[-1])

    # ------------------------------------------------------------------
    # Factory: load from Parquet + DB + Oracle
    # ------------------------------------------------------------------

    @classmethod
    def load(
        cls,
        snapshot_dir: str | Path,
        split_name: str,
        oracle_url: str,
        db_dsn: str,
        scoring_fn: Any | None = None,
    ) -> "SnapshotStore":
        """Synchronous factory. Internally uses asyncio.run() to prefetch data.

        scoring_fn: optional callable(row_dict) -> float to override default scoring.
        Raises RuntimeError if called from within a running event loop.
        """
        try:
            loop = asyncio.get_running_loop()
        except RuntimeError:
            loop = None

        if loop is not None and loop.is_running():
            raise RuntimeError(
                "SnapshotStore.load() cannot be called from a running event loop. "
                "Use SnapshotStore._async_load() directly in async contexts."
            )

        data = asyncio.run(
            cls._async_load(Path(snapshot_dir), split_name, oracle_url, db_dsn, scoring_fn)
        )
        return cls(**data)

    @classmethod
    async def _async_load(
        cls,
        snapshot_dir: Path,
        split_name: str,
        oracle_url: str,
        db_dsn: str,
        scoring_fn: Any | None = None,
    ) -> dict[str, Any]:
        """Async data loading pipeline. Returns kwargs dict for __init__."""
        parquet_path = snapshot_dir / f"{split_name}.parquet"
        if not parquet_path.exists():
            raise FileNotFoundError(f"Parquet file not found: {parquet_path}")

        # Step 1: Read Parquet
        logger.info("snapshot_store_reading_parquet", path=str(parquet_path))
        table = pq.read_table(str(parquet_path))
        rows: list[dict[str, Any]] = table.to_pydict()
        # Convert column-oriented dict to list of row dicts
        num_rows = table.num_rows
        col_names = list(rows.keys())
        row_list: list[dict[str, Any]] = [
            {col: rows[col][i] for col in col_names} for i in range(num_rows)
        ]
        logger.info("snapshot_store_rows_loaded", count=num_rows)

        # Collect event_ids for DB lookup
        event_ids = [str(r.get("event_id", "")) for r in row_list]

        # Step 2–6: DB + Oracle enrichment
        event_meta = await cls._fetch_event_metadata(event_ids, db_dsn)
        unique_symbols = list({m.get("ticker", "") for m in event_meta.values() if m.get("ticker")})

        date_range = cls._compute_date_range(row_list)
        bars_by_symbol, avg_dvol = await cls._fetch_price_data(
            unique_symbols, date_range, oracle_url
        )
        sectors = await cls._fetch_sectors(unique_symbols, oracle_url)
        macro_by_date = await cls._fetch_macro(date_range, db_dsn)

        # Fetch SPY bars for macro regime filter (SMA computation)
        spy_macro = await cls._fetch_spy_macro(date_range, oracle_url)
        for d, spy_data in spy_macro.items():
            macro_by_date.setdefault(d, {}).update(spy_data)

        # Step 7: Build candidates_by_exec_date
        candidates_by_exec_date: dict[dt.date, list[dict[str, Any]]] = {}
        for row in row_list:
            eid = str(row.get("event_id", ""))
            meta = event_meta.get(eid, {})
            ticker = meta.get("ticker")
            if not ticker:
                logger.debug("snapshot_store_skip_no_ticker", event_id=eid)
                continue

            # Map entry_date → execution_date at this boundary
            raw_exec = row.get("entry_date") or row.get("execution_date")
            if raw_exec is None:
                continue
            if isinstance(raw_exec, str):
                exec_date = dt.date.fromisoformat(raw_exec)
            elif isinstance(raw_exec, dt.date):
                exec_date = raw_exec
            else:
                continue

            enriched = dict(row)
            enriched["execution_date"] = exec_date
            enriched["symbol"] = ticker
            enriched["issuer_id"] = meta.get("issuer_id")
            enriched["event_date"] = meta.get("event_date")
            enriched["event_type"] = meta.get("event_type", "")
            enriched["event_timestamp"] = meta.get("event_timestamp")
            enriched["avg_dollar_volume"] = avg_dvol.get(ticker, 0.0)
            enriched["sector"] = sectors.get(ticker, "UNKNOWN")

            # Map Parquet-specific columns to canonical backtest names
            # event_close (reaction-day close) → entry_price_est baseline
            if "entry_price_est" not in enriched and "event_close" in enriched:
                enriched["entry_price_est"] = enriched["event_close"]
            # score: use existing column or compute from market features
            if "score" not in enriched or enriched.get("score") is None:
                if scoring_fn is not None:
                    enriched["score"] = scoring_fn(enriched)
                else:
                    from libs.backtest.scoring import compute_entry_score

                    enriched["score"] = compute_entry_score(enriched)

            candidates_by_exec_date.setdefault(exec_date, []).append(enriched)

        # Build bars_by_symbol_date: symbol -> date -> bar dict
        bars_by_symbol_date: dict[str, dict[dt.date, dict[str, Any]]] = {}
        for sym, date_bars in bars_by_symbol.items():
            bars_by_symbol_date[sym] = date_bars

        logger.info(
            "snapshot_store_built",
            exec_dates=len(candidates_by_exec_date),
            symbols=len(bars_by_symbol_date),
        )
        return {
            "candidates_by_exec_date": candidates_by_exec_date,
            "bars_by_symbol_date": bars_by_symbol_date,
            "macro_by_date": macro_by_date,
        }

    # ------------------------------------------------------------------
    # Internal async helpers
    # ------------------------------------------------------------------

    @staticmethod
    async def _fetch_event_metadata(
        event_ids: list[str],
        db_dsn: str,
    ) -> dict[str, dict[str, Any]]:
        """Batch-query Event + SymbolMaster for event metadata."""
        if not event_ids:
            return {}
        try:
            from sqlalchemy import select
            from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine

            from libs.db.models import Event, SymbolMaster

            engine = create_async_engine(db_dsn, echo=False)
            async_session = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)

            result: dict[str, dict[str, Any]] = {}
            async with async_session() as session:
                # Single IN clause for all event_ids
                stmt = (
                    select(Event, SymbolMaster)
                    .outerjoin(SymbolMaster, Event.symbol_id == SymbolMaster.symbol_id)
                    .where(Event.event_id.in_(event_ids))
                )
                rows = (await session.execute(stmt)).all()
                _UTC = __import__("zoneinfo").ZoneInfo("UTC")
                for event, sym in rows:
                    # Use filed_at_utc if available; fallback to filing_date + 21:00 UTC
                    # (transparent enrichment in the loader — not silent substitution in selector)
                    ts = event.filed_at_utc
                    if ts is None and event.event_date is not None:
                        ts = dt.datetime.combine(
                            event.event_date, dt.time(21, 0), tzinfo=_UTC
                        )
                    result[event.event_id] = {
                        "issuer_id": event.issuer_id,
                        "event_date": event.event_date,
                        "event_type": event.event_type,
                        "event_timestamp": ts,
                        "ticker": sym.ticker if sym else None,
                    }
            await engine.dispose()
            return result
        except Exception as exc:
            logger.warning("snapshot_store_db_fetch_failed", error=str(exc))
            return {}

    @staticmethod
    async def _fetch_price_data(
        symbols: list[str],
        date_range: tuple[dt.date, dt.date] | None,
        oracle_url: str,
        concurrency: int = 16,
    ) -> tuple[dict[str, dict[dt.date, dict[str, Any]]], dict[str, float]]:
        """Fetch daily OHLCV bars and compute avg_dollar_volume per symbol."""
        if not symbols or date_range is None:
            return {}, {}
        try:
            from libs.oracle_client import OracleClient, PriceService

            start_str = date_range[0].isoformat()
            end_str = date_range[1].isoformat()

            bars_by_symbol: dict[str, dict[dt.date, dict[str, Any]]] = {}
            avg_dvol: dict[str, float] = {}
            semaphore = asyncio.Semaphore(concurrency)

            async with OracleClient(base_url=oracle_url) as client:
                svc = PriceService(client)
                async def _fetch_symbol(sym: str) -> tuple[str, dict[dt.date, dict[str, Any]], float]:
                    async with semaphore:
                        try:
                            resp = await svc.get_daily_bars(sym, start=start_str, end=end_str)
                            date_bars: dict[dt.date, dict[str, Any]] = {}
                            dollar_vols: list[float] = []
                            for bar in resp.bars:
                                d = dt.date.fromisoformat(bar.date)
                                b = {
                                    "date": d,
                                    "open": bar.open,
                                    "high": bar.high,
                                    "low": bar.low,
                                    "close": bar.close,
                                    "volume": bar.volume,
                                }
                                date_bars[d] = b
                                dollar_vols.append(bar.close * bar.volume)
                            if dollar_vols:
                                last_20 = dollar_vols[-20:]
                                mean_dvol = sum(last_20) / len(last_20)
                            else:
                                mean_dvol = 0.0
                            return sym, date_bars, mean_dvol
                        except Exception as sym_exc:
                            logger.warning(
                                "snapshot_store_price_fetch_failed",
                                symbol=sym,
                                error=str(sym_exc),
                            )
                            return sym, {}, 0.0

                results = await asyncio.gather(*(_fetch_symbol(sym) for sym in symbols))
                for sym, date_bars, mean_dvol in results:
                    bars_by_symbol[sym] = date_bars
                    avg_dvol[sym] = mean_dvol
            return bars_by_symbol, avg_dvol
        except Exception as exc:
            logger.warning("snapshot_store_oracle_failed", error=str(exc))
            return {}, {}

    @staticmethod
    async def _fetch_sectors(
        symbols: list[str],
        oracle_url: str,
        concurrency: int = 16,
    ) -> dict[str, str]:
        """Fetch company sector for each symbol. Default 'UNKNOWN' if unavailable."""
        if not symbols:
            return {}
        result: dict[str, str] = {}
        semaphore = asyncio.Semaphore(concurrency)
        try:
            from libs.oracle_client import CompanyService, OracleClient

            async with OracleClient(base_url=oracle_url) as client:
                company_svc = CompanyService(client)
                async def _fetch_sector(sym: str) -> tuple[str, str]:
                    async with semaphore:
                        try:
                            info = await company_svc.get_company(sym)
                            return sym, info.sector or "UNKNOWN"
                        except Exception:
                            return sym, "UNKNOWN"

                sector_results = await asyncio.gather(*(_fetch_sector(sym) for sym in symbols))
                for sym, sector in sector_results:
                    result[sym] = sector
        except Exception as exc:
            logger.warning("snapshot_store_sector_fetch_failed", error=str(exc))
        # Default all remaining to UNKNOWN
        for sym in symbols:
            result.setdefault(sym, "UNKNOWN")
        return result

    @staticmethod
    async def _fetch_macro(
        date_range: tuple[dt.date, dt.date] | None,
        db_dsn: str,
    ) -> dict[dt.date, dict[str, Any]]:
        """Load MacroObservation regime data from DB."""
        if date_range is None:
            return {}
        try:
            from sqlalchemy import select
            from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine

            from libs.db.models import MacroObservation

            engine = create_async_engine(db_dsn, echo=False)
            async_session = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)

            result: dict[dt.date, dict[str, Any]] = {}
            async with async_session() as session:
                stmt = select(MacroObservation).where(
                    MacroObservation.observation_date >= date_range[0],
                    MacroObservation.observation_date <= date_range[1],
                )
                rows = (await session.execute(stmt)).scalars().all()
                for obs in rows:
                    d = obs.observation_date
                    result.setdefault(d, {})[obs.series_id] = (
                        float(obs.value) if obs.value is not None else None
                    )
            await engine.dispose()
            return result
        except Exception as exc:
            logger.warning("snapshot_store_macro_fetch_failed", error=str(exc))
            return {}

    @staticmethod
    async def _fetch_spy_macro(
        date_range: tuple[dt.date, dt.date] | None,
        oracle_url: str,
        sma_period: int = 20,
    ) -> dict[dt.date, dict[str, Any]]:
        """Fetch SPY daily bars and compute SMA for macro regime filtering.

        Returns dict: date -> {"spy_close": float, "spy_sma_20": float|None}.
        SMA is None for the first (sma_period - 1) bars.
        """
        if date_range is None:
            return {}
        try:
            from libs.oracle_client import OracleClient, PriceService

            # Extend start date back by sma_period trading days for SMA warm-up
            warmup_days = sma_period * 2  # calendar days (conservative buffer)
            extended_start = date_range[0] - dt.timedelta(days=warmup_days)

            async with OracleClient(base_url=oracle_url) as client:
                svc = PriceService(client)
                resp = await svc.get_daily_bars(
                    "SPY", start=extended_start.isoformat(), end=date_range[1].isoformat()
                )

            # Sort bars by date
            sorted_bars = sorted(resp.bars, key=lambda b: b.date)
            closes: list[tuple[dt.date, float]] = [
                (dt.date.fromisoformat(b.date), float(b.close)) for b in sorted_bars
            ]

            result: dict[dt.date, dict[str, Any]] = {}
            for i, (d, close) in enumerate(closes):
                sma = None
                if i >= sma_period - 1:
                    window = [c for _, c in closes[i - sma_period + 1 : i + 1]]
                    sma = sum(window) / len(window)
                # Only store data within the actual date range
                if d >= date_range[0]:
                    result[d] = {"spy_close": close, "spy_sma_20": sma}

            logger.info(
                "snapshot_store_spy_macro_loaded",
                bars=len(closes),
                dates_with_sma=sum(1 for v in result.values() if v.get("spy_sma_20") is not None),
            )
            return result
        except Exception as exc:
            logger.warning("snapshot_store_spy_macro_failed", error=str(exc))
            return {}

    @staticmethod
    def _compute_date_range(
        rows: list[dict[str, Any]],
    ) -> tuple[dt.date, dt.date] | None:
        """Compute (min_date, max_date) from execution and reaction-date columns."""
        dates: list[dt.date] = []
        for r in rows:
            for raw in (
                r.get("entry_date"),
                r.get("execution_date"),
                r.get("reaction_date"),
            ):
                if raw is None:
                    continue
                if isinstance(raw, str):
                    try:
                        dates.append(dt.date.fromisoformat(raw))
                    except ValueError:
                        pass
                elif isinstance(raw, dt.date):
                    dates.append(raw)
        if not dates:
            return None
        return min(dates), max(dates)

    @staticmethod
    def _build_reaction_index(
        candidates_by_exec_date: dict[dt.date, list[dict[str, Any]]],
    ) -> dict[dt.date, list[dict[str, Any]]]:
        reaction_index: dict[dt.date, list[dict[str, Any]]] = {}
        for rows in candidates_by_exec_date.values():
            for row in rows:
                reaction_date = SnapshotStore._normalize_date(row.get("reaction_date"))
                if reaction_date is None:
                    continue
                reaction_index.setdefault(reaction_date, []).append(row)
        return reaction_index

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