schwab-mcp-custom/schwab_scraper/unified_api.py

"""Unified Schwab data surface with envelope-based async endpoints."""

from __future__ import annotations

from typing import Optional

from .core import AccountOverview, AccountSummary, Envelope, MorningstarData, PortfolioSnapshot, Position, EquityPhase1Data
from .core.models import TransactionData
from .core import ErrorType, fail
from .features.accounts_positions.accounts_scraper import list_accounts as _list_accounts
from .features.accounts_positions.overview_scraper import get_account_overview as _get_account_overview
from .features.accounts_positions.positions_scraper import get_positions as _get_positions
from .features.accounts_positions.portfolio_scraper import get_portfolio_snapshot as _get_portfolio_snapshot
from .features.equity.service import get_morningstar_data as _get_morningstar_data, get_equity_phase1_data as _get_equity_phase1_data
from .features.transactions.service import (
    get_transaction_history as _get_transaction_history,
    get_transaction_history_enhanced as _get_transaction_history_enhanced,
    list_available_accounts as _list_available_accounts,
)
from .browser.session import get_session_status as _get_session_status_impl
from .browser.session import refresh_session as _refresh_session_impl
from .browser.session import set_cookies_from_file as _set_cookies_impl
from .browser.session import export_cookies as _export_cookies_impl


async def get_session_status(debug: bool = False) -> Envelope[dict]:
    try:
        status = await _get_session_status_impl(debug=debug)
        return status  # already returns envelope
    except Exception as exc:
        return fail(str(exc), ErrorType.UNKNOWN, retryable=True)


async def refresh_session(debug: bool = False) -> Envelope[None]:
    try:
        return await _refresh_session_impl(debug=debug)
    except Exception as exc:
        return fail(str(exc), ErrorType.UNKNOWN, retryable=True)


async def set_cookies(cookies_path: str, debug: bool = False) -> Envelope[None]:
    try:
        return await _set_cookies_impl(cookies_path, debug=debug)
    except Exception as exc:
        return fail(str(exc), ErrorType.UNKNOWN, retryable=False)


async def export_cookies(cookies_path: str, debug: bool = False) -> Envelope[None]:
    try:
        return await _export_cookies_impl(cookies_path, debug=debug)
    except Exception as exc:
        return fail(str(exc), ErrorType.UNKNOWN, retryable=False)


async def list_accounts(debug: bool = False) -> Envelope[list[AccountSummary]]:
    envelope = await _list_accounts(debug=debug)
    if not envelope["success"]:
        return envelope
    data = envelope["data"] or []
    summaries: list[AccountSummary] = []
    for item in data:
        if isinstance(item, AccountSummary):
            summaries.append(item)
        else:
            summaries.append(AccountSummary(**item))
    return {
        "success": True,
        "data": summaries,
        "error": None,
        "error_type": None,
        "retryable": False,
    }


async def get_account_overview(
    account: AccountSummary | str | None = None,
    *,
    debug: bool = False,
) -> Envelope[AccountOverview]:
    if isinstance(account, dict):
        account = AccountSummary(**account)
    return await _get_account_overview(account=account, debug=debug)


async def get_positions(
    account: AccountSummary | str | None = None,
    *,
    include_non_equity: bool = False,
    debug: bool = False,
) -> Envelope[list[Position]]:
    if isinstance(account, dict):
        account = AccountSummary(**account)
    return await _get_positions(account=account, include_non_equity=include_non_equity, debug=debug)


async def get_portfolio_snapshot(
    account: AccountSummary | str | None = None,
    *,
    aggregate_by_symbol: bool = True,
    include_non_equity: bool = False,
    debug: bool = False,
) -> Envelope[PortfolioSnapshot]:
    if isinstance(account, dict):
        account = AccountSummary(**account)
    return await _get_portfolio_snapshot(
        account=account,
        aggregate_by_symbol=aggregate_by_symbol,
        include_non_equity=include_non_equity,
        debug=debug,
    )


async def get_morningstar_data(ticker: str, debug: bool = False) -> Envelope[MorningstarData]:
    return await _get_morningstar_data(ticker, debug=debug)


async def get_equity_phase1_data(ticker: str, debug: bool = False) -> Envelope[EquityPhase1Data]:
    """Get Phase 1 enhanced equity data for a ticker.
    
    Extracts:
    - Quote/Price Data (symbol bar)
    - Enhanced Dividend Information (forward-looking dates)
    - Core Earnings Metrics (EPS, forecasts)
    - Basic Valuation Ratios (P/E, Forward P/E, PEG)
    - Calculated Metrics (payout ratio)
    
    Args:
        ticker: Stock ticker symbol
        debug: Enable debug logging
        
    Returns:
        Envelope containing EquityPhase1Data or error
    """
    return await _get_equity_phase1_data(ticker, debug=debug)


async def list_available_accounts(debug: bool = False) -> Envelope[list[dict]]:
    return await _list_available_accounts(debug=debug)


async def get_transaction_history(
    account: Optional[str] = None,
    start_date: Optional[str] = None,
    end_date: Optional[str] = None,
    time_period: Optional[str] = None,
    debug: bool = False,
) -> Envelope[TransactionData]:
    envelope = await _get_transaction_history(
        account=account,
        start_date=start_date,
        end_date=end_date,
        time_period=time_period,
        debug=debug,
    )
    return envelope


async def get_transaction_history_enhanced(
    account: Optional[str] = None,
    start_date: Optional[str] = None,
    end_date: Optional[str] = None,
    time_period: Optional[str] = None,
    debug: bool = False,
) -> Envelope[TransactionData]:
    envelope = await _get_transaction_history_enhanced(
        account=account,
        start_date=start_date,
        end_date=end_date,
        time_period=time_period,
        debug=debug,
    )
    return envelope

__all__ = [
    "get_session_status",
    "refresh_session",
    "set_cookies",
    "export_cookies",
    "list_accounts",
    "get_account_overview",
    "get_positions",
    "get_portfolio_snapshot",
    "get_morningstar_data",
    "get_equity_phase1_data",
    "list_available_accounts",
    "get_transaction_history",
    "get_transaction_history_enhanced",
]