"""Collecteur continu Polymarket "Up or Down" crypto -> SQLite.

Polymarket sert uniquement de capteur (aucune mise). On enregistre :
  - live    : l'etat du carnet du jeton Up (cote, spread, profondeur) de la fenetre en cours ET de la
              suivante, a chaque etape de leur vie, plus le prix spot Binance au meme instant ;
  - apres cloture : le resultat, le prix de reference et le flux de transactions agrege par 10 s ;
  - --backfill-days N : l'historique des N derniers jours (cote 1 min + flux) pour analyser tout de suite.
Ce qui n'est PAS stocke ici (bougies, indicateurs techniques) se recupere a posteriori dans analyze.py.

Usage :
  python collector.py
  python collector.py --backfill-days 7 --assets btc,eth
  python collector.py --tfs 5m,15m,1h,4h,1d --db D:/data/pm.db
Ctrl+C pour arreter ; relancer reprend sur la meme base.
"""
import argparse
import json
import logging
import os
import threading
import time
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor

from common import (BINANCE, CLOB, DATA, GAMMA, SYMBOLS, TF_SAMPLE, TF_SECONDS, connect, default_db_path, single_instance,
                    http_json, parse_market, series_slug, ts_iso, upsert_market)

log = logging.getLogger("collector")
SPOT_EVERY = 5
DISCOVER_EVERY = 60
MAX_TRADE_OFFSET = 10000


def safe(fn):
    def wrapped(*a):
        try:
            return fn(*a)
        except Exception as e:  # une erreur reseau ne doit jamais arreter la collecte
            log.warning("%s%s: %s", fn.__name__, a[:2] if a else "", e)
            return None
    return wrapped


# ---------------------------------------------------------------- decouverte des fenetres

def discover(con, pairs):
    now = int(time.time())

    @safe
    def one(pair):
        asset, tf = pair
        d = TF_SECONDS[tf]
        evs = http_json(f"{GAMMA}/events", {"series_slug": series_slug(asset, tf), "closed": "false",
                                            "end_date_min": ts_iso(now), "end_date_max": ts_iso(now + 2 * d + 120),
                                            "limit": 10})
        rows = []
        for ev in evs:
            for m in ev.get("markets", []):
                m["_eventMetadata"] = ev.get("eventMetadata")
                r = parse_market(m, asset, tf)
                # fenetre en cours + suivante (la suivante donne l'etape "avant ouverture")
                if r and r["end_ts"] > now and r["start_ts"] <= now + d:
                    rows.append(r)
        return rows

    with ThreadPoolExecutor(6) as ex:
        found = [r for rows in ex.map(one, pairs) if rows for r in rows]
    for r in found:
        r["id"] = upsert_market(con, r)
    con.commit()
    return found


# ---------------------------------------------------------------- carnet + spot

def book_row(b, ts, mkt):
    bids = [(float(x["price"]), float(x["size"])) for x in b.get("bids", [])]
    asks = [(float(x["price"]), float(x["size"])) for x in b.get("asks", [])]
    bid = max((p for p, _ in bids), default=None)
    ask = min((p for p, _ in asks), default=None)
    mid = (bid + ask) / 2 if bid is not None and ask is not None else None
    ref = mid if mid is not None else (bid if bid is not None else ask)

    def depth(levels, side, width):
        if ref is None:
            return None
        if side == "bid":
            return sum(s for p, s in levels if p >= ref - width - 1e-9)
        return sum(s for p, s in levels if p <= ref + width + 1e-9)

    last = b.get("last_trade_price")
    return (ts, mkt, "live", bid, ask, mid, float(last) if last else None,
            sum(s for p, s in bids if p == bid) if bid is not None else None,
            sum(s for p, s in asks if p == ask) if ask is not None else None,
            depth(bids, "bid", 0.05), depth(asks, "ask", 0.05),
            depth(bids, "bid", 0.10), depth(asks, "ask", 0.10),
            len(bids), len(asks), int(b.get("timestamp") or 0))


def sample_books(con, markets):
    by_token = {m["up_token"]: m["id"] for m in markets}
    tokens = list(by_token)
    rows = []
    for i in range(0, len(tokens), 40):
        chunk = tokens[i:i + 40]
        t0 = time.time()
        books = http_json(f"{CLOB}/books", body=[{"token_id": t} for t in chunk])
        ts = (t0 + time.time()) / 2
        for b in books:
            mkt = by_token.get(b.get("asset_id"))
            if mkt:
                rows.append(book_row(b, ts, mkt))
    con.executemany("INSERT INTO book VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)", rows)
    con.commit()
    return len(rows)


