"""Test all Python code snippets from docs/multi-strategy-ops.mdx.

Each snippet is wrapped in a try/except that prints PASS/FAIL.
Snippets requiring network access, real exchange connections, or hz.run() are skipped.
"""

import traceback

results = []


def run_snippet(name, fn):
    """Execute a snippet function and record pass/fail."""
    try:
        fn()
        results.append((name, True, None))
        print(f"  PASS: {name}")
    except Exception as e:
        results.append((name, False, str(e)))
        print(f"  FAIL: {name}")
        traceback.print_exc()
        print()


# ---------------------------------------------------------------------------
# Snippet 1: Import StrategyBook
# docs line 37: from horizon.multi_strategy import StrategyBook
# ---------------------------------------------------------------------------
def snippet_01_import_strategy_book():
    from horizon.multi_strategy import StrategyBook
    assert StrategyBook is not None

run_snippet("01 - Import StrategyBook", snippet_01_import_strategy_book)


# ---------------------------------------------------------------------------
# Snippet 2: Quick Start (StrategyBook add, update, report, intervention, rebalance)
# docs lines 43-62
# ---------------------------------------------------------------------------
def snippet_02_quick_start():
    import horizon as hz
    from horizon.multi_strategy import StrategyBook

    engine_a = hz.Engine()
    engine_b = hz.Engine()

    book = StrategyBook(total_capital=100_000)
    book.add_strategy("momentum", engine_a, risk_budget=0.4)
    book.add_strategy("mean_revert", engine_b, risk_budget=0.6)

    # Each cycle:
    book.update()
    report = book.report()
    print(f"    Portfolio Sharpe: {report.portfolio_sharpe:.2f}")
    print(f"    Total PnL: {report.total_pnl:+.2f}")

    # Check for problems:
    troubled = book.needs_intervention(max_drawdown=0.10)
    if troubled:
        print(f"    Strategies needing attention: {troubled}")

    # Rebalance:
    new_budgets = book.rebalance_capital(method="risk_parity")
    for name, budget in new_budgets.items():
        print(f"      {name}: {budget:.1%}")

    assert isinstance(report.portfolio_sharpe, float)
    assert isinstance(new_budgets, dict)

run_snippet("02 - Quick Start (add, update, report, intervention, rebalance)", snippet_02_quick_start)


# ---------------------------------------------------------------------------
# Snippet 3: Constructor
# docs line 67: book = StrategyBook(total_capital=100_000.0)
# ---------------------------------------------------------------------------
def snippet_03_constructor():
    from horizon.multi_strategy import StrategyBook
    book = StrategyBook(total_capital=100_000.0)
    assert book is not None
    assert len(book) == 0

run_snippet("03 - Constructor", snippet_03_constructor)


# ---------------------------------------------------------------------------
# Snippet 4: add_strategy
# docs lines 81-85
# ---------------------------------------------------------------------------
def snippet_04_add_strategy():
    import horizon as hz
    from horizon.multi_strategy import StrategyBook

    engine_a = hz.Engine()
    book = StrategyBook(total_capital=100_000)
    book.add_strategy(
        name="momentum",
        engine=engine_a,
        risk_budget=0.4,  # 40% of total capital
    )
    assert len(book) == 1

run_snippet("04 - add_strategy", snippet_04_add_strategy)


# ---------------------------------------------------------------------------
# Snippet 5: remove_strategy
# docs line 103: book.remove_strategy("momentum")
# ---------------------------------------------------------------------------
def snippet_05_remove_strategy():
    import horizon as hz
    from horizon.multi_strategy import StrategyBook

    engine_a = hz.Engine()
    book = StrategyBook(total_capital=100_000)
    book.add_strategy("momentum", engine_a, risk_budget=0.4)
    assert len(book) == 1

    book.remove_strategy("momentum")
    assert len(book) == 0

run_snippet("05 - remove_strategy", snippet_05_remove_strategy)


# ---------------------------------------------------------------------------
# Snippet 6: update
# docs line 115: book.update()
# ---------------------------------------------------------------------------
def snippet_06_update():
    import horizon as hz
    from horizon.multi_strategy import StrategyBook

    engine_a = hz.Engine()
    book = StrategyBook(total_capital=100_000)
    book.add_strategy("momentum", engine_a, risk_budget=1.0)
    book.update()
    # Should not raise

run_snippet("06 - update", snippet_06_update)


