# app.py — AP Elections API simulator for Hugging Face Spaces (Docker)
# Simulates:  GET /v2/elections/{date}?statepostal=XX&raceTypeId=G&raceId=0&level=ru
# Returns XML containing <ReportingUnit Name="…"><Candidate …/></ReportingUnit> per county.
#
# Design:
# - On startup, downloads US-Atlas counties TopoJSON to build state→counties list.
# - For each requested state, emits ReportingUnit for each county with AP-like name variants.
# - Vote counts are deterministic but tick upward over time to look "live".
#
# Usage with your app:
#   In your wrapper (BASE_URL), change to:  https://<your-space>.hf.space/v2/elections
#
# Notes:
# - We ignore x-api-key. Your wrapper still sends it; that's fine.  (We don't enforce it.)
# - We accept any date string; we don't filter by raceTypeId/raceId/level (but we keep them for parity).

import os, time, math, json, re, asyncio, hashlib
from typing import Dict, List, Tuple
import httpx
from fastapi import FastAPI, Query, Request
from fastapi.responses import PlainTextResponse, Response

app = FastAPI(title="AP Elections API Simulator")

US_ATLAS_COUNTIES_URL = "https://cdn.jsdelivr.net/npm/us-atlas@3/counties-10m.json"

# Map state FIPS (2 digits) → USPS
STATE_FIPS_TO_USPS = {
    "01":"AL","02":"AK","04":"AZ","05":"AR","06":"CA","08":"CO","09":"CT","10":"DE","11":"DC",
    "12":"FL","13":"GA","15":"HI","16":"ID","17":"IL","18":"IN","19":"IA","20":"KS","21":"KY",
    "22":"LA","23":"ME","24":"MD","25":"MA","26":"MI","27":"MN","28":"MS","29":"MO","30":"MT",
    "31":"NE","32":"NV","33":"NH","34":"NJ","35":"NM","36":"NY","37":"NC","38":"ND","39":"OH",
    "40":"OK","41":"OR","42":"PA","44":"RI","45":"SC","46":"SD","47":"TN","48":"TX","49":"UT",
    "50":"VT","51":"VA","53":"WA","54":"WV","55":"WI","56":"WY",
    # Territories (we won't emit unless statepostal matches)
    "72":"PR"
}

# USPS → state FIPS
USPS_TO_STATE_FIPS = {v:k for k,v in STATE_FIPS_TO_USPS.items()}

# USPS that use "Parish" instead of "County"
PARISH_STATES = {"LA"}

# USPS that often use "city" county-equivalents
INDEPENDENT_CITY_STATES = {"VA"}  # (also some in MO e.g., St. Louis city)

# A few AP-style "wrong" name transforms to exercise your fixups.
# Applied after suffixing County/Parish/city when appropriate.
def apize_name(usps: str, canonical: str) -> str:
    n = canonical

    # Common Saint → St. contraction
    if n.startswith("Saint "):
        n = "St. " + n[6:]

    # Diacritics drops often seen
    n = n.replace("Doña", "Dona").replace("Niobrara", "Niobrara")  # keep example stable

    # Collapsed spacing variants occasionally seen
    n = re.sub(r"\bLa\s+Salle\b", "LaSalle", n)
    n = n.replace("DeKalb", "De Kalb")

    # MO "St. Louis city" is a city-county equivalent; keep 'city'
    # leave as-is if endswith('city')
    return n

def county_suffix(usps: str, name: str) -> str:
    # Heuristics based on state
    if usps in PARISH_STATES:
        # LA data in maps typically already contains "Parish". If not, add it.
        return f"{name}" if name.lower().endswith("parish") else f"{name} Parish"

    # Independent cities (e.g., VA) usually already end with "city" in map data
    # If not, default to County.
    if usps in INDEPENDENT_CITY_STATES and name.lower().endswith("city"):
        return name

    # Default: County
    return name if name.lower().endswith("county") else f"{name} County"

def seeded_rng_u32(seed: str) -> int:
    # Deterministic 32-bit integer from any string (FIPS)
    h = hashlib.blake2b(seed.encode("utf-8"), digest_size=4).digest()
    return int.from_bytes(h, "big")

