import argparse
from pathlib import Path

import httpx
import polars as pl

# UK Parliament publishes candidate-level results for the 2024 General Election.
# One row per candidate per constituency — we aggregate to per-constituency stats.
URL = "https://electionresults.parliament.uk/general-elections/6/candidacies.csv"

# Map party names to a smaller set for vote share columns.
# Only parties that won seats in England are kept; the rest become "Other parties".
PARTY_MAP = {
    "Labour": "Labour",
    "Conservative": "Conservative",
    "Liberal Democrat": "Liberal Democrat",
    "Reform UK": "Reform UK",
    "Green Party": "Green",
}
PARTY_GROUPS = [
    "Labour",
    "Conservative",
    "Liberal Democrat",
    "Reform UK",
    "Green",
    "Other parties",
]


def _convert_results(df: pl.DataFrame) -> pl.DataFrame:
    # Filter to England only (constituency codes starting with E14)
    df = df.filter(pl.col("Constituency geographic code").str.starts_with("E14"))

    # Map party names to our output groups
    df = df.with_columns(
        pl.col("Main party name")
        .replace_strict(PARTY_MAP, default="Other parties")
        .alias("party_group"),
    )

    # ── Per-constituency turnout ──
    turnout = df.filter(pl.col("Candidate result position") == 1).select(
        pl.col("Constituency geographic code").alias("pcon"),
        (pl.col("Election valid vote count") / pl.col("Electorate") * 100)
        .round(1)
        .alias("turnout_pct"),
    )

    # ── Per-party vote share percentages ──
    # Sum votes per party group per constituency, then pivot to wide format
    party_votes = (
        df.group_by("Constituency geographic code", "party_group")
        .agg(pl.col("Candidate vote count").sum())
        .rename({"Constituency geographic code": "pcon"})
    )
    total_votes = (
        df.group_by("Constituency geographic code")
        .agg(pl.col("Candidate vote count").sum().alias("total_votes"))
        .rename({"Constituency geographic code": "pcon"})
    )
    party_pct = (
        party_votes.join(total_votes, on="pcon")
        .with_columns(
            (pl.col("Candidate vote count") / pl.col("total_votes") * 100)
            .round(1)
            .alias("vote_pct"),
        )
        .pivot(on="party_group", index="pcon", values="vote_pct")
    )

    # Rename columns to "% Party" format
    rename_map = {col: f"% {col}" for col in party_pct.columns if col != "pcon"}
    party_pct = party_pct.rename(rename_map)
    for party in PARTY_GROUPS:
        col = f"% {party}"
        if col not in party_pct.columns:
            party_pct = party_pct.with_columns(pl.lit(0.0).alias(col))
    party_pct = party_pct.with_columns(
        [pl.col(f"% {party}").fill_null(0.0) for party in PARTY_GROUPS]
    )

    # Join turnout with party vote shares
    return turnout.join(party_pct, on="pcon", how="left")


def download_and_convert(output_path: Path) -> None:
    print("Downloading 2024 General Election results...")
    response = httpx.get(URL, follow_redirects=True, timeout=60)
    response.raise_for_status()

    df = pl.read_csv(response.content)
    print(f"Raw shape: {df.shape}")

    result = _convert_results(df)

    print(f"Constituencies: {result.height}")
    print(f"Columns: {result.columns}")

    output_path.parent.mkdir(parents=True, exist_ok=True)
    result.write_parquet(output_path, compression="zstd")
    print(f"Saved to {output_path}")


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Download 2024 General Election results by constituency"
    )
    parser.add_argument(
        "--output", type=Path, required=True, help="Output parquet file path"
    )
    args = parser.parse_args()
    download_and_convert(args.output)


if __name__ == "__main__":
    main()