"""
Daily English Wikipedia pageviews for every US research-panel ticker, as a
point-in-time attention proxy (the free source with the longest daily history;
X/Reddit history is not available for free, see research/lab/attention.py).

    uv run python -m research.lab.fetch_wiki_attention            # mapping + pageviews
    uv run python -m research.lab.fetch_wiki_attention --map-only

Ticker -> article, first hit wins:
  0. OVERRIDES below (hand-checked)
  1. "List of S&P 500 companies" constituents table (current members)
  2. "Historical components of the S&P 500" change log (added/removed names,
     most recent entry first, so a reused ticker gets its latest S&P company)
  3. Wikidata: SEC CIK (P5531) of data/research/sec_sic.csv
  4. Wikidata: NYSE/Nasdaq ticker qualifier (P414 / P249)
Titles are resolved to the canonical article; views are summed over the
article and every redirect to it, so a rename (Facebook -> Meta Platforms)
does not break the series. Agent "user" (spiders and automated excluded).

Writes data/research/wiki_pageviews.pkl:
  {"views": DataFrame dates x tickers (canonical + redirects), "complete": bool,
   "views_main": same, canonical title only,
   "map": DataFrame ticker/title/source, "fetched": timestamp}
"""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import pickle
import re
import threading
import time
import urllib.parse
from concurrent.futures import ThreadPoolExecutor

import pandas as pd
import requests
from bs4 import BeautifulSoup

UA = "quant-research-attention/0.1 (academic factor research; gahow@example.com)"
OUT = os.path.join("data", "research", "wiki_pageviews.pkl")
CACHE = os.path.join("data", "research", "wiki_cache")
START, END = "20150701", "20260929"
SPARQL = "https://query.wikidata.org/sparql"
API = "https://en.wikipedia.org/w/api.php"
PV = ("https://wikimedia.org/api/rest_v1/metrics/pageviews/per-article/en.wikipedia/"
      "all-access/user/{t}/daily/{a}/{b}")
# change-log links that point at the wrong article, or names missing from both tables
OVERRIDES = {"ITT": "ITT Inc.", "CNX": "CNX Resources", "PODD": "Insulet", "MBC": "MasterBrand"}
TICKER_RE = re.compile(r"^[A-Z][A-Z0-9.\-]{0,6}$")

_session = threading.local()


def _s() -> requests.Session:
    if not hasattr(_session, "s"):
        _session.s = requests.Session()
        _session.s.headers["User-Agent"] = UA
    return _session.s


def _get(url, **kw):
    for i in range(20):
        try:
            r = _s().get(url, timeout=60, **kw)
        except requests.RequestException:
            time.sleep(2 ** i)
            continue
        if r.status_code == 429 or r.status_code >= 500:
            ra = r.headers.get("Retry-After")
            time.sleep(float(ra) + 1 if ra else 2 ** i)
            continue
        return r
    raise RuntimeError(f"failed: {url}")


def _norm(t: str) -> str:
    return t.strip().upper().replace(".", "-")


def _title(href: str) -> str | None:
    if not href or not href.startswith("./"):
        return None
    t = urllib.parse.unquote(href[2:]).split("#")[0]
    return t.replace("_", " ")


def _page(title: str) -> BeautifulSoup:
    r = _get("https://en.wikipedia.org/api/rest_v1/page/html/" + urllib.parse.quote(title, safe=""))
    return BeautifulSoup(r.text, "lxml")


def _from_wikipedia() -> tuple[dict, dict]:
    cur = {}
    s = _page("List_of_S&P_500_companies")
    for tr in s.find("table", id="constituents").find_all("tr")[1:]:
        td = tr.find_all("td")
        if len(td) >= 2:
            a = td[1].find("a")
            t = _title(a.get("href")) if a else None
            if t:
                cur[_norm(td[0].get_text(strip=True))] = t
    hist = {}
    s = _page("Historical_components_of_the_S&P_500")
    for tr in s.find("table", id="changes").find_all("tr")[2:]:     # newest first
        cells = tr.find_all("td")
        for c, nxt in zip(cells, cells[1:]):
            tk = c.get_text(strip=True)
            a = nxt.find("a")
            if TICKER_RE.match(tk) and a is not None:
                t = _title(a.get("href"))
                if t:
                    hist.setdefault(_norm(tk), t)
    return cur, hist


