#!/usr/bin/env python3
"""
Fetch World Bank API v2 indicator data for sixteen countries and write:
  raw.json  - verbatim API responses, keyed by query
  data.csv  - tidy long table (country, indicator, year, value), nulls dropped

Constants below are the exact query parameters recorded in NOTE.md.
Run:  python3 fetch.py
"""

import csv
import hashlib
import json
import os
import sys
import time
import urllib.error
import urllib.request
from datetime import datetime, timezone

OUT_DIR = (
    "/Volumes/Storage/DeepSeek/workspace/pgx-gap-map/"
    "research-notes/worldbank-health"
)

COUNTRIES = [
    "USA", "GBR", "DEU", "FRA", "JPN", "BRA", "IND", "CHN",
    "ZAF", "NGA", "KEN", "MEX", "IDN", "PAK", "RUS", "AUS",
]

INDICATORS = {
    "SH.XPD.CHEX.PC.CD": "health_expenditure_per_capita_current_usd",
    "SH.DYN.MORT": "under5_mortality_per_1000_live_births",
    "SP.DYN.LE00.IN": "life_expectancy_at_birth_years",
}

BASE = "https://api.worldbank.org/v2/country"
DATE_RANGE = "2015:2022"
PER_PAGE = 1000
SLEEP_SECONDS = 0.4
MAX_REQUESTS = 40
USER_AGENT = (
    "pgx-gap-map-research-notes/1.0 "
    "(public World Bank API v2; python-urllib; contact: research note)"
)

request_count = 0
failures = []


def utc_now():
    return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")


def fetch(url):
    """One GET with one retry. Returns (status, raw_text, headers)."""
    global request_count
    attempts = []
    for attempt in (1, 2):
        if request_count >= MAX_REQUESTS:
            raise RuntimeError(f"request cap {MAX_REQUESTS} reached")
        request_count += 1
        req = urllib.request.Request(url, headers={"User-Agent": USER_AGENT})
        try:
            with urllib.request.urlopen(req, timeout=60) as resp:
                body = resp.read().decode("utf-8")
                status = resp.status
            time.sleep(SLEEP_SECONDS)
            return status, body, attempts
        except urllib.error.HTTPError as exc:
            attempts.append({"attempt": attempt, "http_status": exc.code,
                             "error": f"HTTPError {exc.code}"})
            if attempt == 1:
                time.sleep(SLEEP_SECONDS)
                continue
            raise
        except Exception as exc:  # noqa: BLE001 - network layer
            attempts.append({"attempt": attempt,
                             "error": f"{type(exc).__name__}: {exc}"})
            if attempt == 1:
                time.sleep(SLEEP_SECONDS)
                continue
            raise
    raise RuntimeError("unreachable")


def main():
    os.makedirs(OUT_DIR, exist_ok=True)
    country_path = ";".join(COUNTRIES)

    raw = {
        "meta": {
            "api": "World Bank API v2 (JSON, no key)",
            "base_url": BASE,
            "generated_at_utc": utc_now(),
            "date_parameter": DATE_RANGE,
            "per_page": PER_PAGE,
            "server_sleep_seconds": SLEEP_SECONDS,
            "max_requests": MAX_REQUESTS,
            "user_agent": USER_AGENT,
            "countries_iso3": COUNTRIES,
            "indicators": INDICATORS,
            "request_note": (
                "One request per indicator with all sixteen ISO3 codes joined by "
                "semicolons (supported by API v2). Sixteen single-country calls per "
                "indicator would be 48 requests, above the stated cap of 40."
            ),
        },
        "queries": {},
        "failures": failures,
    }

    all_rows = []

    for ind_code, ind_slug in INDICATORS.items():
        key = f"{ind_code}__16countries__{DATE_RANGE}"
        url = (
            f"{BASE}/{country_path}/indicator/{ind_code}"
            f"?format=json&date={DATE_RANGE}&per_page={PER_PAGE}"
        )
        entry = {
            "url": url,
            "indicator_code": ind_code,
            "indicator_slug": ind_slug,
            "countries": COUNTRIES,
            "date_parameter": DATE_RANGE,
            "requested_at_utc": utc_now(),
        }
        try:
            status, body, attempts = fetch(url)
        except Exception as exc:  # noqa: BLE001
            failures.append({"query_key": key, "url": url,
                             "error": f"{type(exc).__name__}: {exc}",
                             "at_utc": utc_now()})
            entry.update({"status": "failed",
                          "error": f"{type(exc).__name__}: {exc}",
                          "verbatim_response": None})
            raw["queries"][key] = entry
            print(f"FAILED {ind_code}: {exc}", file=sys.stderr)
            continue

        payload = json.loads(body)  # two-element array [metadata, data]
        entry.update({
            "status": status,
            "attempts": attempts or [{"attempt": 1, "http_status": status}],
            "byte_length": len(body.encode("utf-8")),
            "sha256": hashlib.sha256(body.encode("utf-8")).hexdigest(),
            "verbatim_response": payload,
        })
        raw["queries"][key] = entry

        meta = payload[0] if isinstance(payload, list) and payload else {}
        rows = payload[1] if isinstance(payload, list) and len(payload) > 1 else []
        if rows is None:
            rows = []
        entry["page_metadata"] = meta

        kept = 0
        for rec in rows:
            value = rec.get("value")
            if value is None:
                continue
            iso3 = rec.get("countryiso3code")
            if iso3 not in COUNTRIES:
                continue
            all_rows.append({
                "country": iso3,
                "indicator": ind_code,
                "year": int(rec["date"]),
                "value": value,
            })
            kept += 1
        entry["rows_returned_by_api"] = len(rows)
        entry["rows_kept_after_null_drop"] = kept
        print(f"OK {ind_code}: {len(rows)} returned, {kept} kept")

    raw["meta"]["request_count"] = request_count
    raw["meta"]["failure_count"] = len(failures)
    raw["meta"]["completed_at_utc"] = utc_now()

    raw_path = os.path.join(OUT_DIR, "raw.json")
    with open(raw_path, "w", encoding="utf-8") as fh:
        json.dump(raw, fh, indent=2, ensure_ascii=False, sort_keys=False)
        fh.write("\n")
    print(f"wrote {raw_path}")

    ind_order = {code: i for i, code in enumerate(INDICATORS)}
    all_rows.sort(key=lambda r: (r["country"], ind_order[r["indicator"]], r["year"]))
    csv_path = os.path.join(OUT_DIR, "data.csv")
    with open(csv_path, "w", encoding="utf-8", newline="") as fh:
        writer = csv.DictWriter(fh, fieldnames=["country", "indicator", "year", "value"])
        writer.writeheader()
        for row in all_rows:
            writer.writerow(row)
    print(f"wrote {csv_path} ({len(all_rows)} rows)")
    print(f"requests: {request_count}  failures: {len(failures)}")


if __name__ == "__main__":
    main()