def sample_spot(con, symbols):
    t0 = time.time()
    data = http_json(f"{BINANCE}/api/v3/ticker/bookTicker", {"symbols": json.dumps(symbols, separators=(",", ":"))})
    ts = (t0 + time.time()) / 2
    rows = [(ts, d["symbol"], float(d["bidPrice"]), float(d["askPrice"]),
             (float(d["bidPrice"]) + float(d["askPrice"])) / 2) for d in data]
    con.executemany("INSERT INTO spot VALUES(?,?,?,?,?)", rows)
    con.commit()
    return len(rows)


# ---------------------------------------------------------------- apres cloture : resultat + flux

def resolve_pending(con):
    now = int(time.time())
    rows = con.execute("""SELECT condition_id, slug, asset, tf, end_ts FROM markets
                          WHERE outcome IS NULL AND end_ts < ? ORDER BY end_ts LIMIT 100""", (now - 60,)).fetchall()
    if not rows:
        return 0

    @safe
    def one(row):
        cid, slug, asset, tf, end_ts = row
        ms = http_json(f"{GAMMA}/markets", {"slug": slug, "closed": "true"})
        return row, (parse_market(ms[0], asset, tf) if ms else None)

    with ThreadPoolExecutor(4) as ex:
        res = [r for r in ex.map(one, rows) if r]
    for (cid, slug, asset, tf, end_ts), r in res:
        if r and r["outcome"] is not None:
            upsert_market(con, r)
        elif end_ts < now - 86400:  # jamais resolu apres 24 h : on abandonne
            con.execute("UPDATE markets SET outcome = -1 WHERE condition_id = ?", (cid,))
    con.commit()
    return len(res)


def trade_key(t):
    return (t.get("transactionHash"), t.get("asset"), t.get("side"), t.get("size"), t.get("price"),
            t.get("proxyWallet"), t.get("timestamp"))


def accumulate(agg, t, up_token):
    """Ajoute une transaction a l'agregat par tranche de 10 s ; retourne la tranche touchee."""
    try:
        sz, px, ts = float(t["size"]), float(t["price"]), int(t["timestamp"])
    except (KeyError, TypeError, ValueError):
        return None
    is_up = t.get("asset") == up_token
    sign = 1 if (t.get("side") == "BUY") == is_up else -1
    b = ts // 10 * 10
    a = agg.setdefault(b, [0.0, 0.0, 0, set(), 0.0, 0.0])
    a[0] += sign * sz * px
    a[1] += sz * px
    a[2] += 1
    a[3].add(t.get("proxyWallet"))
    a[4] += (px if is_up else 1 - px) * sz
    a[5] += sz
    return b


def flow_rows(mkt, agg, buckets):
    return [(mkt, b, agg[b][0], agg[b][1], agg[b][2], len(agg[b][3]), agg[b][4] / agg[b][5] if agg[b][5] else None)
            for b in buckets]


def fetch_flow(mkt, cid, up_token):
    """Apres cloture : toutes les transactions du marche (remplace l'agregat partiel du direct)."""
    trades, seen, offset, truncated = [], set(), 0, False
    while True:
        batch = http_json(f"{DATA}/trades", {"market": cid, "limit": 500, "offset": offset})
        if not isinstance(batch, list):
            break
        for t in batch:
            k = trade_key(t)
            if k not in seen:
                seen.add(k)
                trades.append(t)
        if len(batch) < 500:
            break
        offset += 500
        if offset >= MAX_TRADE_OFFSET:
            truncated = True
            break
    agg = {}
    for t in trades:
        accumulate(agg, t, up_token)
    return flow_rows(mkt, agg, list(agg)), truncated