# ---------------------------------------------------------------------------
# Snippet 7: report (with per-strategy metrics, correlations, and aggregate)
# docs lines 125-141
# ---------------------------------------------------------------------------
def snippet_07_report():
    import horizon as hz
    from horizon.multi_strategy import StrategyBook

    engine_a = hz.Engine()
    engine_b = hz.Engine()

    book = StrategyBook(total_capital=100_000)
    book.add_strategy("momentum", engine_a, risk_budget=0.4)
    book.add_strategy("mean_revert", engine_b, risk_budget=0.6)

    # Build some history so correlations have data
    for _ in range(5):
        book.update()

    report = book.report()

    # Per-strategy metrics
    for s in report.strategies:
        print(f"    {s.name}: PnL={s.pnl:+.2f}, Sharpe={s.sharpe:.2f}, "
              f"DD={s.drawdown:.1%}, budget={s.risk_budget:.1%}")

    # Pairwise correlations
    for c in report.correlations:
        print(f"    {c.strategy_a} vs {c.strategy_b}: "
              f"corr={c.correlation:.3f} (n={c.n_observations})")

    # Aggregate
    print(f"    Portfolio Sharpe: {report.portfolio_sharpe:.2f}")
    print(f"    Total Drawdown: {report.total_drawdown:.1%}")
    print(f"    Capital Utilization: {report.capital_utilization:.1%}")

    assert hasattr(report, 'strategies')
    assert hasattr(report, 'total_pnl')
    assert hasattr(report, 'total_drawdown')
    assert hasattr(report, 'portfolio_sharpe')
    assert hasattr(report, 'correlations')
    assert hasattr(report, 'capital_utilization')

run_snippet("07 - report (per-strategy, correlations, aggregate)", snippet_07_report)


# ---------------------------------------------------------------------------
# Snippet 8: rebalance_capital (all three methods)
# docs lines 180-188
# ---------------------------------------------------------------------------
def snippet_08_rebalance_capital():
    import horizon as hz
    from horizon.multi_strategy import StrategyBook

    engine_a = hz.Engine()
    engine_b = hz.Engine()

    book = StrategyBook(total_capital=100_000)
    book.add_strategy("momentum", engine_a, risk_budget=0.4)
    book.add_strategy("mean_revert", engine_b, risk_budget=0.6)

    # Build some history for risk_parity and performance methods
    for _ in range(5):
        book.update()

    # Equal allocation
    budgets = book.rebalance_capital(method="equal")
    print(f"    Equal: {budgets}")
    assert abs(sum(budgets.values()) - 1.0) < 1e-10

    # Inverse-volatility: lower vol gets more capital
    budgets = book.rebalance_capital(method="risk_parity")
    print(f"    Risk parity: {budgets}")
    assert abs(sum(budgets.values()) - 1.0) < 1e-10

    # Proportional to Sharpe: higher Sharpe gets more capital
    budgets = book.rebalance_capital(method="performance")
    print(f"    Performance: {budgets}")
    assert abs(sum(budgets.values()) - 1.0) < 1e-10

run_snippet("08 - rebalance_capital (equal, risk_parity, performance)", snippet_08_rebalance_capital)


# ---------------------------------------------------------------------------
# Snippet 9: needs_intervention
# docs lines 211-218
# ---------------------------------------------------------------------------
def snippet_09_needs_intervention():
    import horizon as hz
    from horizon.multi_strategy import StrategyBook

    engine_a = hz.Engine()
    engine_b = hz.Engine()

    book = StrategyBook(total_capital=100_000)
    book.add_strategy("momentum", engine_a, risk_budget=0.4)
    book.add_strategy("mean_revert", engine_b, risk_budget=0.6)
    book.update()

    troubled = book.needs_intervention(
        max_drawdown=0.10,  # Flag if drawdown > 10%
        min_sharpe=0.0,     # Flag if Sharpe < 0 (with enough history)
    )

    for name in troubled:
        print(f"    Strategy '{name}' needs attention")

    assert isinstance(troubled, list)

run_snippet("09 - needs_intervention", snippet_09_needs_intervention)


# ---------------------------------------------------------------------------
# Snippet 10: strategy_pipeline
# docs lines 238-246 -- uses hz.run(), so we test the pipeline function creation
# and basic call instead of the full hz.run() integration.
# SKIP hz.run() but test the returned callable
# ---------------------------------------------------------------------------
def snippet_10_strategy_pipeline():
    import horizon as hz
    from horizon.multi_strategy import StrategyBook

    engine_a = hz.Engine()
    book = StrategyBook(total_capital=100_000)
    book.add_strategy("momentum", engine_a, risk_budget=1.0)

    pipeline_fn = book.strategy_pipeline("momentum")
    assert callable(pipeline_fn)
    assert "strategy_pipeline_momentum" in pipeline_fn.__name__

    # Simulate a call with a mock context
    class MockCtx:
        params = {}

    ctx = MockCtx()
    result = pipeline_fn(ctx, signal=0.5)
    assert result == 0.5  # Passthrough
    assert "strategy_pnl" in ctx.params
    assert "strategy_drawdown" in ctx.params
    assert "strategy_sharpe" in ctx.params

