use std::collections::BTreeMap;
use std::f64::consts::PI;
use std::io::{BufWriter, Cursor, Read};
use std::time::SystemTime;
use eccodes::{CodesHandle, FallibleStreamingIterator, ProductKind};
use flate2::read::GzDecoder;
use ndarray::Zip;
use ordered_float::OrderedFloat;
use png::{BitDepth, ColorType, Encoder};
use tokio::sync::RwLock;
use wxbox_pal::{Color, ColorPalette, Palette};
use crate::grib2::{closest_key, LookupTable2D};
use crate::LutKey;
use crate::pixmap::Pixmap;

pub async fn needs_reload(lct: &RwLock<BTreeMap<LutKey, SystemTime>>, lutkey: &LutKey, valid_for: u64) -> bool {
    let lct_reader = lct.read().await;

    if let Some(t) = lct_reader.get(&lutkey) {
        let dur = SystemTime::now().duration_since(*t).expect("time went backwards").as_secs();
        if dur > valid_for {
            return true;
        }
    } else {
        return true;
    }

    false
}

pub async fn load(url: &str, is_gzipped: bool) -> Vec<u8> {
    let f = reqwest::get(url).await.unwrap();
    let data = f.bytes().await.unwrap().to_vec();

    let out;

    if is_gzipped {
        let b = Cursor::new(data);
        let mut decoder = GzDecoder::new(b);

        let mut data = vec![];
        decoder.read_to_end(&mut data).expect("decoding failure");

        out = data;
    } else {
        out = data;
    }

    out
}

pub fn eccodes_remap(data: Vec<u8>) -> LookupTable2D {
    let product_kind = ProductKind::GRIB;

    let mut handle = CodesHandle::new_from_memory(data.clone(), product_kind).unwrap();
    let msg = handle.next().expect("empty grib2 file").expect("empty grib2 file").try_clone().unwrap();

    let mut map = LookupTable2D::new();

    let arr = msg.to_lons_lats_values().unwrap();

    Zip::from(&arr.latitudes)
        .and(&arr.longitudes)
        .and(&arr.values)
        .for_each(|&lat, &long, &value| {
            let lat = OrderedFloat(lat);

            let mut long = long;
            if long > 180.0 {
                long -= 360.0;
            }

            let long = OrderedFloat(long);

            if let Some(row) = map.get_mut(&lat) {
                row.insert(long, value);
            } else {
                map.insert(lat, BTreeMap::from([
                    (long, value)
                ]));
            }
        });

    map
}

pub async fn reload_if_required(from: &str, needs_gzip: bool, valid_for: u64, lut_key: LutKey, lct_cache: &RwLock<BTreeMap<LutKey, SystemTime>>, lut_cache: &RwLock<BTreeMap<LutKey, LookupTable2D>>) {
    if needs_reload(&lct_cache, &lut_key, valid_for).await {
        let mut lct_writer = lct_cache.write().await;

        let map = eccodes_remap(load(from, needs_gzip).await);

        lut_cache.write().await.insert(lut_key, map);
        lct_writer.insert(lut_key, SystemTime::now());
    }
}

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

pub fn render(xtile: f64, ytile: f64, z: i32, tilesize: usize, pal: Palette, map: &LookupTable2D) -> Vec<u8> {
    let mut image: Pixmap = Pixmap::new();

    let denominator = 2.0_f64.powi(z) * tilesize as f64;
    let xtile_times_tilesize = xtile * tilesize as f64;
    let ytile_times_tilesize = ytile * tilesize as f64;

    for x in 0..tilesize {
        for y in 0..tilesize {
            let tx = (xtile_times_tilesize + x as f64) / denominator;
            let ty = (ytile_times_tilesize + y as f64) / denominator;

            let lon = (TWO_PI * tx - PI).to_degrees();
            let lat = ((PI - TWO_PI * ty).exp().atan() * 2.0_f64 - HALF_PI).to_degrees();


            let closest_lat_in_map = &closest_key(map, lat).unwrap();

            if closest_lat_in_map.1 > 0.01 {
                image.set(y, x, Color { red: 0, green: 0, blue: 0, alpha: 30 });
                continue;
            }

            let row = map.get(&closest_lat_in_map.0).unwrap();
            let closest_long_in_row = closest_key(row, lon).unwrap();

            let value = row.get(&closest_long_in_row.0).unwrap();

            if closest_long_in_row.1 > 0.01 {
                image.set(y, x, Color { red: 0, green: 0, blue: 0, alpha: 30 });
                continue;
            }

            let color = match value {
                -999.0 => Color { red: 0, green: 0, blue: 0, alpha: 30 },
                -99.0 => Color { red: 0, green: 0, blue: 0, alpha: 0 },
                value_at_pos => {
                    pal.colorize(*value_at_pos)
                }
            };
            image.set(y, x, color);
        }
    }

    let mut buf: Vec<u8> = vec![];
    // borrow checker insanity
    {
        let mut cur: Cursor<_> = Cursor::new(&mut buf);
        let w = &mut BufWriter::new(&mut cur);
        let mut encoder = Encoder::new(w, tilesize as u32, tilesize as u32);
        encoder.set_color(ColorType::Rgba);
        encoder.set_depth(BitDepth::Eight);
        encoder.set_source_gamma(png::ScaledFloat::from_scaled(45455));
        encoder.set_source_gamma(png::ScaledFloat::new(1.0 / 2.2));
        let source_chromaticities = png::SourceChromaticities::new(
            (0.31270, 0.32900),
            (0.64000, 0.33000),
            (0.30000, 0.60000),
            (0.15000, 0.06000)
        );
        encoder.set_source_chromaticities(source_chromaticities);
        let mut writer = encoder.write_header().unwrap();
        writer.write_image_data(&image.to_raw()).expect("failed to encode png");
        writer.finish().unwrap();
    }

    buf
}