def live_flow_loop(db, get_tracked, stop, every=10):
    """En direct : flux des fenetres en cours, pour que le moteur de signal dispose des memes facteurs
    qu'a l'entrainement. Peut sous-compter en cas de rafale (> 500 transactions / 10 s) ; l'agregat
    complet recalcule apres cloture le remplace."""
    con = connect(db)
    seen, aggs = defaultdict(set), defaultdict(dict)
    while not stop.wait(every):
        now = time.time()
        live = [m for m in get_tracked() if m["start_ts"] - TF_SECONDS[m["tf"]] <= now < m["end_ts"]]
        by_cid = {m["condition_id"]: (m["id"], m["up_token"]) for m in live if m.get("id")}
        touched = defaultdict(set)
        cids = list(by_cid)
        try:
            for i in range(0, len(cids), 4):
                chunk = cids[i:i + 4]
                batch = http_json(f"{DATA}/trades", {"market": ",".join(chunk), "limit": 500})
                for t in batch if isinstance(batch, list) else []:
                    cid = t.get("conditionId")
                    if cid not in by_cid:
                        continue
                    k = trade_key(t)
                    if k in seen[cid]:
                        continue
                    seen[cid].add(k)
                    b = accumulate(aggs[cid], t, by_cid[cid][1])
                    if b is not None:
                        touched[cid].add(b)
            for cid, bs in touched.items():
                con.executemany("INSERT OR REPLACE INTO flow VALUES(?,?,?,?,?,?,?)",
                                flow_rows(by_cid[cid][0], aggs[cid], bs))
            con.commit()
        except Exception as e:
            log.warning("flux direct: %s", e)
        for cid in [c for c in aggs if c not in by_cid]:
            aggs.pop(cid, None)
            seen.pop(cid, None)


def aggregate_flows(con, batch=40):
    now = int(time.time())
    pending = con.execute("""SELECT id, condition_id, up_token FROM markets
                             WHERE flow_done = 0 AND outcome IS NOT NULL AND end_ts < ?
                             ORDER BY end_ts DESC LIMIT ?""", (now - 120, batch)).fetchall()
    if not pending:
        return 0

    @safe
    def one(p):
        return p[0], fetch_flow(*p)

    with ThreadPoolExecutor(4) as ex:
        res = [r for r in ex.map(one, pending) if r]
    for mkt, (rows, truncated) in res:
        con.executemany("INSERT OR REPLACE INTO flow VALUES(?,?,?,?,?,?,?)", rows)
        con.execute("UPDATE markets SET flow_done = ? WHERE id = ?", (2 if truncated else 1, mkt))
    con.commit()
    return len(pending)


def resolver_loop(db, stop):
    con = connect(db)
    while not stop.is_set():
        busy = False
        try:
            resolve_pending(con)
            busy = aggregate_flows(con) >= 40
        except Exception as e:
            log.warning("resolver: %s", e)
        if not busy:
            stop.wait(60)


# ---------------------------------------------------------------- historique

def backfill(db, pairs, days, stop):
    con = connect(db)
    now = int(time.time())
    t0 = int(now - days * 86400)
    for asset, tf in pairs:
        # jour par jour : l'API refuse les decalages (offset) trop grands (HTTP 422 au-dela d'environ 2000)
        n, day = 0, t0
        while day < now and not stop.is_set():
            d_end = min(day + 86400, now)
            offset = 0
            while not stop.is_set():
                try:
                    evs = http_json(f"{GAMMA}/events", {"series_slug": series_slug(asset, tf), "closed": "true",
                                                        "end_date_min": ts_iso(day), "end_date_max": ts_iso(d_end),
                                                        "order": "endDate", "ascending": "false",
                                                        "limit": 100, "offset": offset})
                except Exception as e:
                    log.warning("backfill %s %s %s: %s", asset, tf, ts_iso(day)[:10], e)
                    break
                for ev in evs:
                    for m in ev.get("markets", []):
                        m["_eventMetadata"] = ev.get("eventMetadata")
                        r = parse_market(m, asset, tf)
                        if r:
                            upsert_market(con, r)
                            n += 1
                con.commit()
                if len(evs) < 100:
                    break
                offset += 100
            day = d_end
        log.info("backfill : %d marches %s %s listes", n, asset, tf)

    todo = con.execute("""SELECT id, up_token, tf, start_ts, end_ts FROM markets
                          WHERE hist_done = 0 AND end_ts BETWEEN ? AND ?""", (t0, now)).fetchall()
    log.info("backfill : %d historiques de cote a recuperer", len(todo))

    @safe
    def one(row):
        mkt, tok, tf, start, end = row
        h = http_json(f"{CLOB}/prices-history", {"market": tok, "startTs": start - TF_SECONDS[tf],
                                                 "endTs": end, "fidelity": 1})
        return mkt, [(p["t"], mkt, "hist", None, None, p["p"], None, None, None, None, None, None, None,
                      None, None, None) for p in h.get("history", [])]

    done = 0
    with ThreadPoolExecutor(4) as ex:
        for i in range(0, len(todo), 200):
            if stop.is_set():
                break
            # telecharger d'abord, ecrire ensuite : ne jamais garder la base verrouillee pendant un appel reseau
            for r in list(ex.map(one, todo[i:i + 200])):
                if r:
                    con.executemany("INSERT INTO book VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)", r[1])
                    con.execute("UPDATE markets SET hist_done = 1 WHERE id = ?", (r[0],))
            con.commit()
            done = min(i + 200, len(todo))
            log.info("backfill : %d/%d historiques", done, len(todo))
    log.info("backfill termine (le flux de transactions est agrege en arriere-plan)")


