//! Python bindings for hittekaart.
//!
//! This module provides simple-to-use bindings for hittekaart generation from Python scripts.
use hittekaart::gpx::{self, Compression, Coordinates};
use hittekaart::renderer::{self, Renderer};
use pyo3::create_exception;
use pyo3::exceptions::PyTypeError;
use pyo3::prelude::*;
use std::error::Error;
use std::ffi::OsStr;
use std::fmt::Write as _;
use std::os::unix::ffi::OsStrExt as _;
use std::path::PathBuf;

create_exception!(hittekaart_py, HitteError, pyo3::exceptions::PyException);

/// Converts an error to a Python error.
///
/// This basically maps everything to [`HitteError`] and provides a stringified error explanation.
///
/// This recursively uses `::source()` to walk the chain.
fn err_to_py(mut error: &dyn Error) -> PyErr {
    let mut text = error.to_string();
    loop {
        match error.source() {
            None => break,
            Some(e) => error = e,
        }
        write!(&mut text, "\ncaused by: {error}").unwrap();
    }
    HitteError::new_err(text)
}

/// Python representation of a track.
#[pyclass]
#[derive(Debug, Clone)]
struct Track {
    inner: Vec<Coordinates>,
}

#[pymethods]
impl Track {
    /// Load a track from a file.
    ///
    /// Compression - if given - should be one of the strings "gzip" or "brotli".
    #[staticmethod]
    fn from_file(path: &[u8], compression: Option<&str>) -> PyResult<Track> {
        let compression = match compression {
            None | Some("") => Compression::None,
            Some("gzip") => Compression::Gzip,
            Some("brotli") => Compression::Brotli,
            Some(x) => return Err(HitteError::new_err(format!("invalid compression: {x}"))),
        };
        let track = gpx::extract_from_file(OsStr::from_bytes(path), compression)
            .map_err(|e| err_to_py(&e))?;
        Ok(Track { inner: track })
    }

    /// Load a track from the given coordinates.
    #[staticmethod]
    fn from_coordinates(coordinates: Vec<(f64, f64)>) -> Track {
        Track {
            inner: coordinates
                .iter()
                .map(|(lon, lat)| Coordinates::new(*lon, *lat))
                .collect(),
        }
    }
}

impl Track {
    fn coordinates(&self) -> Vec<(f64, f64)> {
        self.inner
            .iter()
            .map(|c| (c.longitude, c.latitude))
            .collect()
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
enum StorageType {
    Folder(PathBuf),
    Sqlite(PathBuf),
}

/// Python representation of a storage target.
#[pyclass]
#[derive(Debug, Clone, PartialEq, Eq)]
struct Storage(StorageType);

#[pymethods]
impl Storage {
    /// Output to the given folder.
    #[staticmethod]
    #[pyo3(name = "Folder")]
    fn folder(path: &[u8]) -> Self {
        let path = OsStr::from_bytes(path);
        Storage(StorageType::Folder(path.into()))
    }

    /// Output to the given sqlite file.
    #[staticmethod]
    #[pyo3(name = "Sqlite")]
    fn sqlite(path: &[u8]) -> Self {
        let path = OsStr::from_bytes(path);
        Storage(StorageType::Sqlite(path.into()))
    }
}

impl Storage {
    fn open(&self) -> PyResult<Box<dyn hittekaart::storage::Storage>> {
        match self.0 {
            StorageType::Folder(ref path) => {
                let storage = hittekaart::storage::Folder::new(path.clone());
                Ok(Box::new(storage))
            }
            StorageType::Sqlite(ref path) => {
                let storage = hittekaart::storage::Sqlite::connect(path.clone())
                    .map_err(|e| err_to_py(&e))?;
                Ok(Box::new(storage))
            }
        }
    }
}

/// Python representation of a heatmap renderer.
#[pyclass]
struct HeatmapRenderer {
    inner: renderer::heatmap::Renderer,
}

#[pymethods]
impl HeatmapRenderer {
    #[new]
    fn new() -> HeatmapRenderer {
        HeatmapRenderer {
            inner: renderer::heatmap::Renderer,
        }
    }
}

/// Generate the heatmap.
///
/// * `items` is an iterable of [`Track`]s
/// * `renderer` should be a renderer (like [`HeatmapRenderer`])
/// * `storage` is the [`Storage`] output
#[pyfunction]
fn generate(
    items: &Bound<'_, PyAny>,
    renderer: &Bound<'_, PyAny>,
    storage: &Bound<'_, Storage>,
) -> PyResult<()> {
    let mut tracks = vec![];

    for item in items.try_iter()? {
        let item = item?;
        tracks.push(item.extract::<Track>()?.inner);
    }

    if let Ok(r) = renderer.downcast::<HeatmapRenderer>() {
        do_generate(tracks, &r.borrow().inner, &mut *storage.borrow().open()?)
    } else {
        Err(PyTypeError::new_err("Expected a HeatmapRenderer"))
    }
}

fn do_generate<R: Renderer>(
    tracks: Vec<Vec<Coordinates>>,
    renderer: &R,
    storage: &mut dyn hittekaart::storage::Storage,
) -> PyResult<()> {
    storage.prepare().map_err(|e| err_to_py(&e))?;

    for zoom in 0..=19 {
        let counter =
            renderer::prepare(renderer, zoom, &tracks, || Ok(())).map_err(|e| err_to_py(&e))?;

        storage.prepare_zoom(zoom).map_err(|e| err_to_py(&e))?;

        renderer::colorize(renderer, counter, |rendered_tile| {
            storage.store(zoom, rendered_tile.x, rendered_tile.y, &rendered_tile.data)?;
            Ok(())
        })
        .map_err(|e| err_to_py(&e))?;
    }
    storage.finish().map_err(|e| err_to_py(&e))?;

    Ok(())
}

#[pyfunction]
fn set_threads(threads: usize) -> PyResult<()> {
    rayon::ThreadPoolBuilder::new()
        .num_threads(threads)
        .build_global()
        .map_err(|e| err_to_py(&e))
}

/// A Python module implemented in Rust.
#[pymodule]
fn hittekaart_py(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<Track>()?;
    m.add_class::<HeatmapRenderer>()?;
    m.add_class::<Storage>()?;
    m.add_function(wrap_pyfunction!(generate, m)?)?;
    m.add_function(wrap_pyfunction!(set_threads, m)?)?;
    m.add("HitteError", py.get_type::<HitteError>())?;
    Ok(())
}