import duckdb
import jinja2
from pathlib import Path
import click
from tqdm import tqdm
from posix import mkdir
from datetime import datetime, date

sql_tmpl = jinja2.Template('''
LOAD spatial;

CREATE TABLE IF NOT EXISTS city_road(
    city_name TEXT,
    relation_id BIGINT,
    points DOUBLE[][],
    tags MAP(VARCHAR, VARCHAR),
    highway TEXT,
    extract_date Date,
    length DOUBLE
);

CREATE TEMPORARY VIEW osm_quebec AS (
    SELECT * FROM ST_READOSM ('{{ file_name }}')
);

CREATE TEMPORARY VIEW interresting_ways AS (
    SELECT * FROM osm_quebec
    WHERE tags['highway'][1] IS NOT NULL
);

CREATE TEMPORARY VIEW relations AS (
    SELECT id,
           tags,
           UNNEST(refs) AS refs,
           UNNEST(RANGE(LEN(refs))) as index
    FROM interresting_ways
);

CREATE TEMPORARY VIEW road_points AS (
    SELECT osm_quebec.id AS id,
           osm_quebec.lon AS lon,
           osm_quebec.lat AS lat,
           relations.tags AS tags,
           relations.index AS relation_index,
           relations.id AS relation_id
    FROM relations
        JOIN  osm_quebec
            ON osm_quebec.id = relations.refs
    WHERE
       osm_quebec.kind = 'node'
    ORDER BY relation_id, relation_index
);

CREATE TEMPORARY VIEW road_points_with_distance AS (
    SELECT *,
           st_point(lat, lon) AS point,
           lag(st_point(lat, lon)) OVER (PARTITION BY relation_id ORDER BY relation_index) AS lag_point
    FROM road_points
);


CREATE TEMPORARY VIEW data AS (
    SELECT
        relation_id AS relation_id,
        LIST([road_points_with_distance.lat, road_points_with_distance.lon]) AS points,
        ANY_VALUE(tags) AS tags,
        ANY_VALUE(tags)['highway'][1] AS highway,
        SUM(ST_Distance_Spheroid(point, lag_point)) AS length
    FROM road_points_with_distance
    GROUP BY relation_id
);



INSERT INTO city_road(city_name, extract_date, relation_id, points, tags, highway, length) (
    SELECT '{{ city }}' AS city_name,
           '{{extract_date}}' AS extract_date,
           relation_id,
           points,
           tags,
           highway,
           length
    FROM data
);

''')


def extract_city(file_name: str) -> str:
    return file_name.name[:file_name.name.index(".")].replace("_quebec", "")


def sql_safe(w: str) -> str:
    return str(w).replace("'", r"\\'")


def append_data_to_duckdb(conn: duckdb.DuckDBPyConnection, osm_pbf_path: Path, extract_date: date) -> None:
    city = extract_city(osm_pbf_path)
    sql = sql_tmpl.render(city=sql_safe(city),
                          file_name=sql_safe(osm_pbf_path),
                          extract_date=str(extract_date))
    conn.execute(sql)

def append_data_to_duckdb_file(duckdb_file: Path, osm_pbf_path: Path, extract_date: date) -> None:
    with duckdb.connect(str(duckdb_file)) as conn:
        append_data_to_duckdb(conn, osm_pbf_path, extract_date)

@click.command()
@click.option('--output-file', type=click.Path(exists=False, file_okay=True, dir_okay=False, writable=True, path_type=Path))
@click.option('--input-dir', type=click.Path(exists=False, file_okay=False, dir_okay=True, readable=True, path_type=Path))
@click.option('--date', '--date-end', type=click.DateTime(formats=["%Y-%m-%d"]))
def create_duckdb_file(output_file: Path, input_dir: Path, date: datetime) -> None:

    output_file.parent.mkdir(parents=True, exist_ok=True)
    files = [f for f in input_dir.glob('**/*osm.pbf') if f.is_file()]

    for osm_file in tqdm(files):
        if "'" in str(osm_file):
            continue
        try:
            append_data_to_duckdb_file(output_file, osm_file, date.date())
        except Exception as e:
            click.echo(f"Could not process {osm_file}")



if __name__ == '__main__':
    create_duckdb_file()