"""Fetch GDELT news-volume timelines for the S&P 500 panel.

One GDELT DOC 2.0 timelinevol query per company returns the full daily
series since 2017 (news volume intensity = share of all monitored news).
Company names come from the SEC ticker map (company_tickers.json); legal
suffixes are trimmed for the phrase query. The public API is burst-limited
(429s), so requests are paced with adaptive backoff.

Cache: data/research/gdelt/<ticker>.json (resume-safe). Output panel:
data/research/gdelt_volume.pkl {"volume": DataFrame} on daily dates.

    PYTHONPATH=. uv run python -m research.lab.fetch_gdelt [--limit N]
"""
from __future__ import annotations

import json
import os
import pickle
import re
import sys
import time

import pandas as pd
import requests

from research.lab.fetch_us import _proxy_env

OUT_DIR = os.path.join("data", "research", "gdelt")
PANEL = os.path.join("data", "research", "gdelt_volume.pkl")
TICKERS_CACHE = "data/research/company_tickers.json"
UA = {"User-Agent": "quant-research/0.1 gahow@example.com"}
BASE = ("https://api.gdeltproject.org/api/v2/doc/doc?query={q}%20sourcelang:english"
        "&mode=timelinevol&format=json&startdatetime=20170101000000&enddatetime=20261007000000")
SUFFIX = re.compile(r",?\s*(inc|incorporated|corp|corporation|co|company|ltd|limited|plc|"
                    r"holdings?|group|the|nv|sa|ag|llc|lp|trust|reit)\.?$", re.I)


def clean_name(title: str) -> str:
    name = title.strip().rstrip(".")
    prev = None
    while prev != name:
        prev = name
        name = SUFFIX.sub("", name).strip().rstrip(",")
    return name


def query(name: str) -> list | None:
    q = requests.utils.quote(f'"{name}"')
    url = BASE.format(q=q)
    delay = 5.0
    for attempt in range(5):
        try:
            r = requests.get(url, headers=UA, timeout=90)
            if r.status_code == 429:
                time.sleep(delay)
                delay = min(delay * 2, 60)
                continue
            r.raise_for_status()
            d = r.json()
            series = d.get("timeline", [{}])[0].get("data", [])
            return series
        except Exception as exc:  # noqa: BLE001
            if attempt == 4:
                print(f"  FAILED {name}: {exc}")
            else:
                time.sleep(delay)
                delay = min(delay * 2, 60)
    return None


def main() -> None:
    _proxy_env()
    os.makedirs(OUT_DIR, exist_ok=True)
    limit = None
    if "--limit" in sys.argv:
        limit = int(sys.argv[sys.argv.index("--limit") + 1])
    close = pd.read_csv("data/us.csv", index_col=0, nrows=0)
    tickers = list(close.columns)
    raw = json.load(open(TICKERS_CACHE))
    title_by_ticker = {}
    for v in raw.values():
        t = str(v["ticker"]).upper().replace(".", "-")
        title_by_ticker.setdefault(t, str(v["title"]))
    todo = [t for t in tickers if t in title_by_ticker and
            not os.path.exists(os.path.join(OUT_DIR, f"{t}.json"))]
    if limit:
        todo = todo[:limit]
    print(f"tickers: {len(tickers)}, cached: {len(tickers)-len(todo)}, fetching: {len(todo)}")
    for i, t in enumerate(todo):
        series = query(clean_name(title_by_ticker[t]))
        if series is None:
            continue
        with open(os.path.join(OUT_DIR, f"{t}.json"), "w") as fh:
            json.dump(series, fh)
        if (i + 1) % 25 == 0:
            print(f"  {i+1}/{len(todo)} (last {t}, {len(series)} points)", flush=True)
        time.sleep(12.0)

    # assemble panel
    frames = {}
    for t in tickers:
        path = os.path.join(OUT_DIR, f"{t}.json")
        if not os.path.exists(path):
            continue
        try:
            series = json.load(open(path))
            idx = pd.to_datetime([p["date"][:8] for p in series], format="%Y%m%d")
            vals = [p["value"] for p in series]
            frames[t] = pd.Series(vals, index=idx)
        except Exception:  # noqa: BLE001
            continue
    if frames:
        df = pd.DataFrame(frames).sort_index()
        df = df[~df.index.duplicated(keep="last")]
        with open(PANEL, "wb") as fh:
            pickle.dump({"volume": df.astype("float32")}, fh)
        print(f"panel {df.shape} -> {PANEL}")
    else:
        print("no cached series yet")


if __name__ == "__main__":
    main()