"""Test script for docs examples batch 3.

Extracts Python code snippets from:
  1. docs/multi-exchange.mdx (multi-leg / cross-exchange)
  2. docs/examples/cross-exchange-arb.mdx (cross-exchange arb example)
  3. docs/markov.mdx (regime detection)
  4. docs/dashboard.mdx (TUI dashboard)
  5. docs/integrations/mcp-server.mdx (MCP server)

Tests each snippet without running hz.run() or other blocking calls.
Prints PASS/FAIL for each test.
"""

import sys
import traceback

results = []


def run_test(name, fn):
    """Run a single test, catching all exceptions."""
    try:
        fn()
        results.append((name, "PASS", ""))
        print(f"  PASS: {name}")
    except Exception as e:
        tb = traceback.format_exc()
        results.append((name, "FAIL", str(e)))
        print(f"  FAIL: {name}")
        print(f"        {e}")


# =========================================================================
# 1. docs/multi-exchange.mdx - Multi-Exchange
# =========================================================================
print("\n=== docs/multi-exchange.mdx ===")


def test_multi_exchange_imports():
    """Test that multi-exchange related imports work."""
    import horizon as hz
    # Classes/functions used in multi-exchange docs
    assert hasattr(hz, "Polymarket")
    assert hasattr(hz, "Kalshi")
    assert hasattr(hz, "Context")
    assert hasattr(hz, "Risk")
    assert hasattr(hz, "BinanceWS")
    assert hasattr(hz, "PolymarketBook")
    assert hasattr(hz, "KalshiBook")
    assert hasattr(hz, "quotes")
    assert hasattr(hz, "run")

run_test("multi_exchange_imports", test_multi_exchange_imports)


def test_multi_exchange_exchange_classes():
    """Test Polymarket/Kalshi exchange classes can be referenced."""
    from horizon.exchanges import Polymarket, Kalshi
    # These are dataclass/config classes, just verify they exist
    assert Polymarket is not None
    assert Kalshi is not None

run_test("multi_exchange_exchange_classes", test_multi_exchange_exchange_classes)


def test_multi_exchange_context_feeddata():
    """Snippet: ctx.feeds.get("poly", hz.context.FeedData())"""
    import horizon as hz
    from horizon.context import FeedData
    fd = FeedData()
    assert fd.price == 0.0
    assert fd.bid == 0.0
    assert fd.ask == 0.0

run_test("multi_exchange_context_feeddata", test_multi_exchange_context_feeddata)


def test_multi_exchange_quotes_helper():
    """Snippet: hz.quotes(fair, spread, size=5)"""
    import horizon as hz
    result = hz.quotes(0.5, 0.02, size=5)
    assert isinstance(result, list)
    assert len(result) == 1
    q = result[0]
    assert hasattr(q, "bid")
    assert hasattr(q, "ask")
    assert hasattr(q, "size")
    assert abs(q.bid - 0.49) < 1e-9
    assert abs(q.ask - 0.51) < 1e-9

run_test("multi_exchange_quotes_helper", test_multi_exchange_quotes_helper)


def test_multi_exchange_cross_strategy_pattern():
    """Snippet: fair_value and quoter functions from cross-exchange strategy pattern."""
    import horizon as hz
    from horizon.context import FeedData

    # Reproduce the fair_value function from docs
    def fair_value(ctx):
        poly = ctx.feeds.get("poly", FeedData())
        kalshi = ctx.feeds.get("kalshi", FeedData())
        prices = []
        if poly.price > 0:
            prices.append(poly.price)
        if kalshi.price > 0:
            prices.append(kalshi.price)
        return sum(prices) / len(prices) if prices else 0.5

    # Reproduce the quoter function from docs
    def quoter(ctx, fair):
        poly = ctx.feeds.get("poly", FeedData()).price or fair
        kalshi = ctx.feeds.get("kalshi", FeedData()).price or fair
        divergence = abs(poly - kalshi)
        spread = 0.02 + divergence * 0.5
        return hz.quotes(fair, spread, size=5)

    # Verify functions are callable
    assert callable(fair_value)
    assert callable(quoter)

run_test("multi_exchange_cross_strategy_pattern", test_multi_exchange_cross_strategy_pattern)


def test_multi_exchange_engine_api():
    """Test Engine class has multi-exchange methods."""
    from horizon._horizon import Engine
    # Check that Engine has the documented methods
    assert hasattr(Engine, "add_exchange")
    assert hasattr(Engine, "exchange_names")
    assert hasattr(Engine, "exchange_count")
    assert hasattr(Engine, "exchange_name")
    assert hasattr(Engine, "set_netting_pair")
    assert hasattr(Engine, "netting_pairs")

