"""Optimal Stopping pipeline functions.

Pipeline functions:
    exit_optimizer - Optimal exit timing using Longstaff-Schwartz
"""

from __future__ import annotations

from typing import Any, Callable

from horizon._horizon import generate_gbm_paths, longstaff_schwartz, should_exit
from horizon.context import Context


def exit_optimizer(
    feed: str,
    n_paths: int = 1000,
    vol: float = 0.3,
    n_steps: int = 20,
    strike_offset: float = 0.0,
) -> Callable[[Context], dict[str, Any]]:
    """Create an optimal exit timing pipeline function.

    Uses Longstaff-Schwartz backward induction to compute optimal
    exercise boundaries for position exit.

    Args:
        feed: Feed name for price data.
        n_paths: Number of Monte Carlo paths for training.
        vol: Volatility assumption for path generation.
        n_steps: Number of time steps.
        strike_offset: Strike offset from current price.

    Returns:
        Pipeline function: (Context) -> dict with keys:
            should_exit, boundary_price, expected_payoff, time_step
    """
    from horizon._horizon import auth_require_pro
    auth_require_pro()

    boundaries: dict[str, list[float]] = {}
    payoffs: dict[str, float] = {}
    step_counts: dict[str, int] = {}

    def _exit_optimizer(ctx: Context) -> dict[str, Any]:
        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

        result = {
            "should_exit": False,
            "boundary_price": 0.0,
            "expected_payoff": 0.0,
            "time_step": step_counts.get(market_id, 0),
        }

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

        # Compute boundary if not yet done
        if market_id not in boundaries:
            try:
                paths = generate_gbm_paths(
                    fd.price, 0.0, vol, 1.0 / n_steps,
                    n_steps, n_paths, 42,
                )
                strike = fd.price + strike_offset
                policy = longstaff_schwartz(paths, strike, 3)
                boundaries[market_id] = list(policy.exercise_boundary)
                payoffs[market_id] = policy.expected_payoff
            except (ValueError, RuntimeError):
                boundaries[market_id] = []
                payoffs[market_id] = 0.0

        boundary = boundaries[market_id]
        step = step_counts[market_id]

        if boundary:
            result["expected_payoff"] = payoffs.get(market_id, 0.0)
            exit_now = should_exit(fd.price, boundary, step)
            result["should_exit"] = exit_now
            if step < len(boundary):
                result["boundary_price"] = boundary[step]
            result["time_step"] = step

        step_counts[market_id] = step + 1

        return result

    _exit_optimizer.__name__ = "exit_optimizer"
    return _exit_optimizer
