"""Real point-in-time multi-rebalance backtest engine.

True PIT honesty requires a `score_fn(score_by_symbol, as_of)` that returns the
combined scores *as they were known at `as_of`*. The default (current board) has
no historical factor vintages, so it is labeled non-PIT (`leakage_guard=False`).
When a PIT scorer is supplied, `leakage_guard=True`.

At each rebalance date the engine:
  - resolves the combined score as-of that date (price data is genuinely
    point-in-time w.r.t. price through `_latest_close`),
  - marks the current portfolio to market,
  - re-allocates the 50/20/30 dividend buckets over the current value,
  - reconciles holdings (sells names that leave, buys/upsizes names that enter),
so `rebalances` reflects real re-trades, not a single allocate-once.
"""

from __future__ import annotations

import datetime as dt
from dataclasses import dataclass, field
from typing import Callable, Optional

from .simulation import (
    allocate_capital, load_price_snapshot,
)

# score_fn contract: (symbols: list[str], as_of: str|None) -> {sym: meta dict}
ScoreFn = Callable[[list[str], Optional[str]], dict]


def _bar_date(s: Optional[str]) -> dt.date:
    if not s:
        return dt.date.min
    return dt.date.fromisoformat(str(s)[:10])


def _bars_up_to(series: dict, sym: str, date: dt.date) -> list:
    bars = series.get(sym, {}).get("bars", [])
    return [b for b in bars if _bar_date(b.get("date")) <= date]


def _latest_close(series: dict, sym: str, date: dt.date) -> Optional[float]:
    bars = _bars_up_to(series, sym, date)
    if bars:
        return float(bars[-1]["adjusted_close"])
    return None


def momentum_at(series: dict, sym: str, date: dt.date,
                lookback_days: int = 252, skip_days: int = 21) -> Optional[float]:
    """True 12-1 momentum as of `date`: close at ~1 month ago / close ~12 months
    before that, minus 1 — skipping the most recent month to avoid short-term
    reversal. Only uses bars known up to `date` (no lookahead)."""
    bars = _bars_up_to(series, sym, date)
    if len(bars) < lookback_days + skip_days + 1:
        return None
    try:
        ref = float(bars[-1 - skip_days]["adjusted_close"])   # ~1m ago
        base = float(bars[-1 - skip_days - lookback_days]["adjusted_close"])
    except (KeyError, TypeError, ValueError, IndexError):
        return None
    if ref <= 0 or base <= 0:
        return None
    return round((ref / base) - 1.0, 4)


@dataclass
class BacktestResult:
    start: str
    end: str
    capital: float
    final_value: float = 0.0
    price_pnl: float = 0.0
    dividend_income: float = 0.0
    net_return: float = 0.0
    trades: int = 0
    rebalances: int = 0          # actual number of re-allocations executed
    planned_rebalances: int = 0  # number of rebalance windows
    holdings: dict = field(default_factory=dict)  # final {sym: qty}
    leakage_guard: bool = False  # True only when a PIT score_fn was supplied
    dividend_method: str = "final_holdings_yield_proxy"  # which dividend model

    def to_dict(self) -> dict:
        return {
            "start": self.start, "end": self.end, "capital": self.capital,
            "final_value": round(self.final_value, 2),
            "price_pnl": round(self.price_pnl, 2),
            "dividend_income": round(self.dividend_income, 2),
            "dividend_method": self.dividend_method,
            "net_return": round(self.net_return, 4),
            "trades": self.trades,
            "rebalances": self.rebalances,
            "planned_rebalances": self.planned_rebalances,
            "holdings": self.holdings,
            "leakage_guard": self.leakage_guard,
        }


class BacktestError(Exception):
    pass


def _rebalance_dates(start: str, end: str, freq: str = "monthly") -> list[str]:
    s = dt.date.fromisoformat(start)
    e = dt.date.fromisoformat(end)
    if s >= e:
        raise BacktestError("end must be after start")
    dates = []
    cur = s
    if freq == "monthly":
        while cur <= e:
            dates.append(cur.isoformat())
            y, m = (cur.year + 1, 1) if cur.month == 12 else (cur.year, cur.month + 1)
            cur = dt.date(y, m, 1)
    elif freq == "quarterly":
        while cur <= e:
            dates.append(cur.isoformat())
            q = (cur.month - 1) // 3 + 1
            if q == 4:
                cur = dt.date(cur.year + 1, 1, 1)
            else:
                cur = dt.date(cur.year, q * 3 + 1, 1)
    else:
        raise BacktestError(f"unsupported freq {freq}")
    return dates


def _resolve_scores(score_fn, syms: list[str], as_of: Optional[str]) -> tuple[dict, bool]:
    """Return (score_by_symbol, is_pit).

    A supplied score_fn marks ``leakage_guard`` ONLY when the returned scores
    carry a ``pit_meta`` proving they were built point-in-time:
      - ``pit_meta = {"pit": true, ...}``  -> leakage_guard = True
      - ``pit_meta`` present but ``pit=false`` (blocked/partial/fallback) -> False
      - no ``pit_meta`` at all (an arbitrary caller-provided fn) -> False

    This replaces the old behaviour that set leakage_guard=True for ANY supplied
    callable, which could not distinguish a genuine PIT scorer from one that
    silently reused the current board.
    """
    if score_fn is None:
        from .dashboard import default_scores
        return default_scores(syms) or {}, False
    out = score_fn(syms, as_of) or {}
    if not out:
        return out, False
    # is_pit: the scores themselves assert PIT integrity via pit_meta.
    any_pit = any(
        isinstance(m, dict) and isinstance(m.get("pit_meta"), dict)
        and bool(m.get("pit_meta", {}).get("pit"))
        for m in out.values()
    )
    return out, any_pit