run_test("multi_exchange_engine_api", test_multi_exchange_engine_api)


def test_multi_exchange_engine_submit_order_signature():
    """Verify Engine.submit_order and submit_quotes have exchange param."""
    from horizon._horizon import Engine
    import inspect
    # These are Rust-implemented, can't always inspect, but check they exist
    assert hasattr(Engine, "submit_order")
    assert hasattr(Engine, "submit_quotes")
    assert hasattr(Engine, "sync_positions")

run_test("multi_exchange_engine_submit_order_signature", test_multi_exchange_engine_submit_order_signature)


# =========================================================================
# 2. docs/examples/cross-exchange-arb.mdx - Cross-Exchange Arbitrage
# =========================================================================
print("\n=== docs/examples/cross-exchange-arb.mdx ===")


def test_cross_arb_full_imports():
    """Full code snippet imports: hz, Context, FeedData, quotes, Polymarket, Kalshi."""
    import horizon as hz
    from horizon.context import FeedData
    assert hasattr(hz, "Polymarket")
    assert hasattr(hz, "Kalshi")
    assert hasattr(hz, "PolymarketBook")
    assert hasattr(hz, "KalshiBook")
    assert hasattr(hz, "Risk")
    assert hasattr(hz, "quotes")
    assert hasattr(hz, "run")

run_test("cross_arb_full_imports", test_cross_arb_full_imports)


def test_cross_arb_fair_value_fn():
    """Snippet: fair_value function from cross-exchange arb example."""
    import horizon as hz
    from horizon.context import FeedData

    def fair_value(ctx):
        poly_price = ctx.feeds.get("poly", FeedData()).price or 0.5
        kalshi_price = ctx.feeds.get("kalshi", FeedData()).price or 0.5
        return (poly_price + kalshi_price) / 2

    assert callable(fair_value)

run_test("cross_arb_fair_value_fn", test_cross_arb_fair_value_fn)


def test_cross_arb_quoter_fn():
    """Snippet: quoter function from cross-exchange arb example."""
    import horizon as hz
    from horizon.context import FeedData

    def quoter(ctx, fair):
        poly_price = ctx.feeds.get("poly", FeedData()).price or fair
        kalshi_price = ctx.feeds.get("kalshi", FeedData()).price or fair
        divergence = abs(poly_price - kalshi_price)
        base_spread = 0.02
        spread = base_spread + divergence * 0.5
        return hz.quotes(fair, spread, size=5)

    assert callable(quoter)

run_test("cross_arb_quoter_fn", test_cross_arb_quoter_fn)


def test_cross_arb_directional_variant():
    """Snippet: directional arb quoter variant."""
    import horizon as hz
    from horizon.context import FeedData

    def quoter(ctx, fair):
        poly = ctx.feeds.get("poly", FeedData()).price or fair
        kalshi = ctx.feeds.get("kalshi", FeedData()).price or fair
        if poly - kalshi > 0.05:
            return hz.quotes(fair, spread=0.01, size=10)
        elif kalshi - poly > 0.05:
            return hz.quotes(fair, spread=0.01, size=10)
        else:
            return hz.quotes(fair, spread=0.06, size=2)

    assert callable(quoter)

run_test("cross_arb_directional_variant", test_cross_arb_directional_variant)


def test_cross_arb_arb_scanner():
    """Snippet: hz.arb_scanner() pipeline factory."""
    import horizon as hz
    assert hasattr(hz, "arb_scanner")
    # Test that arb_scanner is callable with the documented params
    scanner = hz.arb_scanner(
        market_id="will-btc-hit-100k-by-end-of-2025",
        exchanges=["polymarket", "kalshi"],
        feed_map={"polymarket": "poly", "kalshi": "kalshi"},
        min_edge=0.02,
        auto_execute=True,
        cooldown=10.0,
    )
    assert callable(scanner)

run_test("cross_arb_arb_scanner", test_cross_arb_arb_scanner)


def test_cross_arb_arb_sweep():
    """Snippet: hz.arb_sweep() import check."""
    import horizon as hz
    assert hasattr(hz, "arb_sweep")
    assert callable(hz.arb_sweep)

run_test("cross_arb_arb_sweep", test_cross_arb_arb_sweep)


def test_cross_arb_arb_result():
    """Snippet: ArbResult type import."""
    import horizon as hz
    assert hasattr(hz, "ArbResult")

run_test("cross_arb_arb_result", test_cross_arb_arb_result)


# =========================================================================
# 3. docs/markov.mdx - Markov Regime Detection
# =========================================================================
print("\n=== docs/markov.mdx ===")


