broker-sync/broker_sync/providers/trading212.py

from __future__ import annotations

import asyncio
import random
import re
from collections.abc import AsyncIterator
from datetime import datetime
from decimal import Decimal
from pathlib import Path
from typing import Any
from urllib.parse import parse_qs, urlparse

import httpx

from broker_sync.models import Account, Activity, ActivityType
from broker_sync.providers._checkpoint import Checkpoint

_BASE_URL = "https://live.trading212.com"
_ORDERS_PATH = "/api/v0/equity/history/orders"
_PAGE_LIMIT = 50

_SUFFIX_RE = re.compile(r"(?:_US)?(?:[a-z])?_EQ$")

# Retry config for 429 without Retry-After and for 5xx.
_BACKOFF_INITIAL = 10.0
_BACKOFF_CAP = 120.0
_MAX_RETRIES = 3


def _normalise_ticker(raw: str) -> str:
    """Strip T212's exchange-suffix decoration from a ticker."""
    return _SUFFIX_RE.sub("", raw)


def _jitter(base: float) -> float:
    """Return a jittered backoff value in [base/2, base]."""
    return base * (0.5 + random.random() / 2)


class Trading212Error(Exception):
    """Any non-retryable Trading212 API failure."""


class Trading212AuthError(Trading212Error):
    """HTTP 401 from Trading212 — API key is invalid or revoked."""


class Trading212Provider:
    """Concrete Provider for Trading212.

    One instance serves every T212 wrapper the user owns (ISA + Invest)
    — the caller hands over (Account, api_key) pairs and `fetch()` walks
    each account's history in turn. Pagination is cursor-based via the
    `nextPagePath` field in each response; the provider saves the cursor
    that points to the NEXT page only after the current page has been
    fully yielded to the caller, so a crash mid-stream resumes at the
    start of the unfinished page rather than halfway through.
    """

    name = "trading212"

    def __init__(
        self,
        *,
        accounts: list[tuple[Account, str]],
        checkpoint_dir: Path,
        transport: httpx.AsyncBaseTransport | None = None,
    ) -> None:
        self._accounts = accounts
        self._checkpoint_dir = checkpoint_dir
        self._client = httpx.AsyncClient(
            base_url=_BASE_URL,
            timeout=30.0,
            transport=transport,
        )

    def accounts(self) -> list[Account]:
        return [acc for acc, _ in self._accounts]

    async def close(self) -> None:
        await self._client.aclose()

    async def fetch(
        self,
        *,
        since: datetime | None = None,
        before: datetime | None = None,
    ) -> AsyncIterator[Activity]:
        for account, api_key in self._accounts:
            async for activity in self._fetch_account(account, api_key, since, before):
                yield activity

    async def _fetch_account(
        self,
        account: Account,
        api_key: str,
        since: datetime | None,
        before: datetime | None,
    ) -> AsyncIterator[Activity]:
        checkpoint = Checkpoint(
            self._checkpoint_dir,
            provider=self.name,
            account_id=account.id,
        )
        cursor: str | None = None
        while True:
            page = await self._get_page_with_retry(api_key, cursor)
            items = page.get("items", [])
            saw_too_old = False
            for item in items:
                activity = _item_to_activity(item, account)
                if activity is None:
                    continue
                if since is not None and activity.date < since:
                    saw_too_old = True
                    continue
                if before is not None and activity.date >= before:
                    continue
                yield activity

            next_path = page.get("nextPagePath")
            if isinstance(next_path, str) and next_path:
                checkpoint.save(next_path)
            if not isinstance(next_path, str) or not next_path or saw_too_old:
                return
            cursor = _extract_cursor(next_path)

    async def _get_page_with_retry(
        self,
        api_key: str,
        cursor: str | None,
    ) -> dict[str, Any]:
        attempts = 0
        backoff = _BACKOFF_INITIAL
        while True:
            resp = await self._request_page(api_key, cursor)
            if resp.status_code == 200:
                raw = resp.json()
                assert isinstance(raw, dict)
                return raw
            if resp.status_code == 401:
                raise Trading212AuthError("Trading212 rejected API key (HTTP 401)")
            retryable = resp.status_code == 429 or 500 <= resp.status_code < 600
            if not retryable:
                raise Trading212Error(f"Trading212 /orders HTTP {resp.status_code}: {resp.text}")
            if attempts >= _MAX_RETRIES:
                raise Trading212Error(
                    f"Trading212 /orders HTTP {resp.status_code} after {attempts} retries")
            sleep_for = _sleep_after(resp, backoff)
            await asyncio.sleep(sleep_for)
            attempts += 1
            backoff = min(backoff * 2, _BACKOFF_CAP)

    async def _request_page(self, api_key: str, cursor: str | None) -> httpx.Response:
        params: dict[str, str | int] = {"limit": _PAGE_LIMIT}
        if cursor is not None:
            params["cursor"] = cursor
        return await self._client.get(
            _ORDERS_PATH,
            params=params,
            headers={"Authorization": api_key},
        )


def _sleep_after(resp: httpx.Response, backoff: float) -> float:
    if resp.status_code == 429:
        retry_after = resp.headers.get("Retry-After")
        if retry_after is not None:
            try:
                return float(retry_after)
            except ValueError:
                pass
    return _jitter(backoff)


def _extract_cursor(next_page_path: str) -> str | None:
    """Pull the `cursor` query param out of a nextPagePath URL fragment."""
    parsed = urlparse(next_page_path)
    q = parse_qs(parsed.query)
    cursor_values = q.get("cursor")
    if not cursor_values:
        return None
    return cursor_values[0]


def _item_to_activity(item: dict[str, Any], account: Account) -> Activity | None:
    fill = item.get("fill")
    if fill is None:
        return None
    order = item["order"]
    quantity_raw = Decimal(str(fill["quantity"]))
    activity_type = ActivityType.BUY if quantity_raw > 0 else ActivityType.SELL
    return Activity(
        external_id=f"t212:fill:{fill['id']}",
        account_id=account.id,
        account_type=account.account_type,
        date=_parse_iso(fill["filledAt"]),
        activity_type=activity_type,
        symbol=_normalise_ticker(order["ticker"]),
        quantity=abs(quantity_raw),
        unit_price=Decimal(str(fill["price"])),
        currency=order["currency"],
    )


def _parse_iso(ts: str) -> datetime:
    # T212 always emits `...Z`; datetime.fromisoformat handles `+00:00`.
    return datetime.fromisoformat(ts.replace("Z", "+00:00"))