import argparse
import tempfile
import polars as pl
from pathlib import Path

from pipeline.utils import download, extract_zip

# Ofcom Connected Nations 2025 - Fixed broadband performance (output area & local authority level)
# Source: https://www.ofcom.org.uk/phones-and-broadband/coverage-and-speeds/connected-nations-20252/data-downloads-2025
PERFORMANCE_URL = "https://www.ofcom.org.uk/siteassets/resources/documents/research-and-data/multi-sector/infrastructure-research/connected-nations-2025/202507_fixed_broadband_coverage_r01.zip?v=407830"


def convert_to_parquet(extract_dir: Path, parquet_path: Path) -> None:
    # Find CSV files in the extracted directory
    csv_files = list(extract_dir.rglob("*.csv"))
    if not csv_files:
        raise FileNotFoundError(f"No CSV files found in {extract_dir}")

    print(f"Found {len(csv_files)} CSV files: {[f.name for f in csv_files]}")

    frames = []
    for csv_file in sorted(csv_files):
        print(f"Reading {csv_file.name}...")
        df = pl.read_csv(csv_file, infer_schema_length=10000, encoding="utf8-lossy")
        print(f"  Shape: {df.shape}")
        frames.append(df)

    combined = pl.concat(frames, how="diagonal_relaxed")
    print(f"Combined shape: {combined.shape}")

    parquet_path.parent.mkdir(parents=True, exist_ok=True)
    combined.write_parquet(parquet_path, compression="zstd")
    print(f"Saved {parquet_path} ({combined.shape[0]} rows)")


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Download Ofcom broadband performance data"
    )
    parser.add_argument(
        "--output",
        type=Path,
        required=True,
        help="Output parquet file path",
    )
    args = parser.parse_args()

    with tempfile.TemporaryDirectory() as cache_dir:
        cache = Path(cache_dir)
        zip_path = cache / "broadband_performance.zip"
        extract_dir = cache / "extracted"
        extracted_again_dir = cache / "extracted-again"

        download(PERFORMANCE_URL, zip_path)
        extract_zip(zip_path, extract_dir)
        extract_zip(
            extract_dir
            / "202507_fixed_coverage_r01"
            / "202507_fixed_pc_coverage_r01.zip",
            extracted_again_dir,
        )

        convert_to_parquet(extracted_again_dir, args.output)


if __name__ == "__main__":
    main()