#!/usr/bin/env python3
import os
import requests
import sqlite3
from pyrosm import OSM
import numpy as np
from shapely.geometry import Point

# -------------------------------
# 1️ Download latest OSM Map
# -------------------------------
def download_osm(url, output_dir="data"):
    os.makedirs(output_dir, exist_ok=True)
    filename = os.path.join(output_dir, url.split("/")[-1])
    if os.path.exists(filename):
        print(f"{filename} already exists. Skipping download.")
        return filename
    print(f"Downloading {url} ...")
    with requests.get(url, stream=True) as r:
        r.raise_for_status()
        with open(filename, "wb") as f:
            for chunk in r.iter_content(chunk_size=8192):
                f.write(chunk)
    print("Download complete!")
    return filename

fl_url = "https://download.geofabrik.de/north-america/us/florida-latest.osm.pbf"
osm_file = download_osm(fl_url)

# ----------------------------------------
# 2️ Prepare SQLite database with R-tree
# ----------------------------------------
os.makedirs("output", exist_ok=True)
conn = sqlite3.connect("output/fl_roads.sqlite")
c = conn.cursor()

# Table for road attributes
c.execute("""
CREATE TABLE IF NOT EXISTS roads (
    id INTEGER PRIMARY KEY,
    name TEXT,
    highway TEXT,
    maxspeed TEXT
);
""")

# R-tree table for spatial queries
c.execute("""
CREATE VIRTUAL TABLE IF NOT EXISTS roads_index USING rtree(
    id, minx, maxx, miny, maxy
);
""")
conn.commit()

# ----------------------------------------
# 3️ Load OSM and iterate through roads
# ----------------------------------------
print("Loading OSM...")
osm = OSM(osm_file)

# Use generator to avoid loading everything into memory
roads_generator = osm.get_data_by_custom_criteria(
    custom_filter={"highway": True},
    filter_type="keep",
    as_generator=True
)

print("Processing roads and writing to database...")
for idx, road in enumerate(roads_generator):
    geom = road.geometry
    if geom is None:
        continue

    minx, miny, maxx, maxy = geom.bounds
    road_id = int(road.id)
    name = road.get("name", "")
    highway = road.get("highway", "")
    maxspeed = road.get("maxspeed", "")

    # Insert into R-tree
    c.execute("INSERT INTO roads_index VALUES (?, ?, ?, ?, ?)",
              (road_id, minx, maxx, miny, maxy))

    # Insert into main table
    c.execute("INSERT INTO roads VALUES (?, ?, ?, ?)",
              (road_id, name, highway, maxspeed))

    # Commit periodically to avoid large transactions
    if idx % 1000 == 0:
        conn.commit()
        print(f"Processed {idx} roads...")

conn.commit()
conn.close()
print("Finished! SQLite database saved at output/fl_roads.sqlite")