mod db;

use axum::{
    routing::get,
    Router,
    extract::{State, Path},
    Json,
};
use scylla::{Session, SessionBuilder};
use std::sync::Arc;
use tower_http::services::ServeDir;
use tower_http::cors::CorsLayer;
use serde::Serialize;

struct AppState {
    scylla_session: Arc<Session>,
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    // Initialize tracing
    tracing_subscriber::fmt::init();

    println!("Connecting to ScyllaDB...");
    println!("Starting backend with landuse support...");
    let uri = std::env::var("SCYLLA_URI").unwrap_or_else(|_| "127.0.0.1:9042".to_string());
    
    let session = SessionBuilder::new()
        .known_node(uri)
        .build()
        .await?;
    
    // Initialize schema and seed data
    db::initialize_schema(&session).await?;
    db::seed_data(&session).await?;
    
    let session = Arc::new(session);
    println!("Connected to ScyllaDB!");

    let state = Arc::new(AppState {
        scylla_session: session,
    });

    let app = Router::new()
        .route("/health", get(health_check))
        .route("/api/tiles/:z/:x/:y", get(get_tile))
        .route("/api/tiles/:z/:x/:y/ways", get(get_tile_ways))
        .route("/api/tiles/:z/:x/:y/buildings", get(get_tile_buildings))
        .route("/api/tiles/:z/:x/:y/landuse", get(get_tile_landuse))
        .route("/api/tiles/:z/:x/:y/water", get(get_tile_water))
        .route("/api/tiles/:z/:x/:y/railways", get(get_tile_railways))
        .nest_service("/", ServeDir::new("static").append_index_html_on_directories(true))
        .layer(CorsLayer::permissive())
        .with_state(state);

    let listener = tokio::net::TcpListener::bind("0.0.0.0:3000").await?;
    println!("Server listening on {}", listener.local_addr()?);
    axum::serve(listener, app).await?;

    Ok(())
}

async fn health_check() -> &'static str {
    "OK"
}

#[derive(Serialize)]
struct MapNode {
    id: i64,
    lat: f64,
    lon: f64,
    tags: std::collections::HashMap<String, String>,
}

async fn get_tile(
    Path((z, x, y)): Path<(i32, i32, i32)>,
    State(state): State<Arc<AppState>>,
) -> Result<Json<Vec<MapNode>>, (axum::http::StatusCode, String)> {
    let query = "SELECT id, lat, lon, tags FROM map_data.nodes WHERE zoom = ? AND tile_x = ? AND tile_y = ?";
    let rows = state.scylla_session.query(query, (z, x, y))
        .await
        .map_err(|e| (axum::http::StatusCode::INTERNAL_SERVER_ERROR, format!("Database error: {}", e)))?
        .rows
        .unwrap_or_default();
    
    let mut nodes = Vec::new();
    for row in rows {
        let (id, lat, lon, tags) = row.into_typed::<(i64, f64, f64, std::collections::HashMap<String, String>)>()
            .map_err(|e| (axum::http::StatusCode::INTERNAL_SERVER_ERROR, format!("Serialization error: {}", e)))?;
        nodes.push(MapNode { id, lat, lon, tags });
    }

    Ok(Json(nodes))
}

#[derive(Serialize)]
struct MapWay {
    id: i64,
    tags: std::collections::HashMap<String, String>,
    points: Vec<Vec<f64>>, // List of [lat, lon]
}

async fn get_tile_ways(
    Path((z, x, y)): Path<(i32, i32, i32)>,
    State(state): State<Arc<AppState>>,
) -> Result<Json<Vec<MapWay>>, (axum::http::StatusCode, String)> {
    let query = "SELECT id, tags, points FROM map_data.ways WHERE zoom = ? AND tile_x = ? AND tile_y = ?";
    let rows = state.scylla_session.query(query, (z, x, y))
        .await
        .map_err(|e| (axum::http::StatusCode::INTERNAL_SERVER_ERROR, format!("Database error: {}", e)))?
        .rows
        .unwrap_or_default();
    
    let mut ways = Vec::new();
    for row in rows {
        let (id, tags, points_blob) = row.into_typed::<(i64, std::collections::HashMap<String, String>, Vec<u8>)>()
            .map_err(|e| (axum::http::StatusCode::INTERNAL_SERVER_ERROR, format!("Serialization error: {}", e)))?;
        
        // Deserialize points blob
        let mut points = Vec::new();
        for chunk in points_blob.chunks(16) {
            if chunk.len() == 16 {
                let lat = f64::from_be_bytes(chunk[0..8].try_into().unwrap());
                let lon = f64::from_be_bytes(chunk[8..16].try_into().unwrap());
                points.push(vec![lat, lon]);
            }
        }

        ways.push(MapWay { id, tags, points });
    }

    Ok(Json(ways))
}

async fn get_tile_buildings(
    Path((z, x, y)): Path<(i32, i32, i32)>,
    State(state): State<Arc<AppState>>,
) -> Result<Json<Vec<MapWay>>, (axum::http::StatusCode, String)> {
    let query = "SELECT id, tags, points FROM map_data.buildings WHERE zoom = ? AND tile_x = ? AND tile_y = ?";
    let rows = state.scylla_session.query(query, (z, x, y))
        .await
        .map_err(|e| (axum::http::StatusCode::INTERNAL_SERVER_ERROR, format!("Database error: {}", e)))?
        .rows
        .unwrap_or_default();
    
    let mut buildings = Vec::new();
    for row in rows {
        let (id, tags, points_blob) = row.into_typed::<(i64, std::collections::HashMap<String, String>, Vec<u8>)>()
            .map_err(|e| (axum::http::StatusCode::INTERNAL_SERVER_ERROR, format!("Serialization error: {}", e)))?;
        
        // Deserialize points blob
        let mut points = Vec::new();
        for chunk in points_blob.chunks(16) {
            if chunk.len() == 16 {
                let lat = f64::from_be_bytes(chunk[0..8].try_into().unwrap());
                let lon = f64::from_be_bytes(chunk[8..16].try_into().unwrap());
                points.push(vec![lat, lon]);
            }
        }

        buildings.push(MapWay { id, tags, points });
    }

    Ok(Json(buildings))
}

async fn get_tile_landuse(
    Path((z, x, y)): Path<(i32, i32, i32)>,
    State(state): State<Arc<AppState>>,
) -> Result<Json<Vec<MapWay>>, (axum::http::StatusCode, String)> {
    println!("Request: get_tile_landuse({}, {}, {})", z, x, y);
    let query = "SELECT id, tags, points FROM map_data.landuse WHERE zoom = ? AND tile_x = ? AND tile_y = ?";
    println!("Executing query...");
    let result = state.scylla_session.query(query, (z, x, y)).await;
    
    match result {
        Ok(res) => {
            println!("Query successful, processing rows...");
            let rows = res.rows.unwrap_or_default();
            let mut landuse = Vec::new();
            for row in rows {
                let (id, tags, points_blob) = row.into_typed::<(i64, std::collections::HashMap<String, String>, Vec<u8>)>()
                    .map_err(|e| {
                        println!("Serialization error: {}", e);
                        (axum::http::StatusCode::INTERNAL_SERVER_ERROR, format!("Serialization error: {}", e))
                    })?;
                
                let mut points = Vec::new();
                for chunk in points_blob.chunks(16) {
                    if chunk.len() == 16 {
                        let lat = f64::from_be_bytes(chunk[0..8].try_into().unwrap());
                        let lon = f64::from_be_bytes(chunk[8..16].try_into().unwrap());
                        points.push(vec![lat, lon]);
                    }
                }

                landuse.push(MapWay { id, tags, points });
            }
            println!("Returning {} landuse items", landuse.len());
            Ok(Json(landuse))
        },
        Err(e) => {
            println!("Query failed: {}", e);
            Err((axum::http::StatusCode::INTERNAL_SERVER_ERROR, format!("Database error: {}", e)))
        }
    }
}

async fn get_tile_water(
    Path((z, x, y)): Path<(i32, i32, i32)>,
    State(state): State<Arc<AppState>>,
) -> Result<Json<Vec<MapWay>>, (axum::http::StatusCode, String)> {
    let query = "SELECT id, tags, points FROM map_data.water WHERE zoom = ? AND tile_x = ? AND tile_y = ?";
    let rows = state.scylla_session.query(query, (z, x, y))
        .await
        .map_err(|e| (axum::http::StatusCode::INTERNAL_SERVER_ERROR, format!("Database error: {}", e)))?
        .rows
        .unwrap_or_default();
    
    let mut water = Vec::new();
    for row in rows {
        let (id, tags, points_blob) = row.into_typed::<(i64, std::collections::HashMap<String, String>, Vec<u8>)>()
            .map_err(|e| (axum::http::StatusCode::INTERNAL_SERVER_ERROR, format!("Serialization error: {}", e)))?;
        
        let mut points = Vec::new();
        for chunk in points_blob.chunks(16) {
            if chunk.len() == 16 {
                let lat = f64::from_be_bytes(chunk[0..8].try_into().unwrap());
                let lon = f64::from_be_bytes(chunk[8..16].try_into().unwrap());
                points.push(vec![lat, lon]);
            }
        }

        water.push(MapWay { id, tags, points });
    }

    Ok(Json(water))
}

async fn get_tile_railways(
    Path((z, x, y)): Path<(i32, i32, i32)>,
    State(state): State<Arc<AppState>>,
) -> Result<Json<Vec<MapWay>>, (axum::http::StatusCode, String)> {
    let query = "SELECT id, tags, points FROM map_data.railways WHERE zoom = ? AND tile_x = ? AND tile_y = ?";
    let rows = state.scylla_session.query(query, (z, x, y))
        .await
        .map_err(|e| (axum::http::StatusCode::INTERNAL_SERVER_ERROR, format!("Database error: {}", e)))?
        .rows
        .unwrap_or_default();
    
    let mut railways = Vec::new();
    for row in rows {
        let (id, tags, points_blob) = row.into_typed::<(i64, std::collections::HashMap<String, String>, Vec<u8>)>()
            .map_err(|e| (axum::http::StatusCode::INTERNAL_SERVER_ERROR, format!("Serialization error: {}", e)))?;
        
        let mut points = Vec::new();
        for chunk in points_blob.chunks(16) {
            if chunk.len() == 16 {
                let lat = f64::from_be_bytes(chunk[0..8].try_into().unwrap());
                let lon = f64::from_be_bytes(chunk[8..16].try_into().unwrap());
                points.push(vec![lat, lon]);
            }
        }

        railways.push(MapWay { id, tags, points });
    }

    Ok(Json(railways))
}