"""HJB Market Making pipeline functions.

Pipeline functions:
    hjb_market_maker - Optimal market making with HJB-solved quotes
"""

from __future__ import annotations

from typing import Any, Callable

from horizon._horizon import HJBConfig, solve_hjb, hjb_quote
from horizon.context import Context


def hjb_market_maker(
    feed: str,
    max_inventory: int = 10,
    gamma: float = 0.1,
    sigma: float = 0.3,
    kappa: float = 1.5,
    alpha: float = 0.01,
    terminal_time: float = 1.0,
    n_time_steps: int = 100,
) -> Callable[[Context], list]:
    """Create an HJB optimal market making pipeline function.

    Pre-solves the Avellaneda-Stoikov HJB PDE and uses the solution
    to generate optimal bid/ask quotes in real-time.

    Args:
        feed: Feed name for price data.
        max_inventory: Maximum inventory (absolute).
        gamma: Risk aversion parameter.
        sigma: Volatility parameter.
        kappa: Order arrival intensity decay.
        alpha: Baseline arrival rate.
        terminal_time: Trading session duration.
        n_time_steps: PDE discretization steps.

    Returns:
        Pipeline function: (Context) -> list[Quote]
    """
    from horizon._horizon import auth_require_ultra, Quote
    auth_require_ultra()

    config = HJBConfig(max_inventory, n_time_steps, gamma, sigma, kappa, alpha, terminal_time)
    solution = solve_hjb(config)

    step_counts: dict[str, int] = {}
    inventories: dict[str, int] = {}

    def _hjb_market_maker(ctx: Context) -> list:
        market_id = ctx.market.id if ctx.market else "__default__"
        fd = ctx.feeds.get(feed)

        if market_id not in step_counts:
            step_counts[market_id] = 0
            inventories[market_id] = 0

        if fd is None or fd.price <= 0:
            return []

        step = step_counts[market_id]
        inv = inventories[market_id]
        time_frac = step / max(n_time_steps, 1)

        try:
            quote = hjb_quote(solution, inv, time_frac)

            mid = fd.price if fd.price > 0 else (fd.bid + fd.ask) / 2.0
            bid = max(0.01, mid - quote.bid_delta)
            ask = min(0.99, mid + quote.ask_delta)

            step_counts[market_id] = step + 1

            if bid < ask:
                return [Quote(bid=bid, ask=ask, size=5.0)]
        except (ValueError, RuntimeError):
            pass

        return []

    _hjb_market_maker.__name__ = "hjb_market_maker"
    return _hjb_market_maker
