#!/usr/bin/env python3
"""
Convert county address shapefile (ZIP) to OSM-formatted GeoJSON.

Applies LIFECYCLE filtering, CRS conversion, field mapping via qgis-functions,
and street name exceptions from /data/exceptions.yml.

Usage:
    python convert-addresses.py /data/sumter/addresses.shp.zip /data/sumter/county-addresses.geojson
    python convert-addresses.py /data/lake/addresses.shp.zip /data/lake/county-addresses.geojson
"""

import argparse
import shutil
import sys
import zipfile
from pathlib import Path

import geopandas as gpd
import pandas as pd
import importlib
import warnings

warnings.filterwarnings('ignore')

qgis_functions = importlib.import_module("qgis-functions")


def load_exceptions(exceptions_path='/data/exceptions.yml'):
    path = Path(exceptions_path)
    if not path.exists():
        return {}
    import yaml
    with open(path) as f:
        data = yaml.safe_load(f) or {}
    return {
        item['from']: item['to']
        for item in data.get('corrections', [])
        if 'from' in item and 'to' in item
    }


def process_address_fields(gdf, exceptions):
    """Map county shapefile fields to OSM address schema."""
    processed = gdf.copy()
    mapping = {}

    # House number
    for field in ['ADD_NUM', 'AddressNum', 'ADDRESS_NUM', 'HOUSE_NUM']:
        if field in processed.columns:
            series = pd.to_numeric(processed[field], errors='coerce')
            mapping['addr:housenumber'] = series.round().astype('Int64')
            break

    # Unit
    for field in ['UNIT', 'UnitNumber', 'UNIT_NUM', 'APT']:
        if field in processed.columns:
            series = processed[field].copy().replace(['nan', 'None', '', None], None)
            mapping['addr:unit'] = series.where(series.notna(), None)
            break

    # Street name
    if 'SADD' in processed.columns:
        # Sumter: full address string in SADD
        mapping['addr:street'] = [
            qgis_functions.title(qgis_functions.getstreetfromaddress(str(v), None, None))
            if pd.notna(v) else None
            for v in processed['SADD']
        ]
    elif 'FullAddres' in processed.columns:
        # Lake: full address string in FullAddres
        mapping['addr:street'] = [
            qgis_functions.title(qgis_functions.getstreetfromaddress(str(v), None, None))
            if pd.notna(v) else None
            for v in processed['FullAddres']
        ]
    elif 'BaseStreet' in processed.columns:
        # Lake alternative: assemble from components
        street_names = []
        for _, row in processed.iterrows():
            parts = []
            for col in ['PrefixDire', 'PrefixType', 'BaseStreet', 'SuffixType']:
                if col in row and pd.notna(row[col]):
                    parts.append(str(row[col]).strip())
            street_names.append(qgis_functions.title(' '.join(parts)) if parts else None)
        mapping['addr:street'] = street_names

    # Apply street name exceptions
    if 'addr:street' in mapping and exceptions:
        mapping['addr:street'] = [
            exceptions.get(name, name) if name is not None else None
            for name in mapping['addr:street']
        ]

    # City
    for field in ['POST_COMM', 'PostalCity', 'CITY', 'Jurisdicti']:
        if field in processed.columns:
            mapping['addr:city'] = [
                qgis_functions.title(str(v)) if pd.notna(v) else None
                for v in processed[field]
            ]
            break

    # Postcode
    for field in ['POST_CODE', 'ZipCode', 'ZIP', 'POSTAL_CODE']:
        if field in processed.columns:
            series = pd.to_numeric(processed[field], errors='coerce')
            mapping['addr:postcode'] = series.round().astype('Int64')
            break

    mapping['addr:state'] = 'FL'

    for key, value in mapping.items():
        processed[key] = value

    return processed


def convert(zip_path, output_path):
    zip_path = Path(zip_path)
    output_path = Path(output_path)

    # Skip if output is newer than input
    if (output_path.exists() and zip_path.exists() and
            output_path.stat().st_mtime > zip_path.stat().st_mtime):
        print(f"Output is up to date: {output_path}")
        return

    print(f"Converting {zip_path} ...")

    exceptions = load_exceptions()

    temp_dir = zip_path.parent / "temp_extract"
    temp_dir.mkdir(exist_ok=True)
    try:
        with zipfile.ZipFile(zip_path, 'r') as zf:
            zf.extractall(temp_dir)

        shp_files = list(temp_dir.glob("*.shp"))
        if not shp_files:
            print("Error: no .shp file found in ZIP", file=sys.stderr)
            sys.exit(1)

        gdf = gpd.read_file(shp_files[0])

        # LIFECYCLE filter (Sumter addresses use 'Current'; Lake has no such field)
        LIFECYCLE_FIELD = 'LIFECYCLE'
        ACTIVE_VALUE = 'Current'
        if LIFECYCLE_FIELD in gdf.columns:
            before = len(gdf)
            gdf = gdf[gdf[LIFECYCLE_FIELD] == ACTIVE_VALUE].copy().reset_index(drop=True)
            filtered = before - len(gdf)
            if filtered:
                print(f"Filtered out {filtered} non-active addresses (LIFECYCLE != '{ACTIVE_VALUE}')")
            else:
                print(f"All {len(gdf)} addresses are active")

        # CRS conversion
        if gdf.crs and gdf.crs != 'EPSG:4326':
            print(f"Converting CRS from {gdf.crs} to EPSG:4326")
            gdf = gdf.to_crs('EPSG:4326')

        gdf = process_address_fields(gdf, exceptions)

        # Points only, must have a house number
        gdf = gdf[gdf.geometry.type == 'Point'].copy()
        gdf = gdf[gdf['addr:housenumber'].notna()].copy()

        osm_fields = ['addr:housenumber', 'addr:unit', 'addr:street',
                      'addr:city', 'addr:postcode', 'addr:state']
        keep = [f for f in osm_fields if f in gdf.columns]
        gdf = gdf[keep + ['geometry']]

        output_path.parent.mkdir(parents=True, exist_ok=True)
        gdf.to_file(output_path, driver='GeoJSON')
        print(f"Saved {len(gdf)} addresses to {output_path}")

    finally:
        if temp_dir.exists():
            shutil.rmtree(temp_dir)


def main():
    parser = argparse.ArgumentParser(description="Convert county address shapefile to OSM-formatted GeoJSON")
    parser.add_argument('input_zip', help='Path to county address ZIP (containing .shp)')
    parser.add_argument('output_geojson', help='Output GeoJSON path')
    args = parser.parse_args()
    convert(args.input_zip, args.output_geojson)


if __name__ == '__main__':
    main()