# ---------------------------------------------------------------- boucle principale

def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--db", default=default_db_path())
    ap.add_argument("--assets", default="btc,eth,sol,xrp")
    ap.add_argument("--tfs", default="5m,15m,1h,4h", help="parmi 5m,15m,1h,4h,1d")
    ap.add_argument("--backfill-days", type=float, default=0)
    args = ap.parse_args()

    assets = [a.strip() for a in args.assets.split(",") if a.strip()]
    tfs = [t.strip() for t in args.tfs.split(",") if t.strip()]
    pairs = [(a, t) for a in assets for t in tfs]
    symbols = [SYMBOLS[a] for a in assets]

    os.makedirs(os.path.dirname(os.path.abspath(args.db)), exist_ok=True)
    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s",
                        handlers=[logging.StreamHandler(),
                                  logging.FileHandler(os.path.join(os.path.dirname(os.path.abspath(args.db)),
                                                                   "collector.log"), encoding="utf-8")])
    log.info("base : %s | actifs %s | horizons %s", args.db, assets, tfs)

    lock = single_instance("collector", args.db)  # noqa: F841 (garde le verrou ouvert)
    con = connect(args.db)
    stop = threading.Event()
    state = {"tracked": {}}
    threads = [threading.Thread(target=resolver_loop, args=(args.db, stop), daemon=True),
               threading.Thread(target=live_flow_loop, args=(args.db, lambda: list(state["tracked"].values()), stop),
                                daemon=True)]
    if args.backfill_days > 0:
        threads.append(threading.Thread(target=backfill, args=(args.db, pairs, args.backfill_days, stop), daemon=True))
    for t in threads:
        t.start()

    tracked, last_sample = {}, {}
    last_disc = last_spot = last_status = 0
    n_book = n_spot = 0
    try:
        while True:
            loop_t = time.time()
            if loop_t - last_disc >= DISCOVER_EVERY:
                found = discover(con, pairs)
                if found:
                    tracked = {r["condition_id"]: r for r in found}
                else:
                    tracked = {k: r for k, r in tracked.items() if r["end_ts"] > loop_t}
                state["tracked"] = tracked
                last_disc = loop_t

            due = [r for r in tracked.values()
                   if r["start_ts"] - TF_SECONDS[r["tf"]] <= loop_t < r["end_ts"]
                   and loop_t - last_sample.get(r["condition_id"], 0) >= TF_SAMPLE[r["tf"]]]
            if due:
                try:
                    n_book += sample_books(con, due)
                    for r in due:
                        last_sample[r["condition_id"]] = loop_t
                except Exception as e:
                    log.warning("carnets: %s", e)
            if loop_t - last_spot >= SPOT_EVERY:
                try:
                    n_spot += sample_spot(con, symbols)
                except Exception as e:
                    log.warning("spot: %s", e)
                last_spot = loop_t
            con.commit()

            if loop_t - last_status >= 60:
                pend = con.execute("SELECT COUNT(*) FROM markets WHERE outcome IS NULL AND end_ts < ?",
                                   (int(loop_t),)).fetchone()[0]
                flows = con.execute("SELECT COUNT(*) FROM markets WHERE flow_done = 0 AND outcome IS NOT NULL").fetchone()[0]
                size = sum(os.path.getsize(p) for p in (args.db, args.db + "-wal") if os.path.exists(p)) / 1e6
                log.info("suivis %d | +%d carnets +%d spot (1 min) | a resoudre %d | flux a agreger %d | base %.1f Mo",
                         len(tracked), n_book, n_spot, pend, flows, size)
                n_book = n_spot = 0
                last_status = loop_t
            time.sleep(max(0.2, 1.0 - (time.time() - loop_t)))
    except KeyboardInterrupt:
        log.info("arret demande")
    finally:
        stop.set()
        con.commit()
        con.close()


if __name__ == "__main__":
    main()