use crate::AppState;
use crate::error::AppError;
use crate::tiles::{DataId, TileId};
use anyhow::{anyhow, bail};
use axum::Json;
use axum::extract::{Path, State};
use axum::http::header;
use axum::response::IntoResponse;
use flate2::read::GzDecoder;
use image::codecs::png::PngEncoder;
use image::{Rgba, RgbaImage};
use moka::future::Cache;
use rayon::iter::IntoParallelIterator;
use std::collections::HashMap;
use std::f64::consts::PI;
use std::io;
use std::io::{Cursor, ErrorKind};
use std::sync::Arc;
use tracing::info_span;
use wxbox_grib2::GribMessage;
use wxbox_grib2::wgs84::LatLong;
use wxbox_pal::ColorPalette;

use rayon::iter::ParallelIterator;
use wxbox_common::{Grib2DataSource, GribMessageTimeMetadata, GribTileMetadata};

pub type Grib2DataCache = Cache<DataId, Arc<GribMessage>>;
pub type Grib2TileCache = Cache<TileId, Arc<Vec<u8>>>;

pub type Grib2DataConfig = HashMap<String, Grib2DataSource>;

#[tracing::instrument(level = "info")]
pub async fn grib2_metadata(
    Path(source): Path<String>,
    State(state): State<AppState>,
) -> Result<Json<wxbox_common::GribTileMetadata>, AppError> {
    // is this even a valid data source?
    let data_id = DataId { source };
    let Some(ds) = state.config.data.grib2.get(&data_id.source) else {
        return Err(anyhow!("invalid/unknown grib2 state").into());
    };

    // ok, so we don't have a tile image yet
    // this means we are going to have to kick off a task to put that in the cache
    // lets check if we have the raw data
    let data = if !state.grib2_data_cache.contains_key(&data_id) {
        // we don't, so let's start by starting a task for that
        load_grib2_data(state.grib2_data_cache, data_id, ds.clone()).await?
    } else {
        state.grib2_data_cache.get(&data_id).await.unwrap()
    };

    Ok(Json(GribTileMetadata {
        message_time: GribMessageTimeMetadata {
            reference_time_significance: data.identification.reference_time_significance,
            year: data.identification.year,
            month: data.identification.month,
            day: data.identification.day,
            hour: data.identification.hour,
            minute: data.identification.minute,
            second: data.identification.second,
        },
        source_configuration: ds.clone(),
    }))
}

#[tracing::instrument(level = "info")]
pub async fn grib2_handler(
    Path((source, z, x, y)): Path<(String, usize, usize, String)>,
    State(state): State<AppState>,
) -> Result<impl IntoResponse, AppError> {
    let mut y = y
        .strip_suffix(".png")
        .ok_or(io::Error::new(ErrorKind::InvalidInput, "invalid"))?;
    let mut size = 256;
    if y.ends_with("@2x") {
        size = 512;
        y = y.strip_suffix("@2x").unwrap();
    }
    let y: usize = y.parse()?;
    let tile_id = TileId {
        source,
        z,
        x,
        y,
        size,
    };

    // do we have a pre-prepared tile? if so, return it immediately
    if let Some(tile) = state.grib2_tile_cache.get(&tile_id).await {
        return Ok(([(header::CONTENT_TYPE, "image/png")], tile.as_ref().clone()));
    }

    // is this even a valid data source?
    let data_id = tile_id.data_id();
    let Some(ds) = state.config.data.grib2.get(&data_id.source) else {
        return Err(anyhow!("invalid/unknown grib2 state").into());
    };

    // ok, so we don't have a tile image yet
    // this means we are going to have to kick off a task to put that in the cache
    // lets check if we have the raw data
    let data = if !state.grib2_data_cache.contains_key(&data_id) {
        // we don't, so let's start by starting a task for that
        load_grib2_data(state.grib2_data_cache, data_id, ds.clone()).await?
    } else {
        state.grib2_data_cache.get(&data_id).await.unwrap()
    };

    // we know we need to build the tile, so let's do that now
    // it also returns it, so we can conveniently return it right now
    let pixel_data =
        render_to_png(state.grib2_tile_cache.clone(), data, tile_id, ds.clone()).await?;

    Ok((
        [(header::CONTENT_TYPE, "image/png")],
        pixel_data.as_ref().clone(),
    ))
}

