#!/usr/bin/env python3
#
# SPDX-FileCopyrightText: ⓒ 2025 Volker Krause <vkrause@kde.org>
# SPDX-License-Identifier: LGPL-2.0-or-later
#
# Generate country bounding box tables from OSM ISO 3166-1 boundaries.
# Outputs the embedded tables as code as well as GeoJSON for QA.
#

import argparse
import json
import os
import requests
import zipfile

parser = argparse.ArgumentParser(description='Generates country boundary boxes.')
parser.add_argument('--code', type=str, required=False, help='File to write the generated code to')
parser.add_argument('--geojson', type=str, required=False, help='File to write the GeoJSON output to')
arguments = parser.parse_args()

ISO3166_1_VERSION = '2024-12-09'
ISO3166_1_URL = f"https://volkerkrause.eu/~vkrause/iso3166-boundaries/iso3166-1-boundaries.geojson-{ISO3166_1_VERSION}.zip"


def iso3166FilePath(name):
    return os.path.join(os.path.dirname(__file__), name)


def downloadIso3166File(url, name):
    fileName = iso3166FilePath(name)
    if os.path.exists(fileName):
        return
    print(f"Downloading {url}…")
    r = requests.get(url)
    if r.status_code < 400:
        with open(fileName, 'wb') as f:
            f.write(r.content)
        with zipfile.ZipFile(fileName, 'r') as z:
            z.extractall(os.path.dirname(__file__))


def boundingBox(poly, bbox, filterBox):
    for coord in poly[0]:
        if coord[0] < filterBox[0][0] or coord[0] > filterBox[1][0] or coord[1] < filterBox[0][1] or coord[1] > filterBox[1][1]:
            continue
        bbox[0][0] = min(bbox[0][0], coord[0])
        bbox[0][1] = min(bbox[0][1], coord[1])
        bbox[1][0] = max(bbox[1][0], coord[0])
        bbox[1][1] = max(bbox[1][1], coord[1])
    return bbox


# special cases to deal with anti-meridian wraparounds and oversea territories we don't need for out purpose here
def filterBoundingBox(code):
    match (code):
        case 'AU': return ([110, -90], [165, 90])
        case 'FJ': return ([0, -90], [180, 90])
        case 'FR': return ([-10, 0], [180, 90])
        case 'KI': return ([0, -90], [180, 90])
        case 'NL': return ([0, 0], [180, 90])
        case 'NO': return ([-180, 0], [180, 90])
        case 'NZ': return ([0, -90], [180, 90])
        case 'RU': return ([0, -90], [180, 90])
        case 'TV': return ([0, -90], [180, 90])
        case 'US': return ([-180, 0], [0, 90])
    return ([-180, -90], [180, 90])


filtedCodes = ['AQ', 'UM']


# encode coordinate to uint16_t for code generation
def encodeCoordinate(coord):
    return int(coord * 100)


# download boundary data
downloadIso3166File(ISO3166_1_URL, f"iso3166-1-boundaries.geojson-{ISO3166_1_VERSION}.zip")
with open(iso3166FilePath('iso3166-1-boundaries.geojson')) as f:
    boundaries = json.load(f)


# generate bounding boxes
countries = {}
for feature in boundaries['features']:
    code = feature['properties']['ISO3166-1']
    if code in filtedCodes:
        continue
    if code not in countries:
        countries[code] = ([180, 90], [-180, -90])
    if feature['geometry']['type'] == 'MultiPolygon':
        for poly in feature['geometry']['coordinates']:
            countries[code] = boundingBox(poly, countries[code], filterBoundingBox(code))
    elif feature['geometry']['type'] == 'Polygon':
        countries[code] = boundingBox(feature['geometry']['coordinates'], countries[code])

# GeoJSON output
if arguments.geojson:
    output = {}
    output['type'] = 'FeatureCollection'
    output['name'] = 'Country bounding boxes'
    output['features'] = []
    for code, bbox in countries.items():
        feature = {}
        feature['type'] = 'Feature'
        feature['properties'] = {}
        feature['properties']['name'] = code
        feature['geometry'] = {}
        feature['geometry']['type'] = 'Polygon'
        feature['geometry']['coordinates'] = [[bbox[0], [bbox[1][0], bbox[0][1]], bbox[1], [bbox[0][0], bbox[1][1]], bbox[0]]]
        output['features'].append(feature)

    with open(arguments.geojson, 'w') as f:
        json.dump(output, f)


# generate code
if arguments.code:
    codes = [code for code in countries]
    codes.sort()

    with open(arguments.code, 'w') as out:
        out.write("""/*
  SPDX-FileCopyrightText: ⓒ OpenStreetMap community
  SPDX-License-Identifier: ODbL-1.0

  DO NOT EDIT!
  This file is autogenerated by src/lib/knowledgedb/generate-country-bounding-boxes.py.
*/

#include "countryboundingbox_p.h"

namespace KPublicTransport {

constexpr const CountryBoundingBoxData country_bounding_boxes[] = {
""")
        for code in codes:
            bbox = countries[code]
            out.write(f"    {{{{'{code[0]}','{code[1]}'}}, {encodeCoordinate(bbox[0][0])}, {encodeCoordinate(bbox[0][1]) + 1}, {encodeCoordinate(bbox[1][0])}, {encodeCoordinate(bbox[1][1]) + 1} }},\n")

        out.write("""};
}
""")
