"""
float_probe.py — one-off probe: does TWS hand us shares outstanding and free float?

Answers the open question from 2026-09-11: the IBKR *connector* has no fundamentals field
(33 snapshot fields, none of them shares-related), but TWS's own API has reqFundamentalData.
If ReportSnapshot comes back, float can come off the same connection as everything else and
Alpha Vantage's daily call limit stops mattering.

Read-only. Touches nothing — no quotes.json, no roster.json, no receiver state.
Uses clientId 31 so it cannot collide with the feed's 17.

Run:   python float_probe.py
       python float_probe.py --symbols RKLB,MSTR,INTC,NVDA
       python float_probe.py --save-xml        (dump the raw XML beside this file for inspection)

Cross-check: Alpha Vantage COMPANY_OVERVIEW gave RKLB float 556,861,000 / shares out 598,460,000
on 2026-09-11. If TWS agrees within a percent or so, the TWS route is good.
"""
import argparse, os, sys, xml.etree.ElementTree as ET

import asyncio
asyncio.set_event_loop(asyncio.new_event_loop())   # Python 3.12+: eventkit expects a loop at import time
try:
    from ib_insync import IB, Stock
except ImportError:
    sys.exit("ib_insync not installed — run: pip install ib_insync")

HERE = os.path.dirname(os.path.abspath(__file__))


def fmt(n):
    return "—" if n is None else f"{n:,.0f}"


def parse_snapshot(xml_text):
    """ReportSnapshot: CoGeneralInfo/SharesOut holds the share count, with free float as an attribute.
    Attribute spelling has varied across Reuters/Refinitiv revisions, so try the known ones."""
    out = {"shares_out": None, "float": None, "as_of": None, "name": None}
    try:
        root = ET.fromstring(xml_text)
    except ET.ParseError as e:
        out["error"] = f"XML parse failed: {e}"
        return out

    for el in root.iter():
        tag = el.tag.split("}")[-1]
        if tag == "SharesOut":
            try:
                out["shares_out"] = float(el.text) if el.text and el.text.strip() else None
            except ValueError:
                pass
            out["as_of"] = el.get("Date") or out["as_of"]
            for key in ("TotalFloat", "Float", "FreeFloat"):
                v = el.get(key)
                if v:
                    try:
                        out["float"] = float(v)
                        break
                    except ValueError:
                        pass
        elif tag == "CoID" and el.get("Type") == "CompanyName":
            out["name"] = (el.text or "").strip()
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--symbols", default="RKLB,MSTR,INTC,AAOI",
                    help="comma-separated tickers to probe")
    ap.add_argument("--tws", type=int, default=7496, help="TWS API socket port")
    ap.add_argument("--client-id", type=int, default=31, help="must differ from the feed's 17")
    ap.add_argument("--save-xml", action="store_true", help="write the raw XML beside this script")
    a = ap.parse_args()

    syms = [s.strip().upper() for s in a.symbols.split(",") if s.strip()]

    ib = IB()
    try:
        ib.connect("127.0.0.1", a.tws, clientId=a.client_id, readonly=True, timeout=10)
    except Exception as e:
        sys.exit(f"could not connect to TWS on port {a.tws}: {e}\n"
                 f"Is TWS running with the API enabled? (Global Config -> API -> Settings)")
    print(f"connected to TWS on {a.tws} (clientId {a.client_id})\n")

    rows, any_ok = [], False
    for sym in syms:
        c = Stock(sym, "SMART", "USD")
        try:
            [c] = ib.qualifyContracts(c)
        except Exception as e:
            rows.append((sym, None, None, None, f"could not qualify contract: {e}"))
            continue

        try:
            xml_text = ib.reqFundamentalData(c, "ReportSnapshot")
        except Exception as e:
            rows.append((sym, None, None, None, f"reqFundamentalData raised: {e}"))
            continue

        if not xml_text:
            rows.append((sym, None, None, None, "empty response — most likely no fundamental-data subscription"))
            continue

        if a.save_xml:
            p = os.path.join(HERE, f"fundamentals_{sym}.xml")
            with open(p, "w", encoding="utf-8") as f:
                f.write(xml_text)
            print(f"  wrote {p}  ({len(xml_text):,} chars)")

        d = parse_snapshot(xml_text)
        if d.get("error"):
            rows.append((sym, None, None, None, d["error"]))
            continue
        if d["shares_out"] is None and d["float"] is None:
            rows.append((sym, None, None, d["as_of"],
                         f"XML returned ({len(xml_text):,} chars) but no SharesOut element — "
                         f"rerun with --save-xml and check the structure"))
            continue

        any_ok = True
        rows.append((sym, d["shares_out"], d["float"], d["as_of"], d["name"] or ""))

    ib.disconnect()

    w = max(6, *(len(r[0]) for r in rows))
    print(f"\n{'SYM'.ljust(w)}  {'SHARES OUT':>16}  {'FLOAT':>16}  {'AS OF':>10}  NOTE")
    print("-" * (w + 70))
    for sym, so, fl, asof, note in rows:
        print(f"{sym.ljust(w)}  {fmt(so):>16}  {fmt(fl):>16}  {(asof or '—'):>10}  {note}")

    print()
    if any_ok:
        print("TWS carries the data. Float can come off this connection —")
        print("no Alpha Vantage call budget needed. Cross-check RKLB against 556,861,000 (AV, 2026-09-11).")
    else:
        print("TWS gave nothing usable. Fundamentals need the Reuters/Refinitiv Worldwide Fundamentals")
        print("subscription in Account Management; without it, float stays on Alpha Vantage.")


if __name__ == "__main__":
    main()