"""
TradeStation MarketData MCP Tools.

Provides tools for fetching market quotes, historical bars, symbol info,
option chains, market depth, and streaming market data.

All tool descriptions are designed so the LLM can correctly route
"What's the stock price of MSFT?" to get_quote_snapshots (not brokerage),
and "Show me a chart for AAPL" to get_bars (not quotes).
"""

import json
import logging
from typing import Any

logger = logging.getLogger(__name__)


def register_market_data_tools(mcp, ts_client):
    """Register all MarketData tools on the given FastMCP server instance.

    Args:
        mcp: FastMCP server instance.
        ts_client: TradeStationClient instance for API calls.
    """

    # ── Quote Snapshots ────────────────────────────────────────────

    @mcp.tool()
    async def get_quote_snapshots(symbols: str) -> str:
        """MARKET DATA: Get current price quotes for one or more stocks, ETFs, futures, or crypto symbols.

        Use this when the user asks about current stock price, last trade price,
        bid/ask spread, volume, 52-week high/low, or net change for any symbol.

        Example questions this tool answers:
        - "What's the stock price of MSFT?"
        - "How is AAPL trading right now?"
        - "What's the bid/ask on TSLA?"
        - "Show me the current price of Bitcoin"
        - "What's the latest quote for GOOG?"

        DO NOT use this for account balances, portfolio value, or positions — those
        are BROKERAGE tools (get_balances, get_positions).

        Args:
            symbols: Comma-separated ticker symbols, e.g. "MSFT" or "AAPL,GOOG,AMZN,BTCUSD". Max 100 symbols.

        Returns:
            JSON with Quotes array containing Last, Bid, Ask, BidSize, AskSize,
            Volume, Open, High, Low, Close, NetChange, NetChangePct,
            High52Week, Low52Week, PreviousClose, TradeTime, and MarketFlags per symbol.
        """
        data = await ts_client.get(f"/v3/marketdata/quotes/{symbols}")
        return json.dumps(data, indent=2)

    # ── Historical Bars ────────────────────────────────────────────

    @mcp.tool()
    async def get_bars(
        symbol: str,
        interval: str = "1",
        unit: str = "Daily",
        barsback: str | None = None,
        firstdate: str | None = None,
        lastdate: str | None = None,
        sessiontemplate: str | None = None,
    ) -> str:
        """MARKET DATA: Get historical price bars (OHLCV candlestick data) for a symbol.

        Use this when the user asks about historical prices, price charts, candlestick data,
        or price history over a time range. NOT for the current/live price — use
        get_quote_snapshots for that instead.

        Example questions this tool answers:
        - "Show me MSFT daily bars for the last month"
        - "Get 5-minute candles for AAPL"
        - "What was TSLA's price history last week?"
        - "Show me weekly bars for SPY going back 52 weeks"

        Args:
            symbol: Single ticker symbol (e.g. "MSFT").
            interval: Bar interval size. Default "1". For minute bars, the number of minutes
                      aggregated in each bar (max 1440). For Daily/Weekly/Monthly, must be "1".
            unit: Time unit — "Minute", "Daily", "Weekly", or "Monthly". Default "Daily".
            barsback: Number of bars to fetch. Mutually exclusive with firstdate.
            firstdate: Start date in ISO format e.g. "2025-01-01T00:00:00Z". Mutually exclusive with barsback.
            lastdate: End date in ISO format. Defaults to current time.
            sessiontemplate: For intraday US equities only — "USEQPre", "USEQPost",
                             "USEQPreAndPost", "USEQ24Hour", or "Default".

        Returns:
            JSON with Bars array. Each bar has Open, High, Low, Close, TotalVolume,
            TimeStamp, DownTicks, DownVolume, UpTicks, UpVolume, TotalTicks, OpenInterest.
        """
        params: dict[str, Any] = {"interval": interval, "unit": unit}
        if barsback:
            params["barsback"] = barsback
        if firstdate:
            params["firstdate"] = firstdate
        if lastdate:
            params["lastdate"] = lastdate
        if sessiontemplate:
            params["sessiontemplate"] = sessiontemplate
        data = await ts_client.get(f"/v3/marketdata/barcharts/{symbol}", params=params)
        return json.dumps(data, indent=2)

    # ── Stream Bars ────────────────────────────────────────────────

    @mcp.tool()
    async def stream_bars(
        symbol: str,
        interval: str = "1",
        unit: str = "Daily",
        barsback: str | None = None,
        sessiontemplate: str | None = None,
        duration_seconds: int = 5,
    ) -> str:
        """MARKET DATA: Stream live price bars in real-time for a symbol.

        Use this when the user wants LIVE/STREAMING/REAL-TIME bar updates, not just
        historical data. For historical-only bars, use get_bars instead.

        Args:
            symbol: Single ticker symbol (e.g. "MSFT").
            interval: Bar interval. Default "1".
            unit: "Minute", "Daily", "Weekly", or "Monthly". Default "Daily".
            barsback: Number of historical bars to include before live data begins.
            sessiontemplate: Intraday session template for US equities.
            duration_seconds: How many seconds to collect streaming data before returning. Default 5.

        Returns:
            JSON array of bar objects received during the streaming window.
        """
        params: dict[str, Any] = {"interval": interval, "unit": unit}
        if barsback:
            params["barsback"] = barsback
        if sessiontemplate:
            params["sessiontemplate"] = sessiontemplate
        data = await ts_client.stream_get(
            f"/v3/marketdata/stream/barcharts/{symbol}",
            params=params,
            duration_seconds=duration_seconds,
        )
        return json.dumps(data, indent=2)

    # ── Crypto Symbol Names ────────────────────────────────────────

    @mcp.tool()
    async def get_crypto_symbol_names() -> str:
        """MARKET DATA: List all available cryptocurrency trading pair symbol names.

        Use this when the user asks what crypto symbols or cryptocurrency pairs
        are available on TradeStation (e.g. BTCUSD, ETHUSD).

        Returns:
            JSON with SymbolNames array of available crypto pair strings.
        """
        data = await ts_client.get("/v3/marketdata/symbollists/cryptopairs/symbolnames")
        return json.dumps(data, indent=2)

    # ── Symbol Details ─────────────────────────────────────────────

    @mcp.tool()
    async def get_symbol_details(symbols: str) -> str:
        """MARKET DATA: Get detailed symbol information, metadata, and price formatting.

        Use this when the user asks about a symbol's description, exchange, category,
        asset type, price format, point value, or minimum tick size. NOT for current
        price — use get_quote_snapshots for that.

        Example questions:
        - "What exchange does MSFT trade on?"
        - "Tell me about symbol XYZ"
        - "What's the point value of ES futures?"
        - "What type of asset is BTCUSD?"

        Args:
            symbols: Comma-separated symbols (max 50), e.g. "MSFT,BTCUSD".

        Returns:
            JSON with Symbols array containing Description, Exchange, Category,
            ExpirationDate, PriceFormat (Fraction/SubFraction), PointValue, MinMove, etc.
        """
        data = await ts_client.get(f"/v3/marketdata/symbols/{symbols}")
        return json.dumps(data, indent=2)

    # ── Option Expirations ─────────────────────────────────────────

    @mcp.tool()
    async def get_option_expirations(
        underlying: str,
        strike_price: float | None = None,
    ) -> str:
        """MARKET DATA: Get available option expiration dates for an equity or index.

        Use this when the user asks about option expiration dates, when options expire,
        or available contract expirations for a symbol.

        Example questions:
        - "When do AAPL options expire?"
        - "What expirations are available for SPY options?"

        Args:
            underlying: The underlying equity/index symbol (e.g. "AAPL").
            strike_price: Optional — filter expirations to a specific strike price.

        Returns:
            JSON with Expirations array containing Date and Type (Weekly/Monthly/etc.).
        """
        params: dict[str, Any] = {}
        if strike_price is not None:
            params["strikePrice"] = strike_price
        data = await ts_client.get(f"/v3/marketdata/options/expirations/{underlying}", params=params or None)
        return json.dumps(data, indent=2)

    # ── Option Risk/Reward ─────────────────────────────────────────

    @mcp.tool()
    async def get_option_risk_reward(
        spread_price: float,
        legs: str,
    ) -> str:
        """MARKET DATA: Analyze risk vs. reward of a potential option spread trade.

        Use this when the user asks about max profit, max loss, breakeven points,
        or risk/reward analysis for an options strategy. Not applicable for spreads
        with different expirations (Calendar/Diagonal).

        Example questions:
        - "What's the max profit on this butterfly spread?"
        - "Analyze the risk/reward for my iron condor"

        Args:
            spread_price: The quoted price for the option spread trade.
            legs: JSON string array of legs. Each leg: {"Symbol": "AAPL 250321C150", "Ratio": 1, "Type": "Call"}.

        Returns:
            JSON with MaxGain, MaxLoss, BreakevenPoints, and whether gain/loss is infinite.
        """
        legs_parsed = json.loads(legs) if isinstance(legs, str) else legs
        body = {"SpreadPrice": spread_price, "Legs": legs_parsed}
        data = await ts_client.post("/v3/marketdata/options/riskreward", json_body=body)
        return json.dumps(data, indent=2)

    # ── Option Spread Types ────────────────────────────────────────

    @mcp.tool()
    async def get_option_spread_types() -> str:
        """MARKET DATA: List all available option spread/strategy types.

        Use this when the user asks what kinds of option spreads or strategies
        are available (butterfly, iron condor, vertical, straddle, calendar, etc.).

        Returns:
            JSON with SpreadTypes array, each with Name, LegCount, and description.
        """
        data = await ts_client.get("/v3/marketdata/options/spreadtypes")
        return json.dumps(data, indent=2)

    # ── Option Strikes ─────────────────────────────────────────────

    @mcp.tool()
    async def get_option_strikes(
        underlying: str,
        spread_type: str = "Single",
        strike_interval: int = 1,
        expiration: str | None = None,
        expiration2: str | None = None,
    ) -> str:
        """MARKET DATA: Get available option strike prices for an underlying symbol.

        Use this when the user asks about available strikes for options on a symbol.

        Args:
            underlying: The underlying equity/index symbol (e.g. "AAPL").
            spread_type: Spread type — default "Single". Use get_option_spread_types for valid values.
            strike_interval: Interval between strikes (1=consecutive, 2=skip one). Default 1.
            expiration: Expiration date, e.g. "12-17-2025". Defaults to next expiration.
            expiration2: Second expiration for Calendar/Diagonal spreads.

        Returns:
            JSON with SpreadType and Strikes array of strike price groups.
        """
        params: dict[str, Any] = {
            "spreadType": spread_type,
            "strikeInterval": strike_interval,
        }
        if expiration:
            params["expiration"] = expiration
        if expiration2:
            params["expiration2"] = expiration2
        data = await ts_client.get(f"/v3/marketdata/options/strikes/{underlying}", params=params)
        return json.dumps(data, indent=2)

    # ── Stream Option Chain ────────────────────────────────────────

    @mcp.tool()
    async def stream_option_chain(
        underlying: str,
        expiration: str | None = None,
        strike_proximity: int = 5,
        spread_type: str = "Single",
        option_type: str = "All",
        strike_range: str = "All",
        enable_greeks: bool = True,
        price_center: float | None = None,
        strike_interval: int = 1,
        risk_free_rate: float | None = None,
        duration_seconds: int = 5,
    ) -> str:
        """MARKET DATA: Stream a live option chain with greeks and pricing for an underlying symbol.

        Use this when the user asks for an option chain, option prices with greeks
        (delta, gamma, theta, vega), implied volatility, or wants to browse available
        option contracts for a symbol.

        Example questions:
        - "Show me the AAPL option chain"
        - "What are the greeks for SPY options?"
        - "Show me call options for MSFT near the money"

        Args:
            underlying: The underlying symbol (e.g. "AAPL").
            expiration: Expiration date filter. Defaults to next expiration.
            strike_proximity: Number of strikes above/below center to show. Default 5.
            spread_type: Default "Single". Use get_option_spread_types for valid values.
            option_type: "All", "Call", or "Put". Default "All".
            strike_range: "All", "ITM" (in-the-money), or "OTM" (out-of-the-money). Default "All".
            enable_greeks: Include delta, gamma, theta, vega, IV. Default true.
            price_center: Center strike price. Defaults to last price of underlying.
            strike_interval: Interval between strikes. Default 1.
            risk_free_rate: Custom risk-free rate as decimal (e.g. 0.05 for 5%). Defaults to $IRX.X.
            duration_seconds: Streaming collection window in seconds. Default 5.

        Returns:
            JSON array of option spread objects with greeks, bid/ask, volume,
            open interest, implied volatility, probability calculations per strike/leg.
        """
        params: dict[str, Any] = {
            "strikeProximity": strike_proximity,
            "spreadType": spread_type,
            "optionType": option_type,
            "strikeRange": strike_range,
            "enableGreeks": str(enable_greeks).lower(),
            "strikeInterval": strike_interval,
        }
        if expiration:
            params["expiration"] = expiration
        if price_center is not None:
            params["priceCenter"] = price_center
        if risk_free_rate is not None:
            params["riskFreeRate"] = risk_free_rate
        data = await ts_client.stream_get(
            f"/v3/marketdata/stream/options/chains/{underlying}",
            params=params,
            duration_seconds=duration_seconds,
        )
        return json.dumps(data, indent=2)

    # ── Stream Option Quotes ───────────────────────────────────────

    @mcp.tool()
    async def stream_option_quotes(
        legs: str,
        enable_greeks: bool = True,
        risk_free_rate: float | None = None,
        duration_seconds: int = 5,
    ) -> str:
        """MARKET DATA: Stream live greeks and pricing for specific option contract legs.

        Use this when the user has specific option contracts and wants live pricing,
        greeks, and analytics for those exact legs. For browsing an entire chain,
        use stream_option_chain instead.

        Args:
            legs: JSON string describing legs. Example:
                  [{"Symbol": "MSFT 220916C305", "Ratio": 1}, {"Symbol": "MSFT 220916C310", "Ratio": -1}]
                  Ratio > 0 = buy, Ratio < 0 = sell. Use %20 for spaces in symbols if needed.
            enable_greeks: Include greeks. Default true.
            risk_free_rate: Custom risk-free rate as decimal (e.g. 0.02 for 2%).
            duration_seconds: Streaming collection window in seconds. Default 5.

        Returns:
            JSON array of option spread objects with Delta, Gamma, Theta, Vega,
            ImpliedVolatility, Bid, Ask, Volume, OpenInterest per leg.
        """
        legs_parsed = json.loads(legs) if isinstance(legs, str) else legs
        params: dict[str, Any] = {
            "enableGreeks": str(enable_greeks).lower(),
        }
        # Build leg query params: legs[0].Symbol, legs[0].Ratio, etc.
        for i, leg in enumerate(legs_parsed):
            params[f"legs[{i}].Symbol"] = leg["Symbol"]
            if "Ratio" in leg:
                params[f"legs[{i}].Ratio"] = leg["Ratio"]
        if risk_free_rate is not None:
            params["riskFreeRate"] = risk_free_rate
        data = await ts_client.stream_get(
            "/v3/marketdata/stream/options/quotes",
            params=params,
            duration_seconds=duration_seconds,
        )
        return json.dumps(data, indent=2)

    # ── Stream Quotes ──────────────────────────────────────────────

    @mcp.tool()
    async def stream_quotes(
        symbols: str,
        duration_seconds: int = 5,
    ) -> str:
        """MARKET DATA: Stream real-time live quote updates for one or more symbols.

        Use this when the user wants LIVE/STREAMING/REAL-TIME price updates that
        continuously update, not just a one-time snapshot. For a one-time current
        price check, use get_quote_snapshots instead.

        Example questions:
        - "Stream live prices for AAPL and MSFT"
        - "Show me real-time updates for TSLA for 10 seconds"

        Args:
            symbols: Comma-separated symbols (max 100), e.g. "MSFT,AAPL".
            duration_seconds: How many seconds to collect streaming quote data. Default 5.

        Returns:
            JSON array of quote update objects received during the collection window,
            each with Symbol, Last, Bid, Ask, Volume, NetChange, TradeTime, etc.
        """
        data = await ts_client.stream_get(
            f"/v3/marketdata/stream/quotes/{symbols}",
            duration_seconds=duration_seconds,
        )
        return json.dumps(data, indent=2)

    # ── Market Depth ───────────────────────────────────────────────

    @mcp.tool()
    async def stream_market_depth_quotes(
        symbol: str,
        maxlevels: int = 20,
        duration_seconds: int = 5,
    ) -> str:
        """MARKET DATA: Stream Level 2 / market depth data with individual market maker quotes.

        Use this when the user asks about market depth, Level 2 data, the order book,
        or individual market maker/participant quotes and sizes.

        Args:
            symbol: Single ticker symbol (e.g. "MSFT").
            maxlevels: Maximum number of price levels to show. Default 20.
            duration_seconds: Streaming collection window in seconds. Default 5.

        Returns:
            JSON array of depth snapshots, each with Bids and Asks arrays containing
            Price, Size, Timestamp, and participant info per level.
        """
        params: dict[str, Any] = {"maxlevels": maxlevels}
        data = await ts_client.stream_get(
            f"/v3/marketdata/stream/marketdepth/quotes/{symbol}",
            params=params,
            duration_seconds=duration_seconds,
        )
        return json.dumps(data, indent=2)

    @mcp.tool()
    async def stream_market_depth_aggregates(
        symbol: str,
        maxlevels: int = 20,
        duration_seconds: int = 5,
    ) -> str:
        """MARKET DATA: Stream aggregated market depth (order book) by price level.

        Use this when the user wants aggregated depth — total size at each price level
        across all participants, rather than individual market maker quotes.
        If they want per-participant data, use stream_market_depth_quotes instead.

        Args:
            symbol: Single ticker symbol (e.g. "MSFT").
            maxlevels: Maximum levels. Default 20.
            duration_seconds: Streaming collection window in seconds. Default 5.

        Returns:
            JSON array of aggregated depth snapshots with Bids and Asks arrays,
            each containing Price and aggregated TotalSize per level.
        """
        params: dict[str, Any] = {"maxlevels": maxlevels}
        data = await ts_client.stream_get(
            f"/v3/marketdata/stream/marketdepth/aggregates/{symbol}",
            params=params,
            duration_seconds=duration_seconds,
        )
        return json.dumps(data, indent=2)

    # ── Symbol Suggest ─────────────────────────────────────────────

    @mcp.tool()
    async def suggest_symbols(
        text: str,
        top: int | None = None,
        filter: str | None = None,
    ) -> str:
        """MARKET DATA: Search/autocomplete symbols by partial name, ticker, or company name.

        Use this when the user types a partial name or company and wants matching
        symbol suggestions. Does NOT return option symbols. For structured lookups
        with category filters, use search_symbols instead.

        Example questions:
        - "Find symbols for Apple"
        - "What's the ticker for Tesla?"
        - "Search for Microsoft stock symbol"
        - "Look up ticker for Amazon"

        Args:
            text: Partial symbol or company name to match (e.g. "App", "Tesla", "micro").
            top: Maximum number of results to return.
            filter: OData filter string, e.g. "Category eq 'Stock'" or "Country eq 'United States'".

        Returns:
            JSON array of matching symbols with Name, Description, Category, Exchange, Country, Currency.
        """
        params: dict[str, Any] = {}
        if top is not None:
            params["$top"] = top
        if filter:
            params["$filter"] = filter
        data = await ts_client.get(f"/v2/data/symbols/suggest/{text}", params=params or None)
        return json.dumps(data, indent=2)

    # ── Symbol Search ──────────────────────────────────────────────

    @mcp.tool()
    async def search_symbols(criteria: str) -> str:
        """MARKET DATA: Advanced structured symbol search with criteria filters.

        Use this for precise symbol lookups with structured criteria — by category
        (Stock, Future, Forex, Options), root symbol, strike count, etc. For
        simple name-based autocomplete, use suggest_symbols instead.

        Args:
            criteria: Key/value search criteria string. Examples:
                Stocks:   "C=Stock&N=MSFT" or "C=Stock&Desc=Microsoft"
                Options:  "C=StockOption&R=AAPL&Stk=5"
                Futures:  "C=Future&R=ES&FT=Electronic"
                Forex:    "C=Forex&N=EUR"
                Include expired: add "&Flg=true" for stocks or "&Exp=true" for futures.

        Returns:
            JSON array of matching symbols with Name, Description, Category,
            Exchange, ExpirationDate, OptionType, StrikePrice, etc.
        """
        data = await ts_client.get(f"/v2/data/symbols/search/{criteria}")
        return json.dumps(data, indent=2)

    # ── Stream Tick Bars ───────────────────────────────────────────

    @mcp.tool()
    async def stream_tick_bars(
        symbol: str,
        interval: int,
        bars_back: int,
    ) -> str:
        """MARKET DATA: Stream tick-based bars (bars formed by N trades, not by time intervals).

        Use this when the user asks for tick bars, trade-count bars, or volume-activity
        bars rather than time-based (minute/daily) bars.

        Args:
            symbol: Ticker symbol (e.g. "MSFT").
            interval: Number of ticks per bar (1-64999).
            bars_back: Number of bars to stream back (1-10).

        Returns:
            JSON array of tick bar objects with Close, TotalVolume, TimeStamp, Status.
        """
        data = await ts_client.stream_get(
            f"/v2/stream/tickbars/{symbol}/{interval}/{bars_back}",
            duration_seconds=10,
        )
        return json.dumps(data, indent=2)
