"""Kelly criterion pipeline integration for hz.run().

Provides a pipeline-compatible sizing function that uses the Rust Kelly
functions for zero-overhead computation.

Usage in pipeline:
    hz.run(
        ...
        pipeline=[estimate_prob, kelly_sizer(fraction=0.25, bankroll=1000.0), to_quotes],
    )
"""

from __future__ import annotations

from typing import Callable

from horizon._horizon import (
    kelly_size,
    liquidity_adjusted_kelly,
)
from horizon.context import Context


def _extract_market_price(ctx: Context) -> float:
    """Extract mid price from first available feed. Returns 0.0 if no data."""
    for fdata in ctx.feeds.values():
        if fdata.bid > 0 and fdata.ask > 0:
            return (fdata.bid + fdata.ask) / 2.0
        if fdata.price > 0:
            return fdata.price
    return 0.0


def kelly_sizer(
    fraction: float = 0.25,
    bankroll: float = 1000.0,
    max_size: float = 100.0,
) -> Callable[[Context, float], float]:
    """Create a pipeline function that sizes positions using Kelly criterion.

    The returned function takes (ctx, estimated_prob) and returns the
    optimal contract size.

    Args:
        fraction: Kelly fraction (0.25 = quarter Kelly). Lower = less variance.
        bankroll: Total capital available.
        max_size: Hard cap on position size.

    Returns:
        Pipeline function: (Context, float) -> float
    """
    def _sizer(ctx: Context, estimated_prob: float) -> float:
        market_price = _extract_market_price(ctx)
        if market_price <= 0.0 or market_price >= 1.0:
            return 0.0

        effective_bankroll = ctx.params.get("bankroll", bankroll)

        return kelly_size(
            estimated_prob,
            market_price,
            effective_bankroll,
            fraction,
            max_size,
        )

    _sizer.__name__ = "kelly_sizer"
    return _sizer


def kelly_sizer_with_liquidity(
    fraction: float = 0.25,
    bankroll: float = 1000.0,
    max_size: float = 100.0,
) -> Callable[[Context, float], float]:
    """Like kelly_sizer but adjusts for available liquidity.

    Uses the ask-side depth from feeds to dampen sizing in thin books.
    """
    def _sizer(ctx: Context, estimated_prob: float) -> float:
        market_price = _extract_market_price(ctx)
        if market_price <= 0.0 or market_price >= 1.0:
            return 0.0

        effective_bankroll = ctx.params.get("bankroll", bankroll)
        available_liq = ctx.params.get("available_liquidity", max_size)

        return liquidity_adjusted_kelly(
            estimated_prob,
            market_price,
            effective_bankroll,
            fraction,
            available_liq,
            max_size,
        )

    _sizer.__name__ = "kelly_sizer_liq"
    return _sizer