#[tracing::instrument(level = "info")]
async fn load_grib2_data(
    cache: Grib2DataCache,
    data_id: DataId,
    data_source: Grib2DataSource,
) -> anyhow::Result<Arc<GribMessage>> {
    let client = reqwest::Client::new();
    let r = client.get(data_source.from.as_str()).send().await?;

    if !r.status().is_success() {
        bail!("grib2 data failed to load: {}", r.status());
    }

    let data = Cursor::new(r.bytes().await?.to_vec());

    let msg = if data_source.needs_gzip {
        let decoder = GzDecoder::new(data);
        GribMessage::new(decoder)
    } else {
        GribMessage::new(data)
    }?;

    let data = Arc::new(msg);
    cache.insert(data_id, data.clone()).await;

    Ok(data)
}

const TWO_PI: f64 = PI * 2.0;
const HALF_PI: f64 = PI / 2.0;

async fn render_to_png(
    cache: Grib2TileCache,
    data: Arc<GribMessage>,
    tile_id: TileId,
    data_source: Grib2DataSource,
) -> anyhow::Result<Arc<Vec<u8>>> {
    let span = info_span!("render_to_png");
    let span = span.enter();
    let mut image = RgbaImage::new(tile_id.size as u32, tile_id.size as u32);

    let n = 2_usize.pow(tile_id.z as u32) as f64 * tile_id.size as f64;
    let tile_x_times_tilesize = tile_id.x as f64 * tile_id.size as f64;
    let tile_y_times_tilesize = tile_id.y as f64 * tile_id.size as f64;

    let generate_pixels_span = info_span!("generate_pixels");
    let pixels = (0..tile_id.size)
        .into_par_iter()
        .map(|x| {
            (0..tile_id.size)
                .into_par_iter()
                .map(|y| {
                    let x_cartesian = (tile_x_times_tilesize + x as f64) / n;
                    let y_cartesian = (tile_y_times_tilesize + y as f64) / n;

                    let long = (TWO_PI * x_cartesian - PI).to_degrees();
                    let lat =
                        ((PI - TWO_PI * y_cartesian).exp().atan() * 2.0_f64 - HALF_PI).to_degrees();

                    let nearest_value = data.value_for(LatLong { lat, long }).map(|u| u as f64);

                    colorize(nearest_value, &data_source).unwrap_or(Rgba::from([0, 0, 0, 0]))
                })
                .collect::<Vec<_>>()
        })
        .collect::<Vec<_>>();
    drop(generate_pixels_span);

    let put_pixels_span = info_span!("put_pixels");

    pixels.iter().enumerate().for_each(|(x, r)| {
        r.iter().enumerate().for_each(|(y, c)| {
            image.put_pixel(x as u32, y as u32, *c);
        });
    });

    drop(put_pixels_span);

    // encode to png bytes and return
    let mut result = vec![];
    let encoder = PngEncoder::new(&mut result);
    image.write_with_encoder(encoder)?;

    let output = Arc::new(result);

    drop(span);

    cache.insert(tile_id, output.clone()).await;

    Ok(output)
}

fn colorize(value: Option<f64>, data_source: &Grib2DataSource) -> anyhow::Result<Rgba<u8>> {
    Ok(match value {
        Some(c) if Some(c) == data_source.no_coverage => Rgba([0, 0, 0, 30]),
        Some(c) if Some(c) == data_source.missing => Rgba([0, 0, 0, 0]),
        Some(c) if Some(c) == data_source.range_folded => Rgba([119, 0, 125, 255]),
        Some(value) => {
            let color = wxbox_pal::parser::parse(&data_source.palette)?.colorize(value);
            Rgba([
                color.red,
                color.green,
                color.blue,
                color.alpha, // guaranteed to be present afaict
            ])
        }
        None => Rgba([0, 0, 0, 30]),
    })
}