#!/usr/bin/env python3
"""Importeer alleen benodigde POI's uit een OSM PBF naar PostGIS.
Veel lichter dan volledige osm2pgsql-import voor heel Nederland.
"""

import argparse
import os

import osmium
import psycopg2
from psycopg2.extras import execute_values

DB_DSN = os.environ.get(
    "ADRES_ANALYSE_DB_DSN",
    "host=127.0.0.1 port=5433 dbname=adresanalyse user=adres password=adres",
)

SCHEMA_SQL = """
CREATE EXTENSION IF NOT EXISTS postgis;
DROP TABLE IF EXISTS osm_pois;
CREATE TABLE osm_pois (
  id bigserial PRIMARY KEY,
  osm_type text NOT NULL,
  osm_id bigint NOT NULL,
  category text NOT NULL,
  name text NOT NULL,
  brand text,
  operator text,
  lon double precision NOT NULL,
  lat double precision NOT NULL,
  geom geometry(Point, 4326) NOT NULL
);
CREATE INDEX osm_pois_category_geom_gix ON osm_pois USING GIST (geom);
CREATE INDEX osm_pois_category_idx ON osm_pois (category);
CREATE INDEX osm_pois_name_idx ON osm_pois (lower(name));

CREATE OR REPLACE VIEW poi_supermarkten AS
SELECT osm_id, name, 'supermarkt'::text AS category, 'supermarket'::text AS shop, NULL::text AS amenity, brand, operator, geom
FROM osm_pois WHERE category = 'supermarkt';

CREATE OR REPLACE VIEW poi_scholen AS
SELECT osm_id, name, 'school'::text AS category, NULL::text AS shop, 'school'::text AS amenity, brand, operator, geom
FROM osm_pois WHERE category = 'school';
"""


def poi_from_tags(tags):
    if tags.get("shop") == "supermarket":
        return "supermarkt", tags.get("name") or tags.get("brand") or "Supermarkt"
    if tags.get("amenity") == "school":
        return "school", tags.get("name") or "School"
    return None, None


class PoiHandler(osmium.SimpleHandler):
    def __init__(self, conn, batch_size=1000):
        super().__init__()
        self.conn = conn
        self.batch_size = batch_size
        self.rows = []
        self.count = 0
        self.factory = osmium.geom.WKBFactory()

    def flush(self):
        if not self.rows:
            return
        with self.conn.cursor() as cur:
            execute_values(
                cur,
                """
                INSERT INTO osm_pois (osm_type, osm_id, category, name, brand, operator, lon, lat, geom)
                VALUES %s
                """,
                self.rows,
            )
        self.conn.commit()
        self.rows.clear()

    def add_row(self, osm_type, osm_id, tags, lon, lat):
        category, name = poi_from_tags(tags)
        if not category:
            return
        self.rows.append((
            osm_type,
            int(osm_id),
            category,
            name,
            tags.get("brand"),
            tags.get("operator"),
            float(lon),
            float(lat),
            f"SRID=4326;POINT({float(lon)} {float(lat)})",
        ))
        self.count += 1
        if len(self.rows) >= self.batch_size:
            self.flush()

    def node(self, n):
        if not n.location.valid():
            return
        self.add_row("node", n.id, n.tags, n.location.lon, n.location.lat)

    def area(self, a):
        category, _ = poi_from_tags(a.tags)
        if not category:
            return
        try:
            wkb = self.factory.create_multipolygon(a)
            with self.conn.cursor() as cur:
                cur.execute("SELECT ST_X(p), ST_Y(p) FROM (SELECT ST_PointOnSurface(ST_GeomFromWKB(%s, 4326)) p) s", (psycopg2.Binary(wkb),))
                lon, lat = cur.fetchone()
            self.add_row("area", a.orig_id(), a.tags, lon, lat)
        except Exception:
            return


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("pbf")
    args = parser.parse_args()

    with psycopg2.connect(DB_DSN) as conn:
        with conn.cursor() as cur:
            cur.execute(SCHEMA_SQL)
        conn.commit()
        handler = PoiHandler(conn)
        handler.apply_file(args.pbf, locations=True)
        handler.flush()
        with conn.cursor() as cur:
            cur.execute("ANALYZE osm_pois")
        conn.commit()
        print(f"Imported {handler.count} POIs into osm_pois")


if __name__ == "__main__":
    main()