#!/usr/bin/env python3
"""Download all script-accessible raw sources into _local/raw/.

The downloader keeps every file inside PARTY2D_RAW_DATA_DIR (default:
_local/raw). It does not write to committed data/. Sources that require an API
key or provider form are handled explicitly and reported rather than silently
skipped.
"""

from __future__ import annotations

import html.parser
import json
import http.cookiejar
import os
import re
import shutil
import sys
import urllib.parse
import urllib.request
import zipfile
from pathlib import Path


UA = "party2d-data-setup/1.0 (+https://git.seimel.app/armin/party2d)"


def raw_dir() -> Path:
    return Path(os.environ.get("PARTY2D_RAW_DATA_DIR", Path.cwd() / "_local" / "raw")).resolve()


def report_dir() -> Path:
    return Path(os.environ.get("PARTY2D_REPORT_DIR", Path.cwd() / "_local" / "reports")).resolve()


def request(url: str, data: bytes | None = None, headers: dict[str, str] | None = None):
    h = {"User-Agent": UA}
    if headers:
        h.update(headers)
    return urllib.request.Request(url, data=data, headers=h)


def download(url: str, dest: Path, *, overwrite: bool = False, headers: dict[str, str] | None = None) -> bool:
    dest.parent.mkdir(parents=True, exist_ok=True)
    if dest.exists() and dest.stat().st_size > 0 and not overwrite:
        print(f"OK existing: {dest}")
        return True
    tmp = dest.with_suffix(dest.suffix + ".tmp")
    print(f"Downloading {url} -> {dest}")
    try:
        with urllib.request.urlopen(request(url, headers=headers), timeout=180) as r, tmp.open("wb") as fh:
            shutil.copyfileobj(r, fh)
        tmp.replace(dest)
        return True
    except Exception as e:
        if tmp.exists():
            tmp.unlink()
        print(f"FAILED {url}: {type(e).__name__}: {e}")
        return False


def dataverse_file_id(doi: str, filename: str, directory: str | None = None) -> int | None:
    api = "https://dataverse.harvard.edu/api/datasets/:persistentId/?persistentId=" + urllib.parse.quote(doi, safe="")
    with urllib.request.urlopen(request(api), timeout=120) as r:
        data = json.load(r)
    for f in data["data"]["latestVersion"]["files"]:
        df = f["dataFile"]
        if df["filename"] == filename and (directory is None or f.get("directoryLabel") == directory):
            return int(df["id"])
    return None


def download_dataverse(doi: str, filename: str, dest: Path, directory: str | None = None) -> bool:
    try:
        fid = dataverse_file_id(doi, filename, directory)
    except Exception as e:
        print(f"FAILED Dataverse lookup {doi} {filename}: {type(e).__name__}: {e}")
        return False
    if fid is None:
        print(f"FAILED Dataverse lookup {doi}: file not found: {directory or ''}/{filename}")
        return False
    return download(f"https://dataverse.harvard.edu/api/access/datafile/{fid}", dest)


class InputParser(html.parser.HTMLParser):
    def __init__(self):
        super().__init__()
        self.inputs: dict[str, str] = {}

    def handle_starttag(self, tag, attrs):
        if tag != "input":
            return
        d = dict(attrs)
        name = d.get("name")
        if name:
            self.inputs[name] = d.get("value", "")


def download_vparty(dest: Path) -> bool:
    """Download V-Party R zip through V-Dem's required form.

    V-Dem requires email, gender, privacy acceptance, and format selection. We do
    not invent personal info: set PARTY2D_VDEM_EMAIL and PARTY2D_VDEM_GENDER
    (woman/man/non_binary/empty for prefer-not-to-say) to enable this download.
    """
    email = os.environ.get("PARTY2D_VDEM_EMAIL")
    if not email:
        print("SKIP V-Party: set PARTY2D_VDEM_EMAIL to use V-Dem's required download form")
        return False
    gender = os.environ.get("PARTY2D_VDEM_GENDER", "")
    page = "https://www.v-dem.net/data/v-party-dataset/country-party-date-v2/"
    cj = http.cookiejar.CookieJar()
    opener = urllib.request.build_opener(urllib.request.HTTPCookieProcessor(cj))
    try:
        with opener.open(request(page), timeout=60) as r:
            html = r.read().decode("utf-8", "replace")
    except Exception as e:
        print(f"FAILED V-Party form load: {type(e).__name__}: {e}")
        return False
    p = InputParser(); p.feed(html)
    csrf = p.inputs.get("csrfmiddlewaretoken", "")
    fields = {
        "csrfmiddlewaretoken": csrf,
        "email": email,
        "gender": gender,
        "accept_terms": "on",
        "dataset_file": "17",  # R format in current V-Dem form
        "website": "",
    }
    data = urllib.parse.urlencode(fields).encode()
    headers = {
        "Content-Type": "application/x-www-form-urlencoded",
        "Referer": page,
        "X-CSRFToken": csrf,
        "X-Requested-With": "XMLHttpRequest",
    }
    dest.parent.mkdir(parents=True, exist_ok=True)
    tmp = dest.with_suffix(".zip")
    print(f"Submitting V-Dem form -> {tmp}")
    try:
        with opener.open(request(page + "#dataset-download", data=data, headers=headers), timeout=180) as r:
            payload = r.read()
            ctype = r.headers.get("content-type", "")
        # The AJAX form may return JSON with a download URL, or the file itself.
        if b"download" in payload[:1000] and "json" in ctype:
            js = json.loads(payload.decode("utf-8"))
            url = js.get("download_url") or js.get("url") or js.get("file")
            if url and url.startswith("/"):
                url = urllib.parse.urljoin(page, url)
            if not url:
                print(f"FAILED V-Party form response did not include download URL: {js}")
                return False
            if not download(url, tmp, overwrite=True, headers={"Referer": page}):
                return False
        else:
            text = payload.decode("utf-8", "replace")
            m = re.search(r'href="([^"]*CPD_V-Party_R_v2\.zip[^"]*)"', text)
            if m:
                url = urllib.parse.urljoin(page, m.group(1))
                if not download(url, tmp, overwrite=True, headers={"Referer": page}):
                    return False
            else:
                tmp.write_bytes(payload)
        if zipfile.is_zipfile(tmp):
            with zipfile.ZipFile(tmp) as z:
                r_files = [n for n in z.namelist() if n.lower().endswith((".rds", ".rda", ".rdata"))]
                if not r_files:
                    print("FAILED V-Party zip contains no R data file")
                    return False
                member = r_files[0]
                with z.open(member) as src, dest.open("wb") as out:
                    shutil.copyfileobj(src, out)
            print(f"Extracted V-Party R data: {dest}")
            return True
        print("FAILED V-Party download was not a zip file")
        return False
    except Exception as e:
        print(f"FAILED V-Party form submit: {type(e).__name__}: {e}")
        return False


