#!/usr/bin/env python3
"""
Merge NSRDB cloudiness metrics into climate-data.csv.

Adds:
- cloudinessIndexPct

Polygon area-weighted values are used first when available; representative-point
values remain the fallback.
"""

from __future__ import annotations

import argparse
import csv
from pathlib import Path


DEFAULT_CLIMATE_DATA = Path("data/climate-data.csv")
DEFAULT_POLYGON_CLOUD_SUMMARY = Path("data/nrel/county_polygon_cloud_summary.csv")
DEFAULT_REPRESENTATIVE_POINT_CLOUD_SUMMARY = Path("data/nrel/county_representative_point_cloud_summary.csv")
METRIC_FIELD = "cloudinessIndexPct"
POLYGON_SOURCE_TAG = "nsrdb-polygon-area-weighted-cloudiness-tmy"
REPRESENTATIVE_POINT_SOURCE_TAG = "nsrdb-representative-point-cloudiness-tmy"
CLOUD_SOURCE_TAGS = {
    POLYGON_SOURCE_TAG,
    REPRESENTATIVE_POINT_SOURCE_TAG,
}


def load_cloud_values(path: Path) -> dict[str, str]:
    """Read cloudiness values keyed by county FIPS."""
    values: dict[str, str] = {}
    with path.open("r", encoding="utf-8-sig", newline="") as handle:
        reader = csv.DictReader(handle)
        fieldnames = reader.fieldnames or []
        required_fields = {"county_fips", METRIC_FIELD}
        missing_fields = sorted(required_fields - set(fieldnames))
        if missing_fields:
            raise ValueError(f"{path} is missing fields: {', '.join(missing_fields)}")

        for row in reader:
            county_fips = (row.get("county_fips") or "").strip().zfill(5)
            value = (row.get(METRIC_FIELD) or "").strip()
            if county_fips and value:
                values[county_fips] = value
    return values


def append_source_tag(source: str, tag: str) -> str:
    """Append a source tag if it is not already present."""
    parts = [
        part.strip()
        for part in source.split("+")
        if part.strip() and part.strip() not in CLOUD_SOURCE_TAGS
    ]
    if tag not in parts:
        parts.append(tag)
    return " + ".join(parts)


def ensure_field_after(fieldnames: list[str], field: str, after_field: str | None = None) -> None:
    """Insert a field into CSV field order if it is missing."""
    if field in fieldnames:
        return
    if after_field and after_field in fieldnames:
        fieldnames.insert(fieldnames.index(after_field) + 1, field)
        return
    insert_at = fieldnames.index("source") if "source" in fieldnames else len(fieldnames)
    fieldnames.insert(insert_at, field)


def merge_metric(
    climate_data: Path,
    polygon_cloud_summary: Path,
    representative_point_cloud_summary: Path,
) -> tuple[int, int, int, int]:
    """Merge cloudiness values into the app climate CSV."""
    polygon_cloudiness_by_fips = (
        load_cloud_values(polygon_cloud_summary)
        if polygon_cloud_summary.exists()
        else {}
    )
    representative_cloudiness_by_fips = load_cloud_values(representative_point_cloud_summary)

    with climate_data.open("r", encoding="utf-8", newline="") as handle:
        reader = csv.DictReader(handle)
        fieldnames = list(reader.fieldnames or [])
        rows = list(reader)

    if "countyFips" not in fieldnames:
        raise ValueError(f"{climate_data} is missing countyFips")

    ensure_field_after(fieldnames, METRIC_FIELD, "avgSolarGhiKwhM2Day")

    polygon_count = 0
    representative_count = 0
    missing_count = 0
    for row in rows:
        county_fips = (row.get("countyFips") or "").strip().zfill(5)
        polygon_value = polygon_cloudiness_by_fips.get(county_fips, "")
        representative_value = representative_cloudiness_by_fips.get(county_fips, "")
        value = polygon_value or representative_value
        row[METRIC_FIELD] = value
        if polygon_value:
            row["source"] = append_source_tag(row.get("source", ""), POLYGON_SOURCE_TAG)
            polygon_count += 1
        elif representative_value:
            row["source"] = append_source_tag(row.get("source", ""), REPRESENTATIVE_POINT_SOURCE_TAG)
            representative_count += 1
        else:
            missing_count += 1

    with climate_data.open("w", encoding="utf-8", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(rows)

    return len(rows), polygon_count, representative_count, missing_count


def main() -> int:
    parser = argparse.ArgumentParser(description="Merge NSRDB cloudiness metric into climate-data.csv.")
    parser.add_argument("--climate-data", type=Path, default=DEFAULT_CLIMATE_DATA)
    parser.add_argument("--polygon-cloud-summary", type=Path, default=DEFAULT_POLYGON_CLOUD_SUMMARY)
    parser.add_argument(
        "--representative-point-cloud-summary",
        type=Path,
        default=DEFAULT_REPRESENTATIVE_POINT_CLOUD_SUMMARY,
    )
    args = parser.parse_args()

    row_count, polygon_count, representative_count, missing_count = merge_metric(
        args.climate_data,
        args.polygon_cloud_summary,
        args.representative_point_cloud_summary,
    )
    print(f"Updated {args.climate_data} with {METRIC_FIELD}")
    print(f"Rows: {row_count}")
    print(f"Polygon area-weighted rows: {polygon_count}")
    print(f"Representative-point fallback rows: {representative_count}")
    print(f"Missing {METRIC_FIELD}: {missing_count}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())