"""
uc_wages_downloader.py
=======================
Download UC employee wage data from the ucinvestments/UC-wages repository
(https://github.com/ucinvestments/UC-wages) — a public dataset (not a
scraper/API) of one JSON file per campus per year, sourced from the UC
Annual Wage site (https://ucannualwage.ucop.edu/wage/).

No GitHub API calls are made (avoids the 60 req/hour unauthenticated rate
limit) — files are fetched directly from raw.githubusercontent.com, which
has no such limit. A combination that doesn't exist upstream (e.g. ASUCLA
before 2018) returns HTTP 404 and is logged + skipped rather than treated
as an error.

USAGE
-----
  # Download every campus for every year (2010..current year; anything not
  # yet published just 404-skips):
  python uc_wages_downloader.py

  # One campus, specific years:
  python uc_wages_downloader.py --years 2023 2024 --campuses UCB UCLA

  # Dry run (print URLs, download nothing):
  python uc_wages_downloader.py --dry-run

OUTPUT
------
One JSON file per campus per year, mirroring the upstream layout:
  uc_wages_data/<Campus>/wages_<year>.json
Already-downloaded files are left in place (idempotent).
"""

import argparse
import datetime
import logging
import sys
import time
from pathlib import Path

import requests

# ---------------------------------------------------------------------------
# LOCATIONS  (UC-wages folder name -> UC abbreviation, and back)
# ---------------------------------------------------------------------------
LOCATIONS = {
    "ASUCLA":     "ASUCLA",
    "UCB":        "Berkeley",
    "UCD":        "Davis",
    "UCI":        "Irvine",
    "UCLA":       "Los_Angeles",
    "UCM":        "Merced",
    "UCR":        "Riverside",
    "UCSD":       "San_Diego",
    "UCSF":       "San_Francisco",
    "UCSB":       "Santa_Barbara",
    "UCSC":       "Santa_Cruz",
    "UCOP":       "UCOP",
    "UC_SF_LAW":  "UC_SF_Law",
}

RAW_BASE = "https://raw.githubusercontent.com/ucinvestments/UC-wages/main/data"
DEFAULT_YEARS = list(range(2010, datetime.date.today().year + 1))

# ---------------------------------------------------------------------------
# LOGGING
# ---------------------------------------------------------------------------
logging.basicConfig(
    format="%(asctime)s [%(levelname)s] %(message)s",
    level=logging.INFO,
    handlers=[
        logging.FileHandler("uc_wages_downloader.log", encoding="utf-8"),
        logging.StreamHandler(sys.stdout),
    ],
)
log = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# DOWNLOAD
# ---------------------------------------------------------------------------

def file_url(campus: str, year: int) -> str:
    folder = LOCATIONS[campus]
    return f"{RAW_BASE}/{folder}/wages_{year}.json"


def out_path(campus: str, year: int, out_dir: Path) -> Path:
    folder = LOCATIONS[campus]
    return out_dir / folder / f"wages_{year}.json"


def ensure_file(campus: str, year: int, out_dir: Path,
                session: requests.Session,
                retries: int = 3, timeout: int = 60) -> Path | None:
    """
    Guarantee the JSON file for one campus/year is present locally.
    Returns the local path on success, None if the file doesn't exist
    upstream (404) or all retries were exhausted.

    Skips the request entirely if the file already exists (idempotent).
    """
    path = out_path(campus, year, out_dir)
    if path.exists():
        log.info("Already exists, skipping: %s", path)
        return path

    url = file_url(campus, year)
    log.info("Downloading %s %s ...", campus, year)

    for attempt in range(1, retries + 1):
        try:
            resp = session.get(url, timeout=timeout)
            if resp.status_code == 404:
                log.info("  %s %s: not available upstream (404), skipping.",
                         campus, year)
                return None
            resp.raise_for_status()

            path.parent.mkdir(parents=True, exist_ok=True)
            path.write_bytes(resp.content)
            log.info("  -> %s (%.1f MB)", path, len(resp.content) / 1_000_000)
            return path

        except requests.RequestException as exc:
            if attempt == retries:
                log.error("%s %s: download failed after %d attempts: %s",
                          campus, year, retries, exc)
                return None
            wait = 10 * attempt
            log.warning("%s %s: attempt %d failed (%s), retrying in %ds ...",
                       campus, year, attempt, exc, wait)
            time.sleep(wait)

    return None


def download_all(years: list[int], campuses: list[str], out_dir: Path,
                 retries: int = 3, timeout: int = 60) -> list[Path]:
    """Download (or reuse cached) files for every year x campus combination."""
    out_dir.mkdir(parents=True, exist_ok=True)
    session = requests.Session()
    session.headers.update({
        "User-Agent": "uc-wages-downloader/1.0",
    })

    paths = []
    for year in sorted(years):
        for campus in campuses:
            path = ensure_file(campus, year, out_dir, session,
                              retries=retries, timeout=timeout)
            if path:
                paths.append(path)
            time.sleep(0.3)  # polite pause between requests
    return paths


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser(
        description="Download UC employee wage JSON from ucinvestments/UC-wages.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )
    parser.add_argument(
        "--years", nargs="+", type=int, default=DEFAULT_YEARS,
        help=f"Calendar year(s) to download (default: 2010-{DEFAULT_YEARS[-1]}).",
    )
    parser.add_argument(
        "--campuses", nargs="+",
        choices=list(LOCATIONS.keys()), default=list(LOCATIONS.keys()),
        metavar="CAMPUS",
        help="Campuses to download (default: all). Choices: " +
             " ".join(LOCATIONS.keys()),
    )
    parser.add_argument(
        "--out-dir", default="uc_wages_data",
        help="Directory to write JSON files into (default: uc_wages_data/)",
    )
    parser.add_argument(
        "--dry-run", action="store_true",
        help="Print download URLs without downloading anything.",
    )
    parser.add_argument(
        "--timeout", type=int, default=60,
        help="HTTP timeout in seconds per request (default: 60)",
    )
    parser.add_argument(
        "--retries", type=int, default=3,
        help="Retry attempts per file on failure (default: 3)",
    )
    args = parser.parse_args()

    out_dir = Path(args.out_dir)

    if args.dry_run:
        print("Dry run - URLs that would be fetched:")
        for year in sorted(args.years):
            for campus in args.campuses:
                print(f"  {out_path(campus, year, out_dir)}  <-  {file_url(campus, year)}")
        return

    paths = download_all(args.years, args.campuses, out_dir,
                        retries=args.retries, timeout=args.timeout)
    log.info("Done. %d files available locally.", len(paths))


if __name__ == "__main__":
    main()