#!/usr/bin/env python3
"""Ticker crowd around each sampled breakout. T = breakout bar open. Windows: pre [T-24h,T) 2 pages, first [T,T+2h) 1 page
(breakout hour + retest hour = what is known at the decision), after [T+2h,T+12h) 2 pages. Usage: ab_xpull.py <budget> <shard> <nshards>"""
import json, os, sys, time, importlib.util, re, glob
sp = importlib.util.spec_from_file_location("xsafe", "/opt/common/xsafe/xsafe.py")
xs = importlib.util.module_from_spec(sp); sp.loader.exec_module(xs)
xs.budget(int(sys.argv[1])); SH, NSH = int(sys.argv[2]), int(sys.argv[3])
EV = json.load(open("/opt/abstudy/xsample.json"))
OUT = f"/opt/abstudy/x_{SH}.jsonl"; done = set()
for fn in glob.glob("/opt/abstudy/x_*.jsonl"):
    for l in open(fn):
        try: done.add(json.loads(l)["k"])
        except Exception: pass
keep = ("id", "text", "createdAt", "likeCount", "retweetCount", "replyCount", "viewCount", "isReply")
def slim(t):
    d = {k: t.get(k) for k in keep}; a = t.get("author") or {}
    d["user"] = a.get("userName"); d["fol"] = a.get("followers"); d["ts"] = xs.ts(t); return d
def q(e):
    sym = (e["sym"] or "").strip().lstrip("$"); nm = re.sub(r'["()]', "", (e["name"] or "").strip())
    return f'(${sym} OR "{nm}")' if nm and nm.lower() != sym.lower() and len(nm) >= 3 else f"${sym}"
n = 0
with open(OUT, "a") as f:
    for idx, e in enumerate(EV):
        if idx % NSH != SH: continue
        k = f'{e["a"]}:{e["t"]}'
        if k in done: continue
        T = e["t"]; rec = {"k": k, "query": q(e)}
        try:
            for lab, a, b, pg in (("pre", T - 86400, T, 2), ("first", T, T + 7200, 1), ("after", T + 7200, T + 43200, 2)):
                r = xs.search(rec["query"], since=a, until=b, pages=pg)
                rec[lab] = {"real": [slim(t) for t in r["real"]], "spam": len(r["spam"]), "complete": r["complete"]}
        except Exception as ex:
            print("ERR", k, str(ex)[:160], flush=True)
            if "run cap" in str(ex): break
            if "brake" in str(ex).lower(): time.sleep(90)
            continue
        f.write(json.dumps(rec) + "\n"); f.flush(); n += 1
        if n % 50 == 0: print(time.strftime("%H:%M:%S"), "done", n, flush=True)
print("FINISHED", n, flush=True)