run_snippet("10 - strategy_pipeline (callable creation + mock call)", snippet_10_strategy_pipeline)


# ---------------------------------------------------------------------------
# Snippet 11: Import reconciliation module
# docs line 267: from horizon.reconciliation import reconcile, auto_reconcile, ...
# ---------------------------------------------------------------------------
def snippet_11_import_reconciliation():
    from horizon.reconciliation import reconcile, auto_reconcile, PositionBreak, ReconciliationReport
    assert reconcile is not None
    assert auto_reconcile is not None
    assert PositionBreak is not None
    assert ReconciliationReport is not None

run_snippet("11 - Import reconciliation", snippet_11_import_reconciliation)


# ---------------------------------------------------------------------------
# Snippet 12: reconcile (with exchange_positions)
# docs lines 292-307
# ---------------------------------------------------------------------------
def snippet_12_reconcile():
    import horizon as hz
    from horizon.reconciliation import reconcile

    engine = hz.Engine()

    report = reconcile(
        engine=engine,
        exchange_positions=[
            {"market_id": "btc-100k", "size": 10.0, "side": "yes"},
            {"market_id": "eth-5k", "size": 5.0, "side": "no"},
        ],
    )

    if not report.is_clean:
        for brk in report.breaks:
            print(f"    [{brk.severity.upper()}] {brk.market_id}: {brk.break_type}")
            print(f"      Engine: size={brk.engine_size}, side={brk.engine_side}")
            print(f"      Exchange: size={brk.exchange_size}, side={brk.exchange_side}")
    else:
        print(f"    All clean: {report.matched} positions matched")

    # With no engine positions and 2 exchange positions, we expect breaks
    assert hasattr(report, 'is_clean')
    assert hasattr(report, 'breaks')
    assert hasattr(report, 'matched')
    assert hasattr(report, 'total_engine')
    assert hasattr(report, 'total_exchange')
    assert hasattr(report, 'timestamp')

run_snippet("12 - reconcile (engine vs exchange_positions)", snippet_12_reconcile)


# ---------------------------------------------------------------------------
# Snippet 13: auto_reconcile (pipeline function)
# docs lines 348-356 -- uses hz.run(), so we test creation only
# SKIP hz.run() but test the returned callable
# ---------------------------------------------------------------------------
def snippet_13_auto_reconcile():
    import horizon as hz
    from horizon.reconciliation import auto_reconcile

    engine = hz.Engine()
    pipeline_fn = auto_reconcile(engine, interval=300.0)  # Every 5 minutes
    assert callable(pipeline_fn)
    assert pipeline_fn.__name__ == "auto_reconcile"

run_snippet("13 - auto_reconcile (callable creation)", snippet_13_auto_reconcile)


# ---------------------------------------------------------------------------
# Snippet 14: Full Multi-Strategy Workflow (partial -- skip hz.run())
# docs lines 406-455
# ---------------------------------------------------------------------------
def snippet_14_full_workflow():
    import horizon as hz
    from horizon.multi_strategy import StrategyBook
    from horizon.reconciliation import reconcile, auto_reconcile

    # 1. Set up engines (use default paper engine)
    engine_a = hz.Engine()
    engine_b = hz.Engine()

    # 2. Build the strategy book
    book = StrategyBook(total_capital=100_000)
    book.add_strategy("momentum", engine_a, risk_budget=0.4)
    book.add_strategy("mean_revert", engine_b, risk_budget=0.6)

    # 3. Skipping hz.run() -- test pipeline creation
    auto_recon_fn = auto_reconcile(engine_a, interval=300.0)
    assert callable(auto_recon_fn)
    pipeline_fn = book.strategy_pipeline("momentum")
    assert callable(pipeline_fn)

    # 4. Periodic oversight (e.g., in a monitoring loop)
    # Run multiple update cycles to build history
    for _ in range(5):
        book.update()

    report = book.report()

    # Check correlations -- high correlation reduces diversification benefit
    for c in report.correlations:
        if c.correlation > 0.8:
            print(f"    High correlation: {c.strategy_a} <-> {c.strategy_b} "
                  f"({c.correlation:.2f})")

    # Check for struggling strategies
    troubled = book.needs_intervention(max_drawdown=0.10, min_sharpe=-0.5)
    if troubled:
        print(f"    Strategies needing review: {troubled}")

        # Rebalance away from struggling strategies
        new_budgets = book.rebalance_capital(method="risk_parity")
        print(f"    Suggested budgets: {new_budgets}")

    # 5. Reconcile against exchange
    recon = reconcile(engine_a, exchange_positions=[])
    if not recon.is_clean:
        print(f"    {len(recon.breaks)} breaks detected!")
    else:
        print(f"    Reconciliation clean: {recon.matched} matched")

    assert isinstance(report, hz.MultiStrategyReport)
    assert isinstance(recon, hz.ReconciliationReport)

