"""US Census Bureau ACS 5-year population extractor.

Fetches city-level (Census place) population from the American Community Survey
5-year estimates. Requires a free API key from api.census.gov.

Env var: CENSUS_API_KEY  (register free at https://api.census.gov/data/key_signup.html)

Landing: {LANDING_DIR}/census_usa/{year}/{month}/acs5_places.json.gz
Output:  {"rows": [{"city_name": "Los Angeles", "state_fips": "06",
                    "place_fips": "0644000", "population": 3990456,
                    "ref_year": 2023, "country_code": "US"}], "count": N}
"""

import json
import os
import sqlite3
from pathlib import Path

import niquests

from ._shared import HTTP_TIMEOUT_SECONDS, run_extractor, setup_logging
from .utils import landing_path, skip_if_current, write_gzip_atomic

logger = setup_logging("padelnomics.extract.census_usa")

EXTRACTOR_NAME = "census_usa"

# ACS 5-year estimates, 2023 vintage — refreshed annually by Census Bureau.
# B01003_001E = total population; NAME = place name + state.
ACS_URL = (
    "https://api.census.gov/data/2023/acs/acs5"
    "?get=B01003_001E,NAME&for=place:*&in=state:*"
)

REF_YEAR = 2023
MIN_POPULATION = 10_000
MAX_RETRIES = 2


def _parse_city_name(full_name: str) -> str:
    """Extract city name from Census place name.

    Examples:
      'Los Angeles city, California' → 'Los Angeles'
      'New York city, New York'      → 'New York'
      'Miami city, Florida'          → 'Miami'
    """
    # Take everything before the first comma
    before_comma = full_name.split(",")[0].strip()
    # Strip common suffixes: ' city', ' town', ' CDP', ' borough', ' village'
    for suffix in (" city", " town", " CDP", " borough", " village", " municipality"):
        if before_comma.lower().endswith(suffix):
            before_comma = before_comma[: -len(suffix)].strip()
            break
    return before_comma


def extract(
    landing_dir: Path,
    year_month: str,
    conn: sqlite3.Connection,
    session: niquests.Session,
) -> dict:
    """Fetch ACS 5-year place population. Skips if already run this month."""
    api_key = os.environ.get("CENSUS_API_KEY", "").strip()
    if not api_key:
        logger.warning("CENSUS_API_KEY not set — writing empty placeholder so SQLMesh models can run")
        year, month = year_month.split("/")
        dest_dir = landing_path(landing_dir, "census_usa", year, month)
        dest = dest_dir / "acs5_places.json.gz"
        if not dest.exists():
            write_gzip_atomic(dest, b'{"rows": [], "count": 0}')
        return {"files_written": 0, "files_skipped": 1, "bytes_written": 0}

    # Skip if we already have data for this month (annual data, monthly cursor)
    skip = skip_if_current(conn, EXTRACTOR_NAME, year_month)
    if skip:
        logger.info("already have data for %s — skipping", year_month)
        return skip

    year, month = year_month.split("/")
    url = f"{ACS_URL}&key={api_key}"

    logger.info("GET ACS 5-year places (vintage %d)", REF_YEAR)
    resp = session.get(url, timeout=HTTP_TIMEOUT_SECONDS * 2)
    resp.raise_for_status()

    raw = resp.json()
    assert isinstance(raw, list) and len(raw) > 1, "ACS response must be a non-empty list"

    # First row is headers: ["B01003_001E", "NAME", "state", "place"]
    headers = raw[0]
    assert "B01003_001E" in headers, f"Population column missing from ACS response: {headers}"
    pop_idx = headers.index("B01003_001E")
    name_idx = headers.index("NAME")
    state_idx = headers.index("state")
    place_idx = headers.index("place")

    rows: list[dict] = []
    for row in raw[1:]:
        try:
            population = int(row[pop_idx])
        except (ValueError, TypeError):
            continue
        if population < MIN_POPULATION:
            continue
        full_name = row[name_idx]
        city_name = _parse_city_name(full_name)
        if not city_name:
            continue
        state_fips = row[state_idx]
        place_fips = state_fips + row[place_idx]
        rows.append({
            "city_name": city_name,
            "state_fips": state_fips,
            "place_fips": place_fips,
            "population": population,
            "ref_year": REF_YEAR,
            "country_code": "US",
        })

    assert len(rows) > 2000, f"Expected >2000 US cities ≥10K pop, got {len(rows)} — parse may have failed"
    logger.info("parsed %d US cities with population ≥%d", len(rows), MIN_POPULATION)

    dest_dir = landing_path(landing_dir, "census_usa", year, month)
    dest = dest_dir / "acs5_places.json.gz"
    payload = json.dumps({"rows": rows, "count": len(rows)}).encode()
    bytes_written = write_gzip_atomic(dest, payload)
    logger.info("written %s bytes compressed", f"{bytes_written:,}")

    return {
        "files_written": 1,
        "files_skipped": 0,
        "bytes_written": bytes_written,
        "cursor_value": year_month,
    }


def main() -> None:
    run_extractor(EXTRACTOR_NAME, extract)


if __name__ == "__main__":
    main()