def _candidates_at(series: dict, syms: list[str], date: dt.date,
                   score_by_symbol: dict) -> list:
    out = []
    for sym in syms:
        price = _latest_close(series, sym, date)
        if not price or price <= 0:
            continue
        meta = score_by_symbol.get(sym, {})
        from .simulation import Candidate
        out.append(Candidate(
            symbol=sym, price=price,
            combined_score=float(meta.get("combined", 0.0)),
            is_dividend=bool(meta.get("is_dividend", False)),
            dividend_yield=float(meta.get("dividend_yield") or 0.0),
        ))
    return out


def run_backtest(
    start: str, end: str,
    capital: float = 1_000_000,
    rebalance_freq: str = "monthly",
    score_fn: Optional[ScoreFn] = None,
    symbols: Optional[list[str]] = None,
    dividend_ledger=None,
) -> BacktestResult:
    """Run a multi-rebalance backtest over [start, end].

    `score_fn(symbols, as_of)` returns {sym: {combined, is_dividend,
    dividend_yield}} as of `as_of`. Default: current board (static, non-PIT ->
    leakage_guard=False). A supplied score_fn sets leakage_guard=True only when
    its scores assert pit_meta.

    `dividend_ledger` (optional DividendLedger) replaces the final-holdings
    yield proxy: dividends are credited as ``per_share * qty`` from the ledger's
    dated per-symbol entries instead of ``final_qty * px * yield%``. When a
    symbol has no ledger entry it earns no dividend (fail closed, no
    fabrication). When `dividend_ledger` is None the legacy proxy is used and
    labelled as such.
    """
    from .dividend_ledger import DividendLedger, credit_dividends
    from .dividend_ledger import DividendLedgerError
    _ledger = dividend_ledger if dividend_ledger is not None else None
    series = load_price_snapshot()
    if not series:
        raise BacktestError("no price snapshot")
    syms = symbols or list(series.keys())

    dates = _rebalance_dates(start, end, rebalance_freq)
    result = BacktestResult(start=start, end=end, capital=capital)
    result.planned_rebalances = len(dates)

    cash = capital
    holdings: dict[str, int] = {}   # sym -> qty
    total_dividend = 0.0
    trades = 0
    actual_rebalances = 0
    leakage_guard = False
    score_by_symbol: dict = {}

    for d_iso in dates:
        d = dt.date.fromisoformat(d_iso)
        score_by_symbol, is_pit = _resolve_scores(score_fn, syms, d.isoformat())
        leakage_guard = leakage_guard or is_pit
        cands = _candidates_at(series, syms, d, score_by_symbol)
        if not cands:
            continue
        # current portfolio value at d
        port_val = cash + sum(
            qty * (px or 0.0)
            for sym, qty in holdings.items()
            if (px := _latest_close(series, sym, d)) is not None
        )
        alloc = allocate_capital(port_val, cands)
        target = {o.symbol: o.qty for o in alloc.orders}

        # sell holdings not in the new target
        for sym, qty in list(holdings.items()):
            tgt = target.get(sym, 0)
            if qty > tgt:
                px = _latest_close(series, sym, d)
                if px is None:
                    continue
                cash += (qty - tgt) * px
                holdings[sym] = tgt
                trades += 1
        # buy / upsize to target
        for o in alloc.orders:
            cur = holdings.get(o.symbol, 0)
            if o.qty > cur:
                cash -= (o.qty - cur) * o.price
                holdings[o.symbol] = o.qty
                trades += 1
        actual_rebalances += 1
        # drop zero-holding entries
        holdings = {k: v for k, v in holdings.items() if v > 0}

    _e = dt.date.fromisoformat(end)
    ending_market_value = 0.0
    dividend_method = "final_holdings_yield_proxy"
    for sym, qty in holdings.items():
        px = _latest_close(series, sym, _e)
        if px:
            ending_market_value += qty * px
            if _ledger is not None:
                # ledger-driven: per_share * qty from dated entries (fail closed
                # if no entry -> no credit).
                try:
                    credit = credit_dividends(_ledger, {sym: float(qty)}, _e)
                except DividendLedgerError:
                    credit = 0.0
                total_dividend += credit
                dividend_method = "dated_ledger" if not _ledger_has_estimate(_ledger, sym) else "dps_annual_proxy"
            else:
                # legacy proxy: yield% * current market value (honest-flagged)
                meta = score_by_symbol.get(sym, {})
                yield_pct = float(meta.get("dividend_yield") or 0.0) / 100.0
                total_dividend += qty * px * yield_pct

    ending_equity_before_dividend = cash + ending_market_value
    final_value = ending_equity_before_dividend + total_dividend

    result.holdings = holdings
    result.rebalances = actual_rebalances
    result.final_value = final_value
    result.dividend_income = total_dividend
    result.price_pnl = ending_equity_before_dividend - capital
    result.net_return = (final_value - capital) / capital if capital else 0.0
    result.trades = trades
    result.leakage_guard = leakage_guard
    # record which dividend model produced `dividend_income`
    result.dividend_method = dividend_method
    return result


def _ledger_has_estimate(ledger, symbol: str) -> bool:
    """True if any ledger entry for `symbol` is a DPS annual proxy estimate."""
    for e in ledger.entries(symbol):
        if e.get("estimate"):
            return True
    return False


def total_investable(cands: list) -> float:
    return sum(c.price for c in cands if c.symbol)