def test_markov_imports():
    """Test all markov-related imports from docs."""
    import horizon as hz
    assert hasattr(hz, "MarkovRegimeModel")
    assert hasattr(hz, "prices_to_returns")
    assert hasattr(hz, "markov_regime")

run_test("markov_imports", test_markov_imports)


def test_markov_constructor():
    """Snippet: model = hz.MarkovRegimeModel(n_states=2)"""
    import horizon as hz
    model = hz.MarkovRegimeModel(n_states=2)
    assert model is not None
    assert model.n_states() == 2
    assert model.is_trained() is False

run_test("markov_constructor", test_markov_constructor)


def test_markov_constructor_3_states():
    """Snippet: model = hz.MarkovRegimeModel(n_states=3) (from offline analysis)"""
    import horizon as hz
    model = hz.MarkovRegimeModel(n_states=3)
    assert model.n_states() == 3

run_test("markov_constructor_3_states", test_markov_constructor_3_states)


def test_markov_fit():
    """Snippet: log_likelihood = model.fit(data, max_iters=100, tol=1e-6)"""
    import horizon as hz
    import math
    model = hz.MarkovRegimeModel(n_states=2)
    # Generate synthetic returns: mix of calm and volatile
    import random
    random.seed(42)
    data = [random.gauss(0, 0.01) for _ in range(100)] + \
           [random.gauss(0, 0.05) for _ in range(100)]
    ll = model.fit(data, max_iters=100, tol=1e-6)
    assert isinstance(ll, float)
    assert model.is_trained() is True

run_test("markov_fit", test_markov_fit)


def test_markov_decode():
    """Snippet: states = model.decode(data)"""
    import horizon as hz
    import random
    random.seed(42)
    data = [random.gauss(0, 0.01) for _ in range(100)] + \
           [random.gauss(0, 0.05) for _ in range(100)]
    model = hz.MarkovRegimeModel(n_states=2)
    model.fit(data, max_iters=100)
    states = model.decode(data)
    assert isinstance(states, list)
    assert len(states) == len(data)
    assert all(s in (0, 1) for s in states)

run_test("markov_decode", test_markov_decode)


def test_markov_filter_step():
    """Snippet: probs = model.filter_step(observation)"""
    import horizon as hz
    import random
    random.seed(42)
    data = [random.gauss(0, 0.01) for _ in range(100)] + \
           [random.gauss(0, 0.05) for _ in range(100)]
    model = hz.MarkovRegimeModel(n_states=2)
    model.fit(data, max_iters=100)
    probs = model.filter_step(0.005)
    assert isinstance(probs, list)
    assert len(probs) == 2
    assert abs(sum(probs) - 1.0) < 1e-6

run_test("markov_filter_step", test_markov_filter_step)


def test_markov_predict():
    """Snippet: next_probs = model.predict()"""
    import horizon as hz
    import random
    random.seed(42)
    data = [random.gauss(0, 0.01) for _ in range(100)] + \
           [random.gauss(0, 0.05) for _ in range(100)]
    model = hz.MarkovRegimeModel(n_states=2)
    model.fit(data, max_iters=100)
    model.filter_step(0.005)  # Need at least one filter step
    next_probs = model.predict()
    assert isinstance(next_probs, list)
    assert len(next_probs) == 2
    assert abs(sum(next_probs) - 1.0) < 1e-6

run_test("markov_predict", test_markov_predict)


def test_markov_smooth():
    """Snippet: smoothed = model.smooth(data)"""
    import horizon as hz
    import random
    random.seed(42)
    data = [random.gauss(0, 0.01) for _ in range(100)] + \
           [random.gauss(0, 0.05) for _ in range(100)]
    model = hz.MarkovRegimeModel(n_states=2)
    model.fit(data, max_iters=100)
    smoothed = model.smooth(data)
    assert isinstance(smoothed, list)
    assert len(smoothed) == len(data)
    # Each element should be a list of probabilities
    for row in smoothed:
        assert isinstance(row, list)
        assert len(row) == 2
        assert abs(sum(row) - 1.0) < 1e-6

run_test("markov_smooth", test_markov_smooth)


def test_markov_transition_matrix():
    """Snippet: tm = model.transition_matrix()"""
    import horizon as hz
    import random
    random.seed(42)
    data = [random.gauss(0, 0.01) for _ in range(100)] + \
           [random.gauss(0, 0.05) for _ in range(100)]
    model = hz.MarkovRegimeModel(n_states=2)
    model.fit(data, max_iters=100)
    tm = model.transition_matrix()
    assert isinstance(tm, list)
    assert len(tm) == 2
    for row in tm:
        assert len(row) == 2
        assert abs(sum(row) - 1.0) < 1e-6

