market_data/tools.py

from __future__ import annotations

import json

from smolagents import tool

from .analytics import (
    black_scholes_greeks,
    classify_volatility_regime,
    rank_current_iv_against_rv,
    realized_volatility,
    summarize_option_chain,
)
from .providers import get_current_quote, get_option_chain, get_price_history, list_option_expirations
from .schemas import VolSnapshot


def json_dumps(payload) -> str:
    return json.dumps(payload, ensure_ascii=False, indent=2, default=str)


@tool
def query_market_asset(symbol: str) -> str:
    """Query the current price and intraday quote data for an asset.

    Args:
        symbol: Yahoo Finance ticker, e.g. AAPL, SPY, ^VIX, BTC-USD, EURUSD=X.
    """
    try:
        return json_dumps({"status": "success", **get_current_quote(symbol).to_dict()})
    except Exception as exc:
        return json_dumps({"status": "error", "symbol": symbol, "message": str(exc)})


@tool
def query_price_history(symbol: str, period: str = "1y", interval: str = "1d") -> str:
    """Query historical OHLCV prices for an asset.

    Args:
        symbol: Yahoo Finance ticker.
        period: Yahoo Finance period such as 1mo, 6mo, 1y, 5y.
        interval: Yahoo Finance interval such as 1d, 1h, 15m.
    """
    try:
        history = get_price_history(symbol, period=period, interval=interval)
        records = history.tail(20).reset_index().to_dict(orient="records")
        return json_dumps(
            {
                "status": "success",
                "symbol": symbol.upper(),
                "period": period,
                "interval": interval,
                "rows_returned": len(records),
                "latest_rows": records,
            }
        )
    except Exception as exc:
        return json_dumps({"status": "error", "symbol": symbol, "message": str(exc)})


@tool
def query_realized_volatility(symbol: str, period: str = "1y") -> str:
    """Calculate realized volatility windows from historical close prices.

    Args:
        symbol: Yahoo Finance ticker.
        period: Yahoo Finance history period.
    """
    try:
        history = get_price_history(symbol, period=period, interval="1d")
        rv = realized_volatility(history["Close"])
        return json_dumps({"status": "success", "symbol": symbol.upper(), "realized_volatility": rv})
    except Exception as exc:
        return json_dumps({"status": "error", "symbol": symbol, "message": str(exc)})


@tool
def query_option_expirations(symbol: str) -> str:
    """List available option expiration dates for an underlying.

    Args:
        symbol: Yahoo Finance ticker.
    """
    try:
        expirations = list_option_expirations(symbol)
        return json_dumps({"status": "success", "symbol": symbol.upper(), "expirations": expirations})
    except Exception as exc:
        return json_dumps({"status": "error", "symbol": symbol, "message": str(exc)})


@tool
def query_option_chain(symbol: str, expiration: str = "") -> str:
    """Query an option chain with liquidity warnings and implied volatility.

    Args:
        symbol: Yahoo Finance ticker.
        expiration: Expiration date in YYYY-MM-DD. Leave empty to use the nearest expiration.
    """
    try:
        chain = get_option_chain(symbol, expiration or None)
        summary = summarize_option_chain(chain)
        payload = chain.to_dict()
        payload["summary"] = summary
        payload["calls"] = payload["calls"][:80]
        payload["puts"] = payload["puts"][:80]
        return json_dumps({"status": "success", **payload})
    except Exception as exc:
        return json_dumps({"status": "error", "symbol": symbol, "message": str(exc)})


@tool
def query_volatility_snapshot(symbol: str, max_expirations: int = 4, history_period: str = "1y") -> str:
    """Summarize realized volatility, ATM IV, IV-RV spread, skew, and term structure.

    Args:
        symbol: Yahoo Finance ticker.
        max_expirations: Number of expirations to sample from the option chain.
        history_period: Yahoo Finance history period for realized volatility.
    """
    try:
        symbol = symbol.strip().upper()
        quote = get_current_quote(symbol)
        history = get_price_history(symbol, period=history_period, interval="1d")
        rv = realized_volatility(history["Close"])
        rv_20d = rv.get("20d")
        expirations = list_option_expirations(symbol)[:max_expirations]

        atm_iv_by_expiration = {}
        iv_rv_spread_by_expiration = {}
        skew_by_expiration = {}
        for expiration in expirations:
            chain = get_option_chain(symbol, expiration)
            summary = summarize_option_chain(chain, realized_vol_20d=rv_20d)
            atm_iv_by_expiration[expiration] = summary["atm_iv"]
            iv_rv_spread_by_expiration[expiration] = summary["iv_rv_spread_20d"]
            skew_by_expiration[expiration] = summary["skew_put_minus_call"]

        valid_term_ivs = [
            value
            for value in atm_iv_by_expiration.values()
            if value is not None
        ]
        current_atm_iv = valid_term_ivs[0] if valid_term_ivs else None
        sampled_skews = [value for value in skew_by_expiration.values() if value is not None]
        front_skew = sampled_skews[0] if sampled_skews else None
        term_structure_slope = (
            float(valid_term_ivs[-1] - valid_term_ivs[0])
            if len(valid_term_ivs) >= 2
            else None
        )
        regime = classify_volatility_regime(
            current_iv=current_atm_iv,
            realized_vol_20d=rv_20d,
            term_structure_slope=term_structure_slope,
            skew=front_skew,
        )
        snapshot = VolSnapshot(
            symbol=symbol,
            current_price=quote.current_price,
            realized_volatility=rv,
            atm_iv_by_expiration=atm_iv_by_expiration,
            iv_rv_spread_by_expiration=iv_rv_spread_by_expiration,
            term_structure_slope=term_structure_slope,
            skew_by_expiration=skew_by_expiration,
        )
        return json_dumps(
            {
                "status": "success",
                **snapshot.to_dict(),
                "front_atm_iv": current_atm_iv,
                "front_skew": front_skew,
                "iv_vs_rv_rank_proxy": rank_current_iv_against_rv(current_atm_iv, rv),
                "volatility_regime": regime,
                "limitations": [
                    "IV rank/percentile is a proxy based on current ATM IV versus realized-volatility windows.",
                    "True historical IV rank requires historical option-chain data from a richer provider.",
                ],
            }
        )
    except Exception as exc:
        return json_dumps({"status": "error", "symbol": symbol, "message": str(exc)})


@tool
def calculate_option_greeks(
    spot: float,
    strike: float,
    time_to_expiry: float,
    volatility: float,
    option_type: str = "call",
    risk_free_rate: float = 0.0,
    dividend_yield: float = 0.0,
) -> str:
    """Calculate Black-Scholes-Merton Greeks for a single option.

    Args:
        spot: Current underlying price.
        strike: Option strike.
        time_to_expiry: Time to expiration in years.
        volatility: Annualized implied volatility as a decimal.
        option_type: call or put.
        risk_free_rate: Annualized risk-free rate as a decimal.
        dividend_yield: Annualized dividend yield as a decimal.
    """
    greeks = black_scholes_greeks(
        spot=spot,
        strike=strike,
        time_to_expiry=time_to_expiry,
        volatility=volatility,
        risk_free_rate=risk_free_rate,
        dividend_yield=dividend_yield,
        option_type=option_type,
    )
    return json_dumps({"status": "success", "greeks": greeks})