def download_manifesto(dest: Path) -> bool:
    key = os.environ.get("MANIFESTO_API_KEY") or os.environ.get("PARTY2D_MANIFESTO_API_KEY")
    if not key:
        print("SKIP Manifesto: set MANIFESTO_API_KEY or PARTY2D_MANIFESTO_API_KEY")
        return False
    url = "https://manifesto-project.wzb.eu/api/v1/get_core?" + urllib.parse.urlencode({"key": "MPDS2025a", "raw": "true"})
    return download(
        url,
        dest,
        overwrite=True,
        headers={"Referer": "https://manifesto-project.wzb.eu/datasets", "API_KEY": key},
    )


def write_report(status: dict[str, bool]) -> None:
    rd = report_dir(); rd.mkdir(parents=True, exist_ok=True)
    path = rd / "download_sources_report.md"
    lines = ["# Download sources report", "", "| source | status |", "| --- | --- |"]
    for k, v in status.items():
        lines.append(f"| {k} | {'ok' if v else 'missing/failed'} |")
    path.write_text("\n".join(lines) + "\n", encoding="utf-8")
    print(f"Wrote download report: {path}")


def main() -> int:
    root = raw_dir()
    root.mkdir(parents=True, exist_ok=True)
    status: dict[str, bool] = {}

    status["PolDem"] = download("https://poldem.eui.eu/downloads/cosa/poldem-election_all.csv", root / "poldem" / "poldem-election_all.csv")
    status["PartyFacts external parties"] = download("https://partyfacts.herokuapp.com/download/external-parties-csv/", root / "partyfacts" / "partyfacts-external-parties.csv")
    status["Manifesto MPDS 2025a"] = download_manifesto(root / "manifesto" / "MPDataset_MPDS2025a.csv")

    ches_base = "https://www.chesdata.eu/s/"
    ches_files = {
        "1999-2019_CHES_dataset_meansv3.csv": "1999-2019_CHES_dataset_means(v3).csv",
        "1999-2024_CHES_dataset_meansV2-3k4l.csv": "1999-2024_CHES_dataset_meansV2-3k4l.csv",
        "CHES_2024_final_v2.csv": "CHES_2024_final_v2.csv",
        "CHES_2024_ALL_Stacked_Expert.csv": "CHES_2024_expert_level.csv",
        "CHES_CA2023.csv": "CHES_CA2023.csv",
        "CHES_CA2023_expert-level.csv": "CHES_CA2023_expert_level.csv",
        "ches_la_2020_aggregate_level_v01.csv": "ches_la_2020_aggregate_level_v01.csv",
        "ches_la_2020_expert_level_v01.csv": "CHES_LA2020_expert_level.csv",
        "CHES_ISRAEL_means_2021_2022.csv": "CHES_ISRAEL_means_2021_2022.csv",
        "CHES_ISRAEL_expert_level_2021_2022.csv": "CHES_IL_expert_level.csv",
    }
    for remote, local in ches_files.items():
        status[f"CHES {local}"] = download(ches_base + urllib.parse.quote(remote), root / "ches" / local)

    status["POPPA integrated v2"] = download_dataverse("doi:10.7910/DVN/RMQREQ", "poppa_integrated_v2.rds", root / "poppa" / "poppa_integrated_v2.rds", "final_data_v2")
    status["GPS 2019 party"] = download_dataverse("doi:10.7910/DVN/WMGTNS", "Global Party Survey by Party SPSS V2_1_Apr_2020-2.tab", root / "gps" / "Global Party Survey by Party SPSS V2_1_Apr_2020-2.tab")
    status["V-Party"] = download_vparty(root / "vparty" / "V-Dem-CPD-Party-V2.rds")

    write_report(status)
    return 0 if all(status.values()) else 2


if __name__ == "__main__":
    raise SystemExit(main())