def _sparql(q: str) -> pd.DataFrame:
    r = _get(SPARQL, params={"query": q}, headers={"Accept": "text/csv"})
    from io import StringIO
    return pd.read_csv(StringIO(r.text), dtype=str)


def _from_wikidata() -> tuple[dict, dict]:
    c = _sparql("""SELECT ?cik ?article WHERE { ?item wdt:P5531 ?cik .
        ?article schema:about ?item ; schema:isPartOf <https://en.wikipedia.org/> . }""")
    by_cik = {}
    for cik, art in zip(c["cik"].str.lstrip("0"), c["article"]):
        by_cik.setdefault(cik, urllib.parse.unquote(art.rsplit("/", 1)[1]).replace("_", " "))
    t = _sparql("""SELECT ?ticker ?article WHERE { VALUES ?exch { wd:Q13677 wd:Q82059 }
        ?item p:P414 ?st . ?st ps:P414 ?exch ; pq:P249 ?ticker .
        ?article schema:about ?item ; schema:isPartOf <https://en.wikipedia.org/> . }""")
    by_tk = {}
    for tk, art in zip(t["ticker"], t["article"]):
        by_tk.setdefault(_norm(tk), urllib.parse.unquote(art.rsplit("/", 1)[1]).replace("_", " "))
    return by_cik, by_tk


def _canonical(titles: list[str]) -> dict[str, str]:
    out = {}
    for i in range(0, len(titles), 50):
        batch = titles[i:i + 50]
        r = _get(API, params={"action": "query", "titles": "|".join(batch), "redirects": 1,
                              "format": "json", "formatversion": 2}).json()["query"]
        norm = {n["from"]: n["to"] for n in r.get("normalized", [])}
        red = {n["from"]: n["to"] for n in r.get("redirects", [])}
        missing = {p["title"] for p in r["pages"] if p.get("missing")}
        for t in batch:
            x = norm.get(t, t)
            x = red.get(x, x)
            if x not in missing:
                out[t] = x
    return out


def _redirects(titles: list[str]) -> dict[str, list[str]]:
    """Redirects into each title, 50 titles per request, sequential (the action
    API throttles parallel clients); cached in CACHE/redirects.json."""
    path = os.path.join(CACHE, "redirects.json")
    os.makedirs(CACHE, exist_ok=True)
    out = json.load(open(path)) if os.path.exists(path) else {}
    todo = [t for t in titles if t not in out]
    for i in range(0, len(todo), 50):
        batch, cont = todo[i:i + 50], {}
        got = {t: [] for t in batch}
        while True:
            r = _get(API, params={"action": "query", "titles": "|".join(batch), "prop": "redirects",
                                  "rdnamespace": 0, "rdlimit": "max", "format": "json",
                                  "formatversion": 2, **cont}).json()
            for p in r["query"]["pages"]:
                got.setdefault(p["title"], []).extend(x["title"] for x in p.get("redirects", []))
            if "continue" not in r:
                break
            cont = r["continue"]
        out.update({t: got.get(t, []) for t in batch})
        with open(path, "w") as fh:
            json.dump(out, fh)
    return {t: out[t] for t in titles}


def build_map(tickers: list[str]) -> pd.DataFrame:
    cur, hist = _from_wikipedia()
    by_cik, by_tk = _from_wikidata()
    sic = pd.read_csv(os.path.join("data", "research", "sec_sic.csv"), dtype={"cik": str})
    cik = dict(zip(sic["ticker"], sic["cik"].str.lstrip("0")))
    rows = []
    for t in tickers:
        n = _norm(t)
        for src, title in (("override", OVERRIDES.get(t)), ("sp500_list", cur.get(n)), ("sp500_changes", hist.get(n)),
                           ("wikidata_cik", by_cik.get(cik.get(t, ""))),
                           ("wikidata_ticker", by_tk.get(n))):
            if title:
                rows.append({"ticker": t, "title": title, "source": src})
                break
        else:
            rows.append({"ticker": t, "title": None, "source": None})
    m = pd.DataFrame(rows)
    canon = _canonical(sorted(m["title"].dropna().unique()))
    m["canonical"] = m["title"].map(canon)
    return m


def _views(title: str) -> pd.Series:
    os.makedirs(CACHE, exist_ok=True)
    key = os.path.join(CACHE, re.sub(r"[^A-Za-z0-9]+", "_", title)[:120] + "_" + hashlib.md5(title.encode()).hexdigest()[:8] + ".json")
    if os.path.exists(key):
        with open(key) as fh:
            d = json.load(fh)
    else:
        url = PV.format(t=urllib.parse.quote(title.replace(" ", "_"), safe=""), a=START, b=END)
        r = _get(url)
        d = r.json().get("items", []) if r.status_code == 200 else []
        with open(key, "w") as fh:
            json.dump(d, fh)
        time.sleep(0.05)
    if not d:
        return pd.Series(dtype=float)
    return pd.Series({pd.Timestamp(x["timestamp"][:8]): float(x["views"]) for x in d})


def _assemble(m, titles, reds, V, with_redirects: bool) -> tuple[pd.DataFrame, pd.DataFrame]:
    idx = pd.date_range(pd.Timestamp(START), pd.Timestamp(END), freq="D")
    main, tot = {}, {}
    for t in titles:
        s0 = V[t].reindex(idx)
        parts = [s0.fillna(0.0)] + ([V[r].reindex(idx).fillna(0.0) for r in reds[t]] if with_redirects else [])
        s = pd.concat(parts, axis=1).sum(axis=1)
        main[t], tot[t] = s0, s.where(s0.notna() | (s > 0))
    tick = m.dropna(subset=["canonical"])
    views = pd.DataFrame({r.ticker: tot[r.canonical] for r in tick.itertuples()})
    views_main = pd.DataFrame({r.ticker: main[r.canonical] for r in tick.itertuples()})
    return views, views_main


def _save(m, views, views_main, complete: bool) -> None:
    with open(OUT, "wb") as fh:
        pickle.dump({"views": views, "views_main": views_main, "map": m, "complete": complete,
                     "fetched": pd.Timestamp.now()}, fh)
    print(f"wrote {OUT} (redirects {'included' if complete else 'NOT included'}): {views.shape}", flush=True)


def fetch_views(m: pd.DataFrame, workers: int = 4) -> None:
    """Canonical articles first (saved without redirects), then every redirect."""
    titles = sorted(m["canonical"].dropna().unique())
    reds = _redirects(titles)
    rest = sorted({r for v in reds.values() for r in v} - set(titles))
    print(f"{len(titles)} articles, {len(rest)} redirects", flush=True)
    done = [0]

    def one(t):
        s = _views(t)
        done[0] += 1
        if done[0] % 500 == 0:
            print(f"  {done[0]}/{len(titles) + len(rest)}", flush=True)
        return t, s

    with ThreadPoolExecutor(workers) as ex:
        V = dict(ex.map(one, titles))
    _save(m, *_assemble(m, titles, reds, V, False), complete=False)
    with ThreadPoolExecutor(workers) as ex:
        V.update(dict(ex.map(one, rest)))
    _save(m, *_assemble(m, titles, reds, V, True), complete=True)


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--map-only", action="store_true")
    ap.add_argument("--workers", type=int, default=4)
    args = ap.parse_args()
    with open(os.path.join("data", "research", "us.panel.pkl"), "rb") as fh:
        tickers = list(pickle.load(fh).close.columns)
    m = build_map(tickers)
    print(m["source"].value_counts(dropna=False).to_string())
    print(f"mapped {m['canonical'].notna().sum()}/{len(m)}; shared articles:",
          m.dropna(subset=["canonical"]).groupby("canonical")["ticker"].agg(list).loc[lambda s: s.str.len() > 1].to_dict())
    m.to_csv(os.path.join("data", "research", "wiki_map.csv"), index=False)
    if args.map_only:
        return
    fetch_views(m, args.workers)


if __name__ == "__main__":
    main()