"""Research probe for simple ex-dividend capture ideas.

Data source:
  - yfinance historical dividend actions (free, easy to cache locally)

Scope:
  - Uses the current backtest universe from a config's SnapshotStore.
  - Tests simple overnight dividend-capture rules:
      buy previous close -> sell ex-dividend open/close.
  - Important: SnapshotStore bars are already dividend-adjusted, so the
    price return across the ex-date already embeds the dividend effect.
    This probe must not add the cash dividend again.
"""

from __future__ import annotations

import argparse
import datetime as dt
import json
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import pandas as pd
import yfinance as yf

from apps.backtester.run import _build_merged_snapshot_store, load_manifest, resolve_config
from libs.backtest.domain import DailyPortfolioState
from libs.backtest.metrics import (
    compute_max_drawdown_pct,
    compute_sharpe_ratio,
    compute_total_return_pct,
)
from libs.common.logging import configure_logging


@dataclass(frozen=True)
class DividendCaptureSpec:
    name: str
    min_dividend_yield_pct: float
    max_dividend_yield_pct: float | None
    min_avg_dollar_volume: float
    max_positions: int
    exit_mode: str  # ex_open | ex_close


def _parse_date(value: str, *, is_end: bool = False) -> dt.date:
    parts = value.split("-")
    if len(parts) == 1 and len(value) == 4 and value.isdigit():
        year = int(value)
        return dt.date(year, 12, 31) if is_end else dt.date(year, 1, 1)
    if len(parts) == 2 and all(part.isdigit() for part in parts):
        year = int(parts[0])
        month = int(parts[1])
        if is_end:
            next_month = dt.date(year + (month // 12), (month % 12) + 1, 1)
            return next_month - dt.timedelta(days=1)
        return dt.date(year, month, 1)
    return dt.date.fromisoformat(value)


def _cache_path(cache_dir: Path, start_date: dt.date, end_date: dt.date) -> Path:
    return cache_dir / f"dividends_{start_date.isoformat()}_{end_date.isoformat()}.json"


def _download_dividend_actions(
    *,
    symbols: list[str],
    start_date: dt.date,
    end_date: dt.date,
    cache_dir: Path,
) -> dict[str, list[dict[str, Any]]]:
    cache_dir.mkdir(parents=True, exist_ok=True)
    path = _cache_path(cache_dir, start_date, end_date)
    if path.exists():
        return json.loads(path.read_text())

    results: dict[str, list[dict[str, Any]]] = {}
    batch_size = 100
    start_str = (start_date - dt.timedelta(days=14)).isoformat()
    end_str = (end_date + dt.timedelta(days=1)).isoformat()
    for offset in range(0, len(symbols), batch_size):
        batch = symbols[offset : offset + batch_size]
        frame = yf.download(
            batch,
            start=start_str,
            end=end_str,
            auto_adjust=False,
            actions=True,
            progress=False,
            group_by="ticker",
            threads=True,
        )
        if frame is None or frame.empty:
            continue
        for symbol in batch:
            if (symbol, "Dividends") not in frame.columns:
                continue
            series = frame[(symbol, "Dividends")]
            if series is None:
                continue
            rows: list[dict[str, Any]] = []
            nonzero = series[series.fillna(0.0) > 0.0]
            for index, value in nonzero.items():
                date = pd.Timestamp(index).date()
                if date < start_date or date > end_date:
                    continue
                rows.append({"date": date.isoformat(), "dividend": float(value)})
            if rows:
                results[symbol] = rows
    path.write_text(json.dumps(results))
    return results


def _default_specs() -> list[DividendCaptureSpec]:
    return [
        DividendCaptureSpec("yield25_max2_adv20m_exopen", 0.25, 2.0, 20_000_000.0, 5, "ex_open"),
        DividendCaptureSpec("yield25_max3_adv20m_exopen", 0.25, 3.0, 20_000_000.0, 5, "ex_open"),
        DividendCaptureSpec("yield50_max3_adv20m_exopen", 0.50, 3.0, 20_000_000.0, 5, "ex_open"),
        DividendCaptureSpec("yield25_max2_adv20m_exclose", 0.25, 2.0, 20_000_000.0, 5, "ex_close"),
    ]


def _simulate_spec(
    *,
    store: Any,
    start_date: dt.date,
    end_date: dt.date,
    dividends_by_symbol: dict[str, list[dict[str, Any]]],
    spec: DividendCaptureSpec,
    capital: float,
) -> tuple[list[DailyPortfolioState], int, dict[str, float]]:
    trading_days = [d for d in store.all_trading_days() if start_date <= d <= end_date]
    trading_index = {date: idx for idx, date in enumerate(trading_days)}
    by_exit_date: dict[dt.date, list[tuple[float, float]]] = {}

    signal_count = 0
    for symbol, rows in dividends_by_symbol.items():
        for row in rows:
            ex_date = dt.date.fromisoformat(row["date"])
            ex_idx = trading_index.get(ex_date)
            if ex_idx is None or ex_idx <= 0:
                continue
            entry_date = trading_days[ex_idx - 1]
            entry_bar = store.get_bar(symbol, entry_date)
            exit_bar = store.get_bar(symbol, ex_date)
            if not entry_bar or not exit_bar:
                continue
            entry_close = float(entry_bar.get("close") or 0.0)
            exit_price = float(
                exit_bar.get("open" if spec.exit_mode == "ex_open" else "close") or 0.0
            )
            if entry_close <= 0 or exit_price <= 0:
                continue
            dividend = float(row["dividend"])
            div_yield_pct = dividend / entry_close * 100.0
            if div_yield_pct < spec.min_dividend_yield_pct:
                continue
            if spec.max_dividend_yield_pct is not None and div_yield_pct > spec.max_dividend_yield_pct:
                continue

            features = store.get_market_features(symbol, entry_date)
            adv = float(features.get("avg_dollar_volume_20d") or 0.0)
            if adv < spec.min_avg_dollar_volume:
                continue

            # Backtest bars are dividend-adjusted; adding the dividend cash
            # again would double-count ex-date carry.
            total_ret = exit_price / entry_close - 1.0
            by_exit_date.setdefault(ex_date, []).append((div_yield_pct, total_ret))
            signal_count += 1

    equity = capital
    peak = capital
    curve: list[DailyPortfolioState] = []
    yearly_pnl: dict[str, float] = {}
    for date in trading_days:
        rows = sorted(by_exit_date.get(date, []), key=lambda item: item[0], reverse=True)
        rows = rows[: spec.max_positions]
        if rows:
            daily_ret = sum(item[1] for item in rows) / len(rows)
            pnl = equity * daily_ret
            yearly_pnl[str(date.year)] = yearly_pnl.get(str(date.year), 0.0) + pnl
            equity *= 1.0 + daily_ret
        peak = max(peak, equity)
        curve.append(
            DailyPortfolioState(
                date=date,
                equity=equity,
                sizing_equity=equity,
                cash_available=equity,
                gross_exposure=0.0,
                net_exposure=0.0,
                reserved_risk_budget=0.0,
                unrealized_pnl=0.0,
                realized_pnl=0.0,
                open_positions=[],
                daily_new_risk_used=0.0,
                peak_equity=peak,
                current_drawdown_pct=((peak - equity) / peak * 100.0) if peak > 0 else 0.0,
            )
        )
    return curve, signal_count, yearly_pnl


def _run_probe(args: argparse.Namespace) -> list[dict[str, Any]]:
    start_date = _parse_date(args.start)
    end_date = _parse_date(args.end, is_end=True)
    rows: list[dict[str, Any]] = []

    for config_path in args.config:
        manifest = load_manifest(config_path)
        config = resolve_config(manifest)
        store = _build_merged_snapshot_store(
            manifest,
            config,
            snapshot_dir_override=None,
        ).slice_by_date_range(start_date, end_date)
        symbols = sorted(str(symbol).upper() for symbol in store._bars.keys())
        dividends_by_symbol = _download_dividend_actions(
            symbols=symbols,
            start_date=start_date,
            end_date=end_date,
            cache_dir=Path(args.cache_dir),
        )
        for spec in _default_specs():
            curve, signal_count, yearly_pnl = _simulate_spec(
                store=store,
                start_date=start_date,
                end_date=end_date,
                dividends_by_symbol=dividends_by_symbol,
                spec=spec,
                capital=args.capital,
            )
            rows.append(
                {
                    "config": config_path,
                    "spec": spec.name,
                    "signals": signal_count,
                    "total_return_pct": round(compute_total_return_pct(curve) or 0.0, 2),
                    "max_drawdown_pct": round(compute_max_drawdown_pct(curve) or 0.0, 2),
                    "sharpe_ratio": round(compute_sharpe_ratio(curve) or 0.0, 3),
                    "yearly_pnl": {year: round(value, 2) for year, value in sorted(yearly_pnl.items())},
                }
            )
    return rows


def main() -> None:
    parser = argparse.ArgumentParser(description="Probe simple dividend capture strategies")
    parser.add_argument("--config", action="append", required=True)
    parser.add_argument("--start", required=True)
    parser.add_argument("--end", required=True)
    parser.add_argument("--capital", type=float, default=10_000.0)
    parser.add_argument("--cache-dir", default="data/cache/dividend_actions")
    parser.add_argument("--json", action="store_true")
    args = parser.parse_args()
    configure_logging("WARNING")
    rows = _run_probe(args)
    if args.json:
        print(json.dumps(rows, indent=2))
    else:
        for row in rows:
            print(row)


if __name__ == "__main__":
    main()