run_test("markov_transition_matrix", test_markov_transition_matrix)


def test_markov_emission_params():
    """Snippet: model.emission_params() -> list of (mean, variance)"""
    import horizon as hz
    import random
    random.seed(42)
    data = [random.gauss(0, 0.01) for _ in range(100)] + \
           [random.gauss(0, 0.05) for _ in range(100)]
    model = hz.MarkovRegimeModel(n_states=2)
    model.fit(data, max_iters=100)
    params = model.emission_params()
    assert isinstance(params, list)
    assert len(params) == 2
    for mean, var in params:
        assert isinstance(mean, float)
        assert isinstance(var, float)
        assert var >= 0.0

run_test("markov_emission_params", test_markov_emission_params)


def test_markov_current_regime():
    """Snippet: model.current_regime() -> int"""
    import horizon as hz
    import random
    random.seed(42)
    data = [random.gauss(0, 0.01) for _ in range(100)] + \
           [random.gauss(0, 0.05) for _ in range(100)]
    model = hz.MarkovRegimeModel(n_states=2)
    model.fit(data, max_iters=100)
    model.filter_step(0.005)
    regime = model.current_regime()
    assert isinstance(regime, int)
    assert regime in (0, 1)

run_test("markov_current_regime", test_markov_current_regime)


def test_markov_filtered_probs():
    """Snippet: model.filtered_probs() -> list[float]"""
    import horizon as hz
    import random
    random.seed(42)
    data = [random.gauss(0, 0.01) for _ in range(100)] + \
           [random.gauss(0, 0.05) for _ in range(100)]
    model = hz.MarkovRegimeModel(n_states=2)
    model.fit(data, max_iters=100)
    model.filter_step(0.005)
    probs = model.filtered_probs()
    assert isinstance(probs, list)
    assert len(probs) == 2
    assert abs(sum(probs) - 1.0) < 1e-6

run_test("markov_filtered_probs", test_markov_filtered_probs)


def test_markov_reset_filter():
    """Snippet: model.reset_filter()"""
    import horizon as hz
    import random
    random.seed(42)
    data = [random.gauss(0, 0.01) for _ in range(100)] + \
           [random.gauss(0, 0.05) for _ in range(100)]
    model = hz.MarkovRegimeModel(n_states=2)
    model.fit(data, max_iters=100)
    model.filter_step(0.005)
    model.reset_filter()
    # After reset, filtered_probs should be initial distribution
    probs = model.filtered_probs()
    assert isinstance(probs, list)
    assert len(probs) == 2

run_test("markov_reset_filter", test_markov_reset_filter)


def test_markov_prices_to_returns():
    """Snippet: returns = hz.prices_to_returns(prices)"""
    import horizon as hz
    prices = [100.0, 101.0, 99.0, 102.0]
    returns = hz.prices_to_returns(prices)
    assert isinstance(returns, list)
    assert len(returns) == 3
    # First return: ln(101/100) ~ 0.00995
    assert abs(returns[0] - 0.00995) < 0.001

run_test("markov_prices_to_returns", test_markov_prices_to_returns)


def test_markov_regime_factory():
    """Snippet: hz.markov_regime(model=None, n_states=2, warmup=100, feed='book', param_name='regime')"""
    import horizon as hz
    fn = hz.markov_regime(
        model=None,
        n_states=2,
        warmup=100,
        feed="book",
        param_name="regime",
    )
    assert callable(fn)

run_test("markov_regime_factory", test_markov_regime_factory)


def test_markov_regime_with_pretrained():
    """Snippet: hz.markov_regime(model=model, feed='book') with pre-trained model."""
    import horizon as hz
    import random
    random.seed(42)
    data = [random.gauss(0, 0.01) for _ in range(100)] + \
           [random.gauss(0, 0.05) for _ in range(100)]
    model = hz.MarkovRegimeModel(n_states=2)
    model.fit(data, max_iters=100)
    fn = hz.markov_regime(model=model, feed="book")
    assert callable(fn)

run_test("markov_regime_with_pretrained", test_markov_regime_with_pretrained)


def test_markov_regime_adaptive_quoter_fn():
    """Snippet: adaptive_quoter function from regime-adaptive market maker example."""
    import horizon as hz

    def adaptive_quoter(ctx):
        regime = ctx.params.get("regime", 0)
        vol_prob = ctx.params.get("regime_vol_state", 0.0)
        base_spread = 0.04
        spread = base_spread + vol_prob * 0.06
        size = 10.0 if regime == 0 else 3.0
        fair = ctx.feed.price
        return hz.quotes(fair=fair, spread=spread, size=size)

    assert callable(adaptive_quoter)

