"""
TradeStation Brokerage MCP Tools.

Provides tools for accessing account information, balances, positions,
orders, and streaming account updates.

Brokerage tools auto-resolve account IDs when not provided — if the user
asks "What's my balance?" without specifying an account, the tool internally
fetches all accounts first.
"""

import json
import logging
from typing import Any

logger = logging.getLogger(__name__)


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

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

    # ── Helper: Auto-resolve account IDs ───────────────────────────

    async def _resolve_accounts(accounts: str | None) -> str:
        """If accounts param is empty, fetch all accounts and build the comma-separated ID list."""
        if accounts and accounts.strip():
            return accounts.strip()
        logger.info("No account IDs provided, fetching all accounts...")
        data = await ts_client.get("/v3/brokerage/accounts")
        account_list = data.get("Accounts", [])
        if not account_list:
            raise ValueError("No brokerage accounts found for this user.")
        ids = ",".join(a["AccountID"] for a in account_list)
        logger.info("Auto-resolved accounts: %s", ids)
        return ids

    # ── Get Accounts ───────────────────────────────────────────────

    @mcp.tool()
    async def get_accounts() -> str:
        """BROKERAGE: List all brokerage accounts for the authenticated user.

        Use this when the user asks about their accounts, account list,
        account numbers, or account types. This is often the first step
        before checking balances, positions, or placing trades.

        Example questions this tool answers:
        - "What accounts do I have?"
        - "Show my account numbers"
        - "List my trading accounts"
        - "What type of accounts do I have?"

        DO NOT use this for stock prices or market data — those are
        MARKET DATA tools (get_quote_snapshots, get_bars).

        Returns:
            JSON with Accounts array, each containing AccountID, AccountType
            (Cash, Margin, Futures, DVP), Status, and alias info.
        """
        data = await ts_client.get("/v3/brokerage/accounts")
        return json.dumps(data, indent=2)

    # ── Get Balances ───────────────────────────────────────────────

    @mcp.tool()
    async def get_balances(accounts: str | None = None) -> str:
        """BROKERAGE: Get current account balances — cash, equity, buying power, margin details.

        Use this when the user asks about their ACCOUNT BALANCE, how much money
        they have, cash available, buying power, equity value, margin usage,
        or overall account value.

        Example questions this tool answers:
        - "What's my account balance?"
        - "How much cash do I have?"
        - "What's my buying power?"
        - "How much equity is in my account?"
        - "What's my margin balance?"

        DO NOT confuse with stock price questions like "What's the price of MSFT?"
        — those should use get_quote_snapshots (MARKET DATA).

        Args:
            accounts: Comma-separated Account IDs (e.g. "61999124,68910124"). Max 25.
                      If omitted, automatically fetches all accounts for the user.

        Returns:
            JSON with Balances array containing CashBalance, Equity, MarketValue,
            BuyingPower, DayTradingBuyingPower, RealizedProfitLoss,
            UnrealizedProfitLoss, MaintenanceRate, and more per account.
        """
        resolved = await _resolve_accounts(accounts)
        data = await ts_client.get(f"/v3/brokerage/accounts/{resolved}/balances")
        return json.dumps(data, indent=2)

    # ── Get BOD Balances ───────────────────────────────────────────

    @mcp.tool()
    async def get_balances_bod(accounts: str | None = None) -> str:
        """BROKERAGE: Get beginning-of-day (BOD) account balances — the snapshot from market open.

        Use this when the user asks what their balance was at the start of the trading
        day, or wants to compare current balance vs opening balance.

        Example questions:
        - "What was my balance at the start of the day?"
        - "Show my beginning-of-day balances"

        Args:
            accounts: Comma-separated Account IDs. Max 25. Auto-fetched if omitted.

        Returns:
            JSON with BODBalances array containing opening cash, equity, and margin values.
        """
        resolved = await _resolve_accounts(accounts)
        data = await ts_client.get(f"/v3/brokerage/accounts/{resolved}/bodbalances")
        return json.dumps(data, indent=2)

    # ── Get Positions ──────────────────────────────────────────────

    @mcp.tool()
    async def get_positions(
        accounts: str | None = None,
        symbol: str | None = None,
    ) -> str:
        """BROKERAGE: Get current open positions (holdings) in brokerage accounts.

        Use this when the user asks about their holdings, portfolio, what stocks
        they own, current positions, unrealized P&L, or market value of their holdings.

        Example questions this tool answers:
        - "What positions do I have?"
        - "Show my portfolio"
        - "What stocks do I own?"
        - "What's my unrealized P&L?"
        - "Am I holding any MSFT?"
        - "Do I have any open positions?"

        DO NOT use this for stock prices — "How much is MSFT worth?" about the stock
        itself should use get_quote_snapshots. Use get_positions when asking about
        the user's OWN holdings.

        Args:
            accounts: Comma-separated Account IDs. Max 25. Auto-fetched if omitted.
            symbol: Optional symbol filter. Comma-separated, supports wildcards:
                    "MSFT" — just MSFT shares
                    "MSFT *" — all MSFT options
                    "MSFT,MSFT *" — MSFT shares and all its options
                    "* 2503*C*" — all call options expiring March 2025

        Returns:
            JSON with Positions array containing Symbol, Quantity, AveragePrice,
            Last, MarketValue, TotalCost, UnrealizedProfitLoss,
            UnrealizedProfitLossPercent, AssetType, LongShort per position.
        """
        resolved = await _resolve_accounts(accounts)
        params: dict[str, Any] = {}
        if symbol:
            params["symbol"] = symbol
        data = await ts_client.get(f"/v3/brokerage/accounts/{resolved}/positions", params=params or None)
        return json.dumps(data, indent=2)

    # ── Get Orders ─────────────────────────────────────────────────

    @mcp.tool()
    async def get_orders(
        accounts: str | None = None,
        page_size: int = 600,
        next_token: str | None = None,
    ) -> str:
        """BROKERAGE: Get today's orders and currently open/pending orders.

        Use this when the user asks about their current orders, pending orders,
        open orders, today's order activity, or recent order status.

        Example questions:
        - "Do I have any open orders?"
        - "What orders did I place today?"
        - "Show my pending orders"
        - "What's the status of my orders?"

        For HISTORICAL orders (past days/weeks), use get_historical_orders instead.

        Args:
            accounts: Comma-separated Account IDs. Max 25. Auto-fetched if omitted.
            page_size: Results per page (1-600). Default 600.
            next_token: Pagination token for subsequent pages (from prior response).

        Returns:
            JSON with Orders array containing OrderID, Symbol, Quantity, OrderType,
            Status, StatusDescription, FilledPrice, Legs, Duration,
            OpenedDateTime, and more per order. Includes NextToken for pagination.
        """
        resolved = await _resolve_accounts(accounts)
        params: dict[str, Any] = {"pageSize": page_size}
        if next_token:
            params["nextToken"] = next_token
        data = await ts_client.get(f"/v3/brokerage/accounts/{resolved}/orders", params=params)
        return json.dumps(data, indent=2)

    # ── Get Orders by ID ───────────────────────────────────────────

    @mcp.tool()
    async def get_orders_by_id(
        order_ids: str,
        accounts: str | None = None,
    ) -> str:
        """BROKERAGE: Get specific orders by their Order IDs (today's and open orders only).

        Use this when the user asks about a specific order using its order number/ID.

        Args:
            order_ids: Comma-separated Order IDs (e.g. "123456789,286179863"). Max 50.
            accounts: Comma-separated Account IDs. Auto-fetched if omitted.

        Returns:
            JSON with matching Orders and any Errors.
        """
        resolved = await _resolve_accounts(accounts)
        data = await ts_client.get(f"/v3/brokerage/accounts/{resolved}/orders/{order_ids}")
        return json.dumps(data, indent=2)

    # ── Get Historical Orders ──────────────────────────────────────

    @mcp.tool()
    async def get_historical_orders(
        since: str,
        accounts: str | None = None,
        page_size: int = 600,
        next_token: str | None = None,
    ) -> str:
        """BROKERAGE: Get historical (past) orders — filled, cancelled, expired. Max 90 days back.

        Use this when the user asks about past trades, order history, filled orders,
        trade history beyond today, or wants to review previous trading activity.

        Example questions:
        - "Show my trades from last week"
        - "What did I trade last month?"
        - "Order history for January"
        - "Show my filled orders from the past 30 days"

        For TODAY'S orders and currently open orders, use get_orders instead.

        Args:
            since: Required. Start date (e.g. "2025-01-13" or "01-13-2025"). Max 90 days back.
            accounts: Comma-separated Account IDs. Max 25. Auto-fetched if omitted.
            page_size: Results per page (1-600). Default 600.
            next_token: Pagination token.

        Returns:
            JSON with Orders array sorted by time closed (descending), with NextToken for pagination.
        """
        resolved = await _resolve_accounts(accounts)
        params: dict[str, Any] = {"since": since, "pageSize": page_size}
        if next_token:
            params["nextToken"] = next_token
        data = await ts_client.get(f"/v3/brokerage/accounts/{resolved}/historicalorders", params=params)
        return json.dumps(data, indent=2)

    # ── Get Historical Orders by ID ────────────────────────────────

    @mcp.tool()
    async def get_historical_orders_by_id(
        order_ids: str,
        since: str,
        accounts: str | None = None,
    ) -> str:
        """BROKERAGE: Get specific historical orders by Order ID.

        Use this when the user asks about a specific past order by its ID number.

        Args:
            order_ids: Comma-separated Order IDs. Max 50.
            since: Required. Date string (e.g. "2025-01-13"). Max 90 days back.
            accounts: Comma-separated Account IDs. Auto-fetched if omitted.

        Returns:
            JSON with matching historical Orders.
        """
        resolved = await _resolve_accounts(accounts)
        params: dict[str, Any] = {"since": since}
        data = await ts_client.get(
            f"/v3/brokerage/accounts/{resolved}/historicalorders/{order_ids}",
            params=params,
        )
        return json.dumps(data, indent=2)

    # ── Stream Orders ──────────────────────────────────────────────

    @mcp.tool()
    async def stream_orders(
        accounts: str | None = None,
        duration_seconds: int = 10,
    ) -> str:
        """BROKERAGE: Stream real-time order status updates for accounts.

        Use this when the user wants live order fill notifications, real-time
        order status changes, or wants to monitor orders as they execute.

        Args:
            accounts: Comma-separated Account IDs. Max 25. Auto-fetched if omitted.
            duration_seconds: How long to collect streaming order updates. Default 10.

        Returns:
            JSON array of order update events received during the collection window.
        """
        resolved = await _resolve_accounts(accounts)
        data = await ts_client.stream_get(
            f"/v3/brokerage/stream/accounts/{resolved}/orders",
            duration_seconds=duration_seconds,
        )
        return json.dumps(data, indent=2)

    # ── Stream Orders by ID ────────────────────────────────────────

    @mcp.tool()
    async def stream_orders_by_id(
        order_ids: str,
        accounts: str | None = None,
        duration_seconds: int = 10,
    ) -> str:
        """BROKERAGE: Stream real-time updates for specific orders by Order ID.

        Use this when the user wants to monitor specific orders as they fill or change status.

        Args:
            order_ids: Comma-separated Order IDs (e.g. "812767578,812941051"). Max 50.
            accounts: Comma-separated Account IDs. Auto-fetched if omitted.
            duration_seconds: Collection window in seconds. Default 10.

        Returns:
            JSON array of order update events for the specified orders.
        """
        resolved = await _resolve_accounts(accounts)
        data = await ts_client.stream_get(
            f"/v3/brokerage/stream/accounts/{resolved}/orders/{order_ids}",
            duration_seconds=duration_seconds,
        )
        return json.dumps(data, indent=2)

    # ── Stream Positions ───────────────────────────────────────────

    @mcp.tool()
    async def stream_positions(
        accounts: str | None = None,
        changes: bool = False,
        duration_seconds: int = 10,
    ) -> str:
        """BROKERAGE: Stream real-time position and holdings updates.

        Use this when the user wants live P&L updates, real-time position value
        changes, or wants to watch their portfolio update in real-time.

        Args:
            accounts: Comma-separated Account IDs. Max 25. Auto-fetched if omitted.
            changes: If true, after the initial full snapshot, only changed fields
                     are streamed (more efficient). Default false.
            duration_seconds: Collection window in seconds. Default 10.

        Returns:
            JSON array of position snapshots/updates received during the window.
            Each includes Symbol, Quantity, MarketValue, UnrealizedProfitLoss, Last, etc.
        """
        resolved = await _resolve_accounts(accounts)
        params: dict[str, Any] = {}
        if changes:
            params["changes"] = "true"
        data = await ts_client.stream_get(
            f"/v3/brokerage/stream/accounts/{resolved}/positions",
            params=params or None,
            duration_seconds=duration_seconds,
        )
        return json.dumps(data, indent=2)