run_snippet("14 - Full Multi-Strategy Workflow (partial, no hz.run())", snippet_14_full_workflow)


# ---------------------------------------------------------------------------
# Snippet 15: ReconciliationReport and PositionBreak type verification
# (testing the dataclass fields documented in the type reference)
# ---------------------------------------------------------------------------
def snippet_15_type_reference():
    from horizon.reconciliation import PositionBreak, ReconciliationReport

    # PositionBreak fields
    brk = PositionBreak(
        market_id="btc-100k",
        break_type="size_mismatch",
        engine_size=10.0,
        exchange_size=12.0,
        engine_side="yes",
        exchange_side="yes",
        severity="warning",
    )
    assert brk.market_id == "btc-100k"
    assert brk.break_type == "size_mismatch"
    assert brk.engine_size == 10.0
    assert brk.exchange_size == 12.0
    assert brk.engine_side == "yes"
    assert brk.exchange_side == "yes"
    assert brk.severity == "warning"

    # ReconciliationReport fields
    rpt = ReconciliationReport(
        breaks=[brk],
        matched=3,
        total_engine=4,
        total_exchange=4,
        is_clean=False,
        timestamp=1700000000.0,
    )
    assert len(rpt.breaks) == 1
    assert rpt.matched == 3
    assert rpt.total_engine == 4
    assert rpt.total_exchange == 4
    assert rpt.is_clean is False
    assert rpt.timestamp == 1700000000.0

run_snippet("15 - PositionBreak and ReconciliationReport type reference", snippet_15_type_reference)


# ---------------------------------------------------------------------------
# Snippet 16: StrategySlot and StrategyCorrelation type verification
# (testing the dataclass fields documented in the type reference)
# ---------------------------------------------------------------------------
def snippet_16_strategy_types():
    from horizon.multi_strategy import StrategySlot, StrategyCorrelation, MultiStrategyReport

    slot = StrategySlot(
        name="momentum",
        engine=None,  # Placeholder
        risk_budget=0.4,
        pnl=150.0,
        sharpe=1.5,
        drawdown=0.05,
        is_active=True,
    )
    assert slot.name == "momentum"
    assert slot.risk_budget == 0.4
    assert slot.pnl == 150.0
    assert slot.sharpe == 1.5
    assert slot.drawdown == 0.05
    assert slot.is_active is True

    corr = StrategyCorrelation(
        strategy_a="momentum",
        strategy_b="mean_revert",
        correlation=0.35,
        n_observations=100,
    )
    assert corr.strategy_a == "momentum"
    assert corr.strategy_b == "mean_revert"
    assert corr.correlation == 0.35
    assert corr.n_observations == 100

    report = MultiStrategyReport(
        strategies=[slot],
        total_pnl=150.0,
        total_drawdown=0.02,
        portfolio_sharpe=1.2,
        correlations=[corr],
        capital_utilization=0.4,
    )
    assert len(report.strategies) == 1
    assert report.total_pnl == 150.0
    assert report.total_drawdown == 0.02
    assert report.portfolio_sharpe == 1.2
    assert len(report.correlations) == 1
    assert report.capital_utilization == 0.4

run_snippet("16 - StrategySlot, StrategyCorrelation, MultiStrategyReport types", snippet_16_strategy_types)


# ===========================================================================
# Summary
# ===========================================================================
print("\n" + "=" * 60)
print("SUMMARY")
print("=" * 60)

total = len(results)
passed = sum(1 for _, ok, _ in results if ok)
failed = sum(1 for _, ok, _ in results if not ok)

for name, ok, err in results:
    status = "PASS" if ok else "FAIL"
    line = f"  [{status}] {name}"
    if err:
        line += f"  -- {err}"
    print(line)

print(f"\nTotal: {total}  |  Passed: {passed}  |  Failed: {failed}")
if failed == 0:
    print("All snippets passed!")
else:
    print(f"{failed} snippet(s) failed.")