run_test("markov_regime_adaptive_quoter_fn", test_markov_regime_adaptive_quoter_fn)


def test_markov_regime_gate_fn():
    """Snippet: regime_gate function from regime-gated trading example."""
    def regime_gate(ctx):
        regime = ctx.params.get("regime", 0)
        if regime != 0:
            return None
        return {"active": True}

    assert callable(regime_gate)

run_test("markov_regime_gate_fn", test_markov_regime_gate_fn)


def test_markov_offline_analysis():
    """Snippet: Offline analysis example (fit, decode, emission_params, transition_matrix, smooth)."""
    import horizon as hz
    import random
    random.seed(42)
    returns = [random.gauss(0, 0.01) for _ in range(100)] + \
              [random.gauss(0, 0.05) for _ in range(100)]

    model = hz.MarkovRegimeModel(n_states=3)
    model.fit(returns, max_iters=200)

    states = model.decode(returns)
    assert len(states) == len(returns)

    for i, (mean, var) in enumerate(model.emission_params()):
        assert isinstance(mean, float)
        assert isinstance(var, float)

    tm = model.transition_matrix()
    for row in tm:
        # format check as done in docs
        formatted = [f"{p:.3f}" for p in row]
        assert len(formatted) == 3

    smoothed = model.smooth(returns)
    assert len(smoothed) == len(returns)

run_test("markov_offline_analysis", test_markov_offline_analysis)


# =========================================================================
# 4. docs/dashboard.mdx - Dashboard (TUI)
# =========================================================================
print("\n=== docs/dashboard.mdx ===")


def test_dashboard_module_import():
    """Test that horizon.dashboard can be imported."""
    import horizon.dashboard
    assert hasattr(horizon.dashboard, "create_app")

run_test("dashboard_module_import", test_dashboard_module_import)


def test_dashboard_tui_import():
    """Test that horizon.dashboard.tui and HorizonApp class exist."""
    from horizon.dashboard.tui import HorizonApp
    assert HorizonApp is not None
    # Verify it's a Textual App subclass
    from textual.app import App
    assert issubclass(HorizonApp, App)

run_test("dashboard_tui_import", test_dashboard_tui_import)


def test_dashboard_panel_classes():
    """Test that panel classes exist in dashboard.tui."""
    from horizon.dashboard import tui
    # Panels referenced in the docs
    assert hasattr(tui, "PnLPanel")
    assert hasattr(tui, "HorizonApp")

run_test("dashboard_panel_classes", test_dashboard_panel_classes)


def test_dashboard_run_kwarg():
    """Docs snippet: hz.run(dashboard=True, ...) - just test run accepts it via inspect."""
    import horizon as hz
    import inspect
    sig = inspect.signature(hz.run)
    param_names = list(sig.parameters.keys())
    assert "dashboard" in param_names, f"'dashboard' not in hz.run params: {param_names}"

run_test("dashboard_run_kwarg", test_dashboard_run_kwarg)


# =========================================================================
# 5. docs/integrations/mcp-server.mdx - MCP Server
# =========================================================================
print("\n=== docs/integrations/mcp-server.mdx ===")


def test_mcp_server_import():
    """Test that horizon.mcp_server can be imported."""
    import horizon.mcp_server
    assert hasattr(horizon.mcp_server, "mcp")

run_test("mcp_server_import", test_mcp_server_import)


def test_mcp_tools_import():
    """Test that horizon.tools can be imported."""
    import horizon.tools
    assert horizon.tools is not None

run_test("mcp_tools_import", test_mcp_tools_import)


def test_mcp_fastmcp_instance():
    """Test that the mcp object is a FastMCP instance."""
    from horizon.mcp_server import mcp
    from mcp.server.fastmcp import FastMCP
    assert isinstance(mcp, FastMCP)

run_test("mcp_fastmcp_instance", test_mcp_fastmcp_instance)


# =========================================================================
# Summary
# =========================================================================
print("\n" + "=" * 60)
total = len(results)
passed = sum(1 for _, status, _ in results if status == "PASS")
failed = sum(1 for _, status, _ in results if status == "FAIL")
print(f"TOTAL: {total}  |  PASSED: {passed}  |  FAILED: {failed}")
print("=" * 60)

if failed > 0:
    print("\nFailed tests:")
    for name, status, err in results:
        if status == "FAIL":
            print(f"  - {name}: {err}")

sys.exit(0 if failed == 0 else 1)
