IB-Dashboard/backend/flex_query.py

"""Utilities for Interactive Brokers Flex Web Service integration.

This module provides the first layer of infrastructure for pulling
historical trade data from a saved Flex Query. It intentionally keeps the
surface area small so the app can validate configuration, trigger a sync,
and persist the raw XML responses locally before the data is wired into the
lot engine.
"""
from __future__ import annotations

import os
import time
import urllib.parse
import urllib.request
import xml.etree.ElementTree as ET
from dataclasses import dataclass
from pathlib import Path


FLEX_BASE_URL = "https://ndcdyn.interactivebrokers.com/AccountManagement/FlexWebService"
FLEX_SEND_REQUEST_URL = f"{FLEX_BASE_URL}/SendRequest"
FLEX_GET_STATEMENT_URL = f"{FLEX_BASE_URL}/GetStatement"


@dataclass
class FlexConfig:
    token: str
    query_id: str
    data_dir: Path
    poll_interval_seconds: float = 2.0
    max_polls: int = 30


class FlexQueryError(RuntimeError):
    pass


def load_flex_config() -> FlexConfig | None:
    token = os.environ.get("IB_FLEX_TOKEN", "").strip()
    query_id = os.environ.get("IB_FLEX_QUERY_ID", "").strip()
    if not token or not query_id:
        return None
    data_dir = Path(os.environ.get("IB_FLEX_DATA_DIR", Path(__file__).resolve().parent.parent / "data" / "flex"))
    return FlexConfig(token=token, query_id=query_id, data_dir=data_dir)


def get_flex_status() -> dict:
    config = load_flex_config()
    if config is None:
        return {
            "configured": False,
            "missing": [
                name
                for name in ("IB_FLEX_TOKEN", "IB_FLEX_QUERY_ID")
                if not os.environ.get(name, "").strip()
            ],
        }
    latest_file = None
    latest_mtime = None
    if config.data_dir.exists():
        xml_files = sorted(config.data_dir.glob("flex_*.xml"))
        if xml_files:
            latest = xml_files[-1]
            latest_file = str(latest)
            latest_mtime = latest.stat().st_mtime
    return {
        "configured": True,
        "query_id": config.query_id,
        "data_dir": str(config.data_dir),
        "latest_file": latest_file,
        "latest_timestamp": latest_mtime,
    }


def _http_get(url: str, params: dict[str, str]) -> bytes:
    full_url = f"{url}?{urllib.parse.urlencode(params)}"
    request = urllib.request.Request(
        full_url,
        headers={"User-Agent": "IB-Dashboard/1.0"},
    )
    with urllib.request.urlopen(request, timeout=30) as response:
        return response.read()


def _parse_send_request(xml_bytes: bytes) -> tuple[str, str]:
    root = ET.fromstring(xml_bytes)
    status = (root.findtext("Status") or "").strip()
    if status.lower() != "success":
        raise FlexQueryError(root.findtext("ErrorMessage") or "Flex send-request failed")
    reference_code = (root.findtext("ReferenceCode") or "").strip()
    if not reference_code:
        raise FlexQueryError("Flex send-request returned no ReferenceCode")
    return status, reference_code


def _parse_statement_status(xml_bytes: bytes) -> tuple[bool, str]:
    root = ET.fromstring(xml_bytes)
    if root.tag == "FlexStatementResponse":
        error_message = (root.findtext("ErrorMessage") or "").strip()
        if error_message:
            raise FlexQueryError(error_message)
        return False, ""
    return True, xml_bytes.decode("utf-8", errors="replace")


def fetch_flex_statement() -> dict:
    config = load_flex_config()
    if config is None:
        raise FlexQueryError("Flex Query is not configured")

    config.data_dir.mkdir(parents=True, exist_ok=True)

    _, reference_code = _parse_send_request(_http_get(
        FLEX_SEND_REQUEST_URL,
        {"t": config.token, "q": config.query_id, "v": "3"},
    ))

    statement_xml = None
    for _ in range(config.max_polls):
        xml_bytes = _http_get(
            FLEX_GET_STATEMENT_URL,
            {"t": config.token, "q": reference_code, "v": "3"},
        )
        done, payload = _parse_statement_status(xml_bytes)
        if done:
            statement_xml = payload
            break
        time.sleep(config.poll_interval_seconds)

    if statement_xml is None:
        raise FlexQueryError("Timed out waiting for Flex statement generation")

    timestamp = time.strftime("%Y%m%d_%H%M%S")
    destination = config.data_dir / f"flex_{timestamp}.xml"
    destination.write_text(statement_xml, encoding="utf-8")

    return {
        "saved_to": str(destination),
        "reference_code": reference_code,
        "bytes": len(statement_xml.encode("utf-8")),
    }