use std::{io, path::PathBuf};

use clap::{Parser, ValueEnum};
use color_eyre::{
    eyre::{bail, eyre, Result},
    Report,
};
use hittekaart::{
    gpx::{self, Compression},
    renderer::{self, heatmap, marktile, tilehunt, Renderer},
    storage::{Folder, Sqlite, Storage},
};
use indicatif::{MultiProgress, ProgressBar, ProgressDrawTarget, ProgressStyle};
use is_terminal::IsTerminal;
use rayon::{
    iter::{IndexedParallelIterator, IntoParallelRefIterator, ParallelIterator},
    ThreadPoolBuilder,
};

/// Tile generation mode.
#[derive(ValueEnum, Debug, Clone, Copy, PartialEq, Eq, Hash)]
enum Mode {
    Heatmap,
    Marktile,
    Tilehunter,
}

#[derive(Parser, Debug, Clone)]
#[command(author, version, about)]
struct Args {
    /// The GPX files to parse.
    #[arg(required = true)]
    files: Vec<PathBuf>,

    /// Minimum zoom level to generate tiles for.
    #[arg(long, default_value_t = 0)]
    min_zoom: u32,

    /// Maximum zoom level to generate tiles for.
    #[arg(long, default_value_t = 19)]
    max_zoom: u32,

    /// Number of threads to use. Set to 0 to use all available CPU cores.
    #[arg(long, short, default_value_t = 0)]
    threads: usize,

    /// The output directory. Will be created if it does not exist. Defaults to "tiles" for the
    /// folder-based storage, and "tiles.sqlite" for the SQLite-based storage.
    #[arg(long, short)]
    output: Option<PathBuf>,

    /// Store the tiles in a SQLite database. If given, `--output` will determine the SQLite
    /// filename.
    #[arg(long)]
    sqlite: bool,

    /// Generation mode.
    #[arg(value_enum, long, short, default_value_t = Mode::Heatmap)]
    mode: Mode,

    /// Zoom level for the tilehunter mode.
    #[arg(long, default_value_t = 14)]
    tilehunter_zoom: u32,
}

fn main() -> Result<()> {
    color_eyre::install()?;

    let args = Args::parse();

    if args.max_zoom < args.min_zoom {
        bail!("Max zoom cannot be smaller than min zoom!");
    }

    ThreadPoolBuilder::new()
        .num_threads(args.threads)
        .build_global()?;

    match args.mode {
        Mode::Heatmap => run(heatmap::Renderer, args),
        Mode::Marktile => run(marktile::Renderer, args),
        Mode::Tilehunter => run(tilehunt::Renderer::new(args.tilehunter_zoom), args),
    }
}

fn run<R: Renderer>(renderer: R, args: Args) -> Result<()> {
    let progress_style =
        ProgressStyle::with_template("[{elapsed}] {prefix:.cyan} {wide_bar} {pos:.green}/{len}")?;
    let zoom_style =
        ProgressStyle::with_template("[{elapsed}] {prefix:.yellow} {wide_bar} {pos:.green}/{len}")?;

    let use_progress_bars = io::stdout().is_terminal();
    let make_bar = |len| {
        if use_progress_bars {
            ProgressBar::new(len)
        } else {
            ProgressBar::hidden()
        }
    };

    let bar = make_bar(args.files.len().try_into().unwrap()).with_style(progress_style.clone());
    bar.set_prefix("Reading GPX files");

    let mut tracks = Vec::new();
    args.files
        .par_iter()
        .map(|file| {
            let compression = Compression::suggest_from_path(file)
                .ok_or_else(|| eyre!("Could not determine format for {file:?}"))?;
            let data = gpx::extract_from_file(file, compression)?;
            bar.inc(1);
            Ok::<_, Report>(data)
        })
        .collect_into_vec(&mut tracks);
    let tracks = tracks.into_iter().collect::<Result<Vec<_>>>()?;
    bar.finish();

    let mut storage: Box<dyn Storage + Send> = if args.sqlite {
        let output = args.output.unwrap_or_else(|| "tiles.sqlite".into());
        Box::new(Sqlite::connect(output)?)
    } else {
        let output = args.output.unwrap_or_else(|| "tiles".into());
        Box::new(Folder::new(output))
    };
    storage.prepare()?;

    let multibar = MultiProgress::new();
    if !use_progress_bars {
        multibar.set_draw_target(ProgressDrawTarget::hidden())
    }
    let zoom_bar = make_bar((args.max_zoom - args.min_zoom + 1).into()).with_style(zoom_style);
    multibar.add(zoom_bar.clone());
    zoom_bar.set_prefix("Zoom levels");

    for zoom in args.min_zoom..=args.max_zoom {
        let bar = make_bar(tracks.len().try_into().unwrap()).with_style(progress_style.clone());
        multibar.insert_from_back(1, bar.clone());
        bar.set_prefix("Rendering heat zones");
        let counter = renderer::prepare(&renderer, zoom, &tracks, || {
            bar.inc(1);
            Ok(())
        })?;
        bar.finish();
        multibar.remove(&bar);

        storage.prepare_zoom(zoom)?;

        let bar = make_bar(renderer.tile_count(&counter)?).with_style(progress_style.clone());
        multibar.insert_from_back(1, bar.clone());
        bar.set_prefix("Saving heat tiles");
        renderer::colorize(&renderer, counter, |rendered_tile| {
            storage.store(zoom, rendered_tile.x, rendered_tile.y, &rendered_tile.data)?;
            bar.inc(1);
            Ok(())
        })?;
        bar.finish();
        multibar.remove(&bar);
        zoom_bar.inc(1);
    }
    storage.finish()?;
    zoom_bar.finish();

    Ok(())
}