schwab-mcp-custom/schwab_scraper/features/accounts_positions/accounts_scraper.py

from __future__ import annotations

import asyncio
import re
from typing import Optional

from ...core import AccountSummary, Envelope, ErrorType, fail, ok
from ...browser.client import connect, new_context, new_page
from ...browser.navigation import goto_with_auth_check
from ...browser.auth import ensure_cookies
from ...core.config import get_playwright_url, load_config

# Use the same URL as transactions feature for consistency and reliability
TRANSACTION_HISTORY_URL = "https://client.schwab.com/app/accounts/history/#/"


def _normalize_account_option(text: str, value: str) -> Optional[AccountSummary]:
    text = text.strip()
    if not text:
        return None

    normalized_text = re.sub(r"\s+", " ", text)

    last4_match = re.search(r"(\d{3,4})", normalized_text.replace(" ", ""))
    last4 = last4_match.group(1)[-4:] if last4_match else None

    type_match = re.search(r"^([A-Za-z&'\- ]+)", normalized_text)
    account_type = (type_match.group(1).strip() if type_match else "Account").replace(" ", "_")

    account_id_candidates = [candidate for candidate in (value.strip(), last4, normalized_text) if candidate]
    account_id = account_id_candidates[0] if account_id_candidates else normalized_text


    label = normalized_text
    is_margin = "margin" in normalized_text.lower()

    return AccountSummary(
        id=account_id,
        label=label,
        type=account_type,
        last4=last4,
        is_margin=is_margin,
    )


async def list_accounts(debug: bool = False) -> Envelope[list[AccountSummary]]:
    """
    Discover accounts from Schwab transaction history page.

    Uses the robust account discovery logic from the transactions feature
    which handles multiple selector patterns and has enhanced reliability.
    """
    cookies = await ensure_cookies()
    if not cookies:
        return fail("Unable to establish Schwab session.", ErrorType.AUTHENTICATION, retryable=False)

    config = load_config()
    playwright_url = get_playwright_url(config)

    playwright = browser = context = page = None
    try:
        playwright, browser = await connect(playwright_url)
        context = await new_context(browser, cookies=cookies)
        page = await new_page(context)

        if not await goto_with_auth_check(page, context, TRANSACTION_HISTORY_URL, debug=debug):
            return fail("Failed to load transaction history for account discovery.", ErrorType.AUTHENTICATION, retryable=True)

        # Allow page to fully load
        await asyncio.sleep(2)

        # Use the robust account discovery from transactions feature
        from ..transactions.scraper import discover_accounts_from_page

        discovered_accounts = await discover_accounts_from_page(page, debug=debug)

        if not discovered_accounts:
            return fail("Account dropdown not found on transaction history page.", ErrorType.PARSING, retryable=True)

        # Convert discovered accounts to AccountSummary objects
        accounts: list[AccountSummary] = []
        seen_ids: set[str] = set()

        for acc in discovered_accounts:
            # Create AccountSummary from discovered account info
            account_id = acc.get('ending', acc.get('label', ''))

            if account_id and account_id not in seen_ids:
                summary = AccountSummary(
                    id=account_id,
                    label=acc.get('label', ''),
                    type=acc.get('type', 'Account'),
                    last4=acc.get('ending', ''),
                    is_margin=False,  # Will be enhanced in future if needed
                )
                accounts.append(summary)
                seen_ids.add(account_id)

        if not accounts:
            return fail("No accounts discovered from Schwab transaction history.", ErrorType.PARSING, retryable=True)

        if debug:
            print(f"DEBUG: Successfully discovered {len(accounts)} accounts:")
            for acc in accounts:
                print(f"DEBUG: - {acc.label} (type: {acc.type}, last4: {acc.last4})")

        return ok(accounts)
    except Exception as exc:
        if debug:
            print(f"DEBUG: Account discovery error: {exc}")
        return fail(str(exc), ErrorType.UNKNOWN, retryable=True)
    finally:
        await _safe_close_page(page)
        await _safe_close_context(context)
        await _safe_close_browser(browser)
        await _safe_stop_playwright(playwright)


async def _safe_close_page(page) -> None:
    if page is None:
        return
    try:
        await page.close()
    except Exception:
        pass


async def _safe_close_context(context) -> None:
    if context is None:
        return
    try:
        await context.close()
    except Exception:
        pass


async def _safe_close_browser(browser) -> None:
    if browser is None:
        return
    try:
        await browser.close()
    except Exception:
        pass


async def _safe_stop_playwright(playwright) -> None:
    if playwright is None:
        return
    try:
        await playwright.stop()
    except Exception:
        pass