def simulated_votes(fips: str) -> Tuple[int,int,int]:
    """
    Deterministic but 'live' votes.
    We derive base from FIPS, then add a minute-tick so numbers grow slowly.
    """
    base_seed = seeded_rng_u32(fips)
    r1 = (base_seed & 0xFFFF)
    r2 = ((base_seed >> 16) & 0xFFFF)

    # minute "tick" so values rise over time
    minute_tick = int(time.time() // 60)

    # Scale to plausible county sizes
    base_total = 2000 + (base_seed % 250000)  # up to ~250k baseline
    # skew rep/dem by seed; keep IND small
    rep = int(base_total * (0.35 + (r1 % 30)/100.0))  # 35–65% of base_total
    dem = base_total - rep
    ind = int(base_total * (0.01 + (r2 % 3)/100.0))   # ~1–4%

    # grow gently each minute
    growth = 50 + (base_seed % 50)  # 50–99 votes per minute
    rep += (growth * (minute_tick % 10)) // 2
    dem += (growth * (minute_tick % 10)) // 2
    ind += (growth * (minute_tick % 10)) // 10

    # cap at non-negative
    rep = max(rep, 0); dem = max(dem, 0); ind = max(ind, 0)
    return rep, dem, ind

# --- In-memory registry: USPS → List[(FIPS, canonical_name, ap_name)] ---
STATE_REGISTRY: Dict[str, List[Tuple[str,str,str]]] = {}

@app.on_event("startup")
async def bootstrap():
    # Build county registry from US-Atlas TopoJSON (same source your frontend uses).
    async with httpx.AsyncClient(timeout=30) as client:
        r = await client.get(US_ATLAS_COUNTIES_URL)
        r.raise_for_status()
        topo = r.json()

    # Extract geometries
    geoms = topo.get("objects", {}).get("counties", {}).get("geometries", [])
    # Some us-atlas builds place names under 'properties.name'. Others under 'properties.NAMELSAD' —
    # we check both and fall back to FIPS if missing.
    for g in geoms:
        fips = str(g.get("id", "")).zfill(5)
        props = g.get("properties", {}) or {}
        name = props.get("name") or props.get("NAMELSAD") or fips
        state_fips = fips[:2]
        usps = STATE_FIPS_TO_USPS.get(state_fips)
        if not usps:
            continue

        # Strip common suffixes from canonical to resemble your client map’s county label.
        canonical = re.sub(r"\s+(County|Parish|city)$", "", name)

        # Build an AP-style reporting name (adds County/Parish and applies quirks)
        apname = county_suffix(usps, canonical)
        apname = apize_name(usps, apname)

        STATE_REGISTRY.setdefault(usps, []).append((fips, canonical, apname))

    # Ensure deterministic order
    for usps in STATE_REGISTRY:
        STATE_REGISTRY[usps].sort(key=lambda t: t[0])

@app.get("/api/ping", response_class=PlainTextResponse)
def ping():
    return "pong"

@app.get("/v2/elections/{date}")
def elections_state_ru(
    request: Request,
    date: str,
    statepostal: str = Query(..., min_length=2, max_length=2),
    raceTypeId: str = Query("G"),
    raceId: str = Query("0"),
    level: str = Query("ru"),
):
    """
    Simulated AP endpoint. We ignore race filters but mirror the URL/shape.
    Response: XML with <ReportingUnit Name="..."><Candidate .../></ReportingUnit>
    Candidates: Trump (REP), Harris (DEM), Kennedy (IND)
    """
    usps = statepostal.upper()
    counties = STATE_REGISTRY.get(usps, [])

    # If we somehow don't have this state, return empty but valid XML (your code handles this).
    if not counties:
        xml = f'<ElectionResults Date="{date}" StatePostal="{usps}"></ElectionResults>'
        return Response(content=xml, media_type="application/xml")

    # Assemble XML
    parts = [f'<ElectionResults Date="{date}" StatePostal="{usps}">']
    for fips, canonical, apname in counties:
        rep, dem, ind = simulated_votes(fips)

        # Example AP attributes your wrapper reads: First, Last, Party, VoteCount
        parts.append(f'  <ReportingUnit Name="{apname}" FIPS="{fips}">')
        parts.append(f'    <Candidate First="Donald" Last="Trump"   Party="REP" VoteCount="{rep}"/>')
        parts.append(f'    <Candidate First="Kamala" Last="Harris"  Party="DEM" VoteCount="{dem}"/>')
        parts.append(f'    <Candidate First="Robert" Last="Kennedy" Party="IND" VoteCount="{ind}"/>')
        parts.append(  "  </ReportingUnit>")
    parts.append("</ElectionResults>")
    xml = "\n".join(parts)
    return Response(content=xml, media_type="application/xml")

# Local run (for dev): uvicorn app:app --reload
if __name__ == "__main__":
    import uvicorn, os
    uvicorn.run(app, host="0.0.0.0", port=int(os.getenv("PORT", "7860")))