#!/usr/bin/env python
"""Generate `registry/countries.csv` (ISO-3166-1 alpha-2, name, UN M49 region/subregion, centroid lat/lon).

Sources (fetched once at generation time; the CSV is committed so nothing is downloaded at runtime):
  * ISO 3166 + UN M49 regions: https://github.com/lukes/ISO-3166-Countries-with-Regional-Codes (all/all.csv, public domain-ish data)
  * Centroids: Wikidata `P297` (ISO alpha-2) → `P625` (coordinate location) via the SPARQL endpoint; a few manual overrides fix items whose
    coordinate is a capital or an outlying island rather than the country centroid.

Usage: .venv/bin/python scripts/seed_countries.py [--out registry/countries.csv]
"""
from __future__ import annotations

import argparse
import csv
import io
import sys
import time
from pathlib import Path

import httpx

ROOT = Path(__file__).resolve().parents[1]
UA = "CompanyAtlasBot/0.1 (contact@spboucher.ai)"
ISO_URL = "https://raw.githubusercontent.com/lukes/ISO-3166-Countries-with-Regional-Codes/master/all/all.csv"
SPARQL = "https://query.wikidata.org/sparql"
NAME_OVERRIDES = {
    "US": "United States", "GB": "United Kingdom", "RU": "Russia", "KR": "South Korea", "KP": "North Korea", "IR": "Iran", "VN": "Vietnam",
    "TW": "Taiwan", "BO": "Bolivia", "VE": "Venezuela", "TZ": "Tanzania", "MD": "Moldova", "SY": "Syria", "LA": "Laos", "BN": "Brunei",
    "CZ": "Czechia", "FM": "Micronesia", "PS": "Palestine", "CD": "DR Congo", "CG": "Republic of the Congo", "CI": "Côte d'Ivoire",
    "VA": "Vatican City", "FK": "Falkland Islands", "VG": "British Virgin Islands", "VI": "U.S. Virgin Islands", "TR": "Türkiye",
    "NL": "Netherlands", "BQ": "Caribbean Netherlands", "SH": "Saint Helena", "MO": "Macao", "HK": "Hong Kong", "AE": "United Arab Emirates",
    "TF": "French Southern Territories", "UM": "U.S. Minor Outlying Islands", "GS": "South Georgia and the South Sandwich Islands",
}
# UN M49 does not assign a region to Taiwan (and Antarctica has none); keep the atlas usable for filtering.
REGION_OVERRIDES = {"TW": ("Asia", "Eastern Asia"), "AQ": ("Antarctica", "Antarctica")}
# Centroid overrides (lat, lon) where Wikidata's P625 is a capital/point rather than a usable centroid, or is missing.
CENTROID_OVERRIDES = {
    "US": (39.8, -98.6), "CA": (56.1, -106.3), "RU": (61.5, 105.3), "FR": (46.6, 2.2), "NO": (64.6, 12.7), "DK": (56.0, 10.0),
    "NL": (52.2, 5.3), "GB": (54.0, -2.5), "AU": (-25.3, 133.8), "NZ": (-41.5, 172.8), "CL": (-35.7, -71.5), "BR": (-14.2, -51.9),
    "ID": (-2.5, 118.0), "JP": (36.2, 138.3), "CN": (35.9, 104.2), "IN": (20.6, 79.0), "KI": (-3.4, -168.7), "FM": (7.4, 150.6),
    "PT": (39.4, -8.2), "ES": (40.5, -3.7), "EC": (-1.8, -78.2), "UM": (19.3, 166.6), "TF": (-49.3, 69.3), "AQ": (-82.9, 135.0),
    "ZA": (-30.6, 22.9), "AR": (-38.4, -63.6), "MX": (23.6, -102.6), "DE": (51.2, 10.4), "IT": (41.9, 12.6), "TW": (23.7, 121.0),
    "HK": (22.3, 114.2), "SG": (1.35, 103.8), "AE": (23.4, 53.8), "SA": (23.9, 45.1), "NG": (9.1, 8.7), "SE": (60.1, 18.6), "CH": (46.8, 8.2),
    "KR": (35.9, 127.8), "IL": (31.0, 34.9), "PS": (31.9, 35.2), "EG": (26.8, 30.8), "TR": (39.0, 35.2), "IR": (32.4, 53.7), "PK": (30.4, 69.3),
    "GL": (71.7, -42.6), "SJ": (77.6, 23.7), "MY": (4.2, 108.0), "PH": (12.9, 121.8), "VN": (14.1, 108.3), "TH": (15.9, 100.9),
    "UA": (48.4, 31.2), "PL": (51.9, 19.1), "KZ": (48.0, 66.9), "MN": (46.9, 103.8), "IE": (53.4, -8.2), "IS": (64.96, -19.0), "FI": (61.9, 25.7),
    "GR": (39.1, 21.8), "MA": (31.8, -7.1), "DZ": (28.0, 1.7), "LY": (26.3, 17.2), "SD": (12.9, 30.2), "ET": (9.1, 40.5), "KE": (-0.02, 37.9),
    "CD": (-4.0, 21.8), "AO": (-11.2, 17.9), "MZ": (-18.7, 35.5), "MG": (-18.8, 46.9), "TZ": (-6.4, 34.9), "ML": (17.6, -4.0), "NE": (17.6, 8.1),
    "TD": (15.5, 18.7), "MR": (21.0, -10.9), "PE": (-9.2, -75.0), "CO": (4.6, -74.3), "VE": (6.4, -66.6), "BO": (-16.3, -63.6), "PY": (-23.4, -58.4),
    "UY": (-32.5, -55.8), "CU": (21.5, -77.8), "GT": (15.8, -90.2), "HN": (15.2, -86.2), "NI": (12.9, -85.2), "PA": (8.5, -80.8), "CR": (9.7, -83.8),
    "AF": (33.9, 67.7), "IQ": (33.2, 43.7), "SY": (34.8, 39.0), "JO": (30.6, 36.2), "OM": (21.5, 55.9), "YE": (15.6, 48.5), "UZ": (41.4, 64.6),
    "TM": (38.97, 59.6), "NP": (28.4, 84.1), "BD": (23.7, 90.4), "MM": (21.9, 95.96), "LK": (7.9, 80.8), "KH": (12.6, 105.0), "LA": (19.9, 102.5),
    "PG": (-6.3, 143.96), "SB": (-9.6, 160.2), "VU": (-15.4, 166.96), "FJ": (-17.7, 178.1), "TO": (-21.2, -175.2), "WS": (-13.8, -172.1),
    "CV": (16.0, -24.0), "MU": (-20.3, 57.6), "SC": (-4.7, 55.5), "MV": (3.2, 73.2), "BH": (26.0, 50.6), "QA": (25.4, 51.2), "KW": (29.3, 47.5),
    "LB": (33.9, 35.9), "CY": (35.1, 33.4), "MT": (35.9, 14.4), "LU": (49.8, 6.1), "BE": (50.5, 4.5), "AT": (47.5, 14.6), "CZ": (49.8, 15.5),
    "SK": (48.7, 19.7), "HU": (47.2, 19.5), "RO": (45.9, 25.0), "BG": (42.7, 25.5), "RS": (44.0, 21.0), "HR": (45.1, 15.2), "SI": (46.2, 15.0),
    "BA": (43.9, 17.7), "ME": (42.7, 19.4), "MK": (41.6, 21.7), "AL": (41.2, 20.2), "XK": (42.6, 20.9), "EE": (58.6, 25.0), "LV": (56.9, 24.6),
    "LT": (55.2, 23.9), "BY": (53.7, 27.95), "MD": (47.4, 28.4), "GE": (42.3, 43.4), "AM": (40.1, 45.0), "AZ": (40.1, 47.6), "KG": (41.2, 74.8),
    "TJ": (38.9, 71.3), "BT": (27.5, 90.4), "TL": (-8.9, 125.7), "BN": (4.5, 114.7), "KP": (40.3, 127.5), "MO": (22.2, 113.55),
}


def fetch_iso() -> list[dict[str, str]]:
    r = httpx.get(ISO_URL, headers={"User-Agent": UA}, timeout=60, follow_redirects=True)
    r.raise_for_status()
    return list(csv.DictReader(io.StringIO(r.text)))


def fetch_centroids() -> dict[str, tuple[float, float]]:
    query = """
    SELECT ?iso ?coord WHERE {
      ?c wdt:P297 ?iso ; wdt:P625 ?coord .
      FILTER NOT EXISTS { ?c wdt:P576 ?dissolved }
    }"""
    for attempt in range(4):
        r = httpx.get(SPARQL, params={"query": query, "format": "json"}, headers={"User-Agent": UA, "Accept": "application/sparql-results+json"},
                      timeout=90)
        if r.status_code == 200:
            break
        time.sleep(5 * (attempt + 1))
    r.raise_for_status()
    out: dict[str, tuple[float, float]] = {}
    for b in r.json()["results"]["bindings"]:
        iso = b["iso"]["value"].upper()
        val = b["coord"]["value"]  # Point(lon lat)
        try:
            lon, lat = val.removeprefix("Point(").removesuffix(")").split()
            out.setdefault(iso, (round(float(lat), 3), round(float(lon), 3)))
        except ValueError:
            continue
    return out


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", default=str(ROOT / "registry" / "countries.csv"))
    args = ap.parse_args()
    iso = fetch_iso()
    time.sleep(2)
    centroids = fetch_centroids()
    rows = []
    for r in iso:
        code = r["alpha-2"].strip().upper()
        if len(code) != 2:
            continue
        name = NAME_OVERRIDES.get(code) or r["name"].split(" (")[0].split(",")[0].strip()
        lat, lon = CENTROID_OVERRIDES.get(code) or centroids.get(code) or ("", "")
        region, subregion = REGION_OVERRIDES.get(code) or (r.get("region") or "", r.get("sub-region") or "")
        rows.append({"code": code, "name": name, "region": region, "subregion": subregion, "lat": lat, "lon": lon})
    # Kosovo is user-assigned (XK) and not in ISO 3166-1; several data sources (and Wikidata companies) use it.
    if not any(x["code"] == "XK" for x in rows):
        rows.append({"code": "XK", "name": "Kosovo", "region": "Europe", "subregion": "Southern Europe", "lat": 42.6, "lon": 20.9})
    rows.sort(key=lambda x: x["code"])
    missing = [x["code"] for x in rows if x["lat"] == ""]
    out = Path(args.out)
    out.parent.mkdir(parents=True, exist_ok=True)
    with out.open("w", newline="", encoding="utf-8") as f:
        w = csv.DictWriter(f, fieldnames=["code", "name", "region", "subregion", "lat", "lon"])
        w.writeheader()
        w.writerows(rows)
    print(f"wrote {len(rows)} countries → {out} (missing centroids: {missing or 'none'})")
    return 0


if __name__ == "__main__":
    sys.exit(main())