pub mod error;
mod nommer;
pub mod wgs84;

use std::fmt::{Debug, Formatter};
use std::io::{Cursor, Read};
use std::time::Instant;
use image::codecs::png::PngDecoder;
use image::{DynamicImage, ImageBuffer, ImageDecoder, ImageFormat, ImageReader, Luma};
use tracing::{debug, warn};
use crate::error::GribError;
use crate::LatLongVectorRelativity::{EasterlyAndNortherly, IncreasingXY};
use crate::nommer::NomReader;
use crate::wgs84::LatLong;

pub const INDICATOR: u32 = u32::from_be_bytes(*b"GRIB");
pub const EDITION: u8 = 2;

#[derive(Debug)]
pub struct GribMessage {
    pub indicator: u32,
    pub reserved: u16,
    pub discipline: u8,
    pub edition: u8,
    pub message_length: u64,
    pub identification: IdentificationSection,
    pub grid_definition: GridDefinitionSection,
    pub product_definition: ProductDefinitionSection,
    pub data_representation: DataRepresentationSection,
    pub bitmap: BitmapSection,
    pub data: DataSection,
}
impl GribMessage {
    pub fn new<R: Read>(reader: R) -> Result<Self, GribError> {
        let mut nommer = NomReader::new(reader);
        let indicator = nommer.read_u32()?;

        if indicator != INDICATOR {
            return Err(GribError::IncorrectIndicator(indicator));
        }

        let reserved = nommer.read_u16()?;
        let discipline = nommer.read_u8()?;
        let edition = nommer.read_u8()?;

        if edition != EDITION {
            return Err(GribError::IncorrectEdition(edition));
        }

        let message_length = nommer.read_u64()?;

        let mut identification: Option<IdentificationSection> = None;
        let mut grid_definition: Option<GridDefinitionSection> = None;
        let mut product_definition: Option<ProductDefinitionSection> = None;
        let mut data_representation: Option<DataRepresentationSection> = None;
        let mut bitmap: Option<BitmapSection> = None;
        let mut data: Option<DataSection> = None;

        loop {
            let length = match nommer.read_u32() {
                Ok(l) => l,
                Err(e) => return match e.kind() {
                    std::io::ErrorKind::UnexpectedEof => break,
                    _ => return Err(e.into())
                }
            };

            if length == u32::from_be_bytes(*b"7777") {
                break;
            }

            let section_number = nommer.read_u8()?;

            match section_number {
                1 => {
                    identification = Some(IdentificationSection::parse(length, &mut nommer)?);
                },
                3 => {
                    grid_definition = Some(GridDefinitionSection::parse(length, &mut nommer)?);
                },
                4 => {
                    product_definition = Some(ProductDefinitionSection::parse(length, &mut nommer)?);
                },
                5 => {
                    data_representation = Some(DataRepresentationSection::parse(length, &mut nommer)?);
                },
                6 => {
                    bitmap = Some(BitmapSection::parse(length, &mut nommer)?);
                },
                7 => {
                    data = Some(DataSection::parse(length, &mut nommer)?)
                },
                _ => {
                    let _ = nommer.read_n((length - 5) as usize);
                    warn!("unimplemented section # {} len {}", section_number, length)
                }
            }
        }

        let identification = identification.ok_or(GribError::MissingIdentification)?;
        let grid_definition = grid_definition.ok_or(GribError::MissingGridDefinition)?;
        let product_definition = product_definition.ok_or(GribError::MissingProductDefinition)?;
        let mut data_representation = data_representation.ok_or(GribError::MissingDataRepresentation)?;
        let bitmap = bitmap.ok_or(GribError::MissingBitmap)?;
        let data = data.ok_or(GribError::MissingData)?;

        data_representation.load_data(data.data.clone())?;

        Ok(Self {
            indicator,
            reserved,
            discipline,
            edition,
            message_length,
            identification,
            grid_definition,
            product_definition,
            data_representation,
            bitmap,
            data
        })
    }

    pub fn value_for(&self, coord: LatLong) -> Option<f32> {
        match self.grid_definition.image_coordinates(coord) {
            Some((i, j)) => self.data_representation.get_image_coordinate(i as u32, j as u32),
            None => None
        }
    }
}

#[derive(Debug)]
pub struct IdentificationSection {
    pub originating_center_id: u16,
    pub originating_subcenter_id: u16,
    pub master_tables_version: u8,
    pub local_tables_version: u8,
    pub reference_time_significance: u8,
    pub year: u16,
    pub month: u8,
    pub day: u8,
    pub hour: u8,
    pub minute: u8,
    pub second: u8,
    pub production_status: u8,
    pub processed_data_type: u8,
    pub reserved: Vec<u8>
}
impl IdentificationSection {
    fn parse<R: Read>(length: u32, nommer: &mut NomReader<R>) -> Result<Self, GribError> {
        let originating_center_id = nommer.read_u16()?;
        let originating_subcenter_id = nommer.read_u16()?;
        let master_tables_version = nommer.read_u8()?;
        let local_tables_version = nommer.read_u8()?;
        let reference_time_significance = nommer.read_u8()?;
        let year = nommer.read_u16()?;
        let month = nommer.read_u8()?;
        let day = nommer.read_u8()?;
        let hour = nommer.read_u8()?;
        let minute = nommer.read_u8()?;
        let second = nommer.read_u8()?;
        let production_status = nommer.read_u8()?;
        let processed_data_type = nommer.read_u8()?;
        let reserved = nommer.read_n((length - 21) as usize)?;
        Ok(Self {
            originating_center_id,
            originating_subcenter_id,
            master_tables_version,
            local_tables_version,
            reference_time_significance,
            year,
            month,
            day,
            hour,
            minute,
            second,
            production_status,
            processed_data_type,
            reserved
        })
    }
}


#[derive(Debug)]
pub struct GridDefinitionSection {
    pub source: u8,
    pub number_of_data_points: u32,
    pub length_of_optional_number_list: u8,
    pub interpretation_of_number_list: u8,
    pub grid_definition_template_number: u16,
    pub grid_definition: GridDefinition
}
impl GridDefinitionSection {
    fn parse<R: Read>(length: u32, nommer: &mut NomReader<R>) -> Result<Self, GribError> {
        let source = nommer.read_u8()?;

        if source != 0 {
            return Err(GribError::NonTable31);
        }

        let number_of_data_points = nommer.read_u32()?;
        let length_of_optional_number_list = nommer.read_u8()?;

        if length_of_optional_number_list != 0 {
            return Err(GribError::ListOfNumbersNotSupported);
        }

        let interpretation_of_number_list = nommer.read_u8()?;
        let grid_definition_template_number = nommer.read_u16()?;


        let grid_definition = match grid_definition_template_number {
            0 => GridDefinition::LatitudeLongitude(LatLongGrid::parse(length, nommer)?),
            _ => return Err(GribError::UnsupportedGrid(grid_definition_template_number)),
        };

        Ok(Self {
            source,
            number_of_data_points,
            length_of_optional_number_list,
            interpretation_of_number_list,
            grid_definition_template_number,
            grid_definition
        })
    }

    fn image_coordinates(&self, at: LatLong) -> Option<(u64, u64)> {
        self.grid_definition.image_coordinates(at)
    }
}

#[derive(Debug)]
pub struct ProductDefinitionSection {
    pub number_of_coordinate_values_after_template: u16,
    pub product_definition_template_number: u16,
    pub product_definition: ProductDefinition,
}
impl ProductDefinitionSection {
    fn parse<R: Read>(length: u32, nommer: &mut NomReader<R>) -> Result<Self, GribError> {
        let number_of_coordinate_values_after_template = nommer.read_u16()?;
        let product_definition_template_number = nommer.read_u16()?;


        if number_of_coordinate_values_after_template != 0 {
            return Err(GribError::ListOfNumbersNotSupported);
        }

        let product_definition = match product_definition_template_number {
            0 => ProductDefinition::HorizontalLayerAtPointInTime(HorizontalLayerInstantaneousProductDef::parse(length, nommer)?),
            _ => return Err(GribError::UnsupportedProductDefTmpl(product_definition_template_number))
        };

        Ok(Self {
            number_of_coordinate_values_after_template,
            product_definition_template_number,
            product_definition
        })
    }
}

#[derive(Debug)]
pub struct DataRepresentationSection {
    pub number_of_data_points: u32,
    pub data_representation_template_number: u16,
    pub data_representation_template: DataRepresentationTemplate
}
impl DataRepresentationSection {
    fn parse<R: Read>(length: u32, nommer: &mut NomReader<R>) -> Result<Self, GribError> {
        let number_of_data_points = nommer.read_u32()?;
        let data_representation_template_number = nommer.read_u16()?;

        let data_representation_template = match data_representation_template_number {
            41 => DataRepresentationTemplate::GridpointPNG(GridpointPNGDataRepresentation::parse(length, nommer)?),
            _ => return Err(GribError::UnsupportedDataRepresentationTemplate(data_representation_template_number))
        };

        Ok(Self {
            number_of_data_points,
            data_representation_template_number,
            data_representation_template
        })
    }

    fn load_data(&mut self, data: Vec<u8>) -> Result<(), GribError> {
        self.data_representation_template.load_data(data)
    }
    fn get_image_coordinate(&self, x: u32, y: u32) -> Option<f32> {
        self.data_representation_template.get_image_coordinate(x, y)
    }
}

#[derive(Debug)]
pub enum DataRepresentationTemplate {
    GridpointPNG(GridpointPNGDataRepresentation)
}
impl DataRepresentationTemplate {
    fn load_data(&mut self, data: Vec<u8>) -> Result<(), GribError> {
        match self {
            Self::GridpointPNG(png) => png.load_data(data)
        }
    }
    fn get_image_coordinate(&self, x: u32, y: u32) -> Option<f32> {
        match self {
            Self::GridpointPNG(png) => png.get_image_coordinate(x, y)
        }
    }
}

#[derive(Debug)]
pub struct GridpointPNGDataRepresentation {
    pub reference_value: f32,
    pub binary_scale_factor: u16,
    pub decimal_scale_factor: u16,
    pub depth: u8,
    pub type_of_values: u8,

    pub image: Option<ImageBuffer<Luma<u16>, Vec<u16>>>
}
impl GridpointPNGDataRepresentation {
    fn parse<R: Read>(_length: u32, nommer: &mut NomReader<R>) -> Result<Self, GribError> {
        let reference_value = nommer.read_f32()?;
        let binary_scale_factor = nommer.read_u16()?;
        let decimal_scale_factor = nommer.read_u16()?;
        let depth = nommer.read_u8()?;
        let type_of_values = nommer.read_u8()?;
        Ok(Self {
            reference_value,
            binary_scale_factor,
            decimal_scale_factor,
            depth,
            type_of_values,
            image: None
        })
    }
    fn load_data(&mut self, data: Vec<u8>) -> Result<(), GribError> {
        let mut image_reader = ImageReader::new(Cursor::new(data));
        image_reader.set_format(ImageFormat::Png);

        let image = image_reader.decode()?;

        self.image = Some(image.to_luma16());

        Ok(())
    }

    fn get_image_coordinate(&self, x: u32, y: u32) -> Option<f32> {
        match &self.image {
            Some(i) => {
                if x >= i.width() || y >= i.height() {
                    None
                } else {
                    let datapoint = i.get_pixel(x, y).0[0] as f32;
                    let diff = datapoint * 2.0_f32.powi(self.binary_scale_factor as i32);
                    let dig_factor = 10_f32.powi(-(self.decimal_scale_factor as i32));
                    let value = (self.reference_value + diff) * dig_factor;
                    Some(value)
                }
            },
            None => None,
        }
    }
}

#[derive(Debug)]
pub struct BitmapSection {
    pub indicator: u8,
    pub bitmap: Vec<u8>
}
impl BitmapSection {
    fn parse<R: Read>(length: u32, nommer: &mut NomReader<R>) -> Result<Self, GribError> {
        let indicator = nommer.read_u8()?;
        if indicator != 255 {
            return Err(GribError::UnsupportedBitmap(indicator));
        }

        let bitmap = nommer.read_n((length - 6) as usize)?;
        Ok(Self {
            indicator,
            bitmap
        })
    }
}

pub struct DataSection {
    pub data: Vec<u8>
}
impl DataSection {
    fn parse<R: Read>(length: u32, nommer: &mut NomReader<R>) -> Result<Self, GribError> {
        let data = nommer.read_n((length - 5) as usize)?;
        Ok(Self {
            data
        })
    }
}
impl Debug for DataSection {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(f, "[data - {} bytes]", self.data.len())
    }
}

#[derive(Debug)]
pub enum ProductDefinition {
    HorizontalLayerAtPointInTime(HorizontalLayerInstantaneousProductDef),
}
#[derive(Debug)]
pub struct HorizontalLayerInstantaneousProductDef {
    pub parameter_category: u8,
    pub parameter_number: u8,
    pub type_of_generating_process: u8,
    pub background_generating_process_id: u8,
    pub type_of_generating_process_identified: u8,
    pub obs_data_cutoff_hours: u16,
    pub obs_data_cutoff_minutes: u8,
    pub indicator_of_unit_in_time_range: u8,
    pub forecast_time: u32,
    pub first_fixed_surface_type: u8,
    pub scale_factor_first_fixed_surface: u8,
    pub scaled_value_first_fixed_surface: u32,
    pub second_fixed_surface_type: u8,
    pub scale_factor_second_fixed_surface: u8,
    pub scaled_value_second_fixed_surface: u32
}
impl HorizontalLayerInstantaneousProductDef {
    fn parse<R: Read>(_length: u32, nommer: &mut NomReader<R>) -> Result<Self, GribError> {
        let parameter_category = nommer.read_u8()?;
        let parameter_number = nommer.read_u8()?;
        let type_of_generating_process = nommer.read_u8()?;
        let background_generating_process_id = nommer.read_u8()?;
        let type_of_generating_process_identified = nommer.read_u8()?;
        let obs_data_cutoff_hours = nommer.read_u16()?;
        let obs_data_cutoff_minutes = nommer.read_u8()?;
        let indicator_of_unit_in_time_range = nommer.read_u8()?;
        let forecast_time = nommer.read_u32()?;
        let first_fixed_surface_type = nommer.read_u8()?;
        let scale_factor_first_fixed_surface = nommer.read_u8()?;
        let scaled_value_first_fixed_surface = nommer.read_u32()?;
        let second_fixed_surface_type = nommer.read_u8()?;
        let scale_factor_second_fixed_surface = nommer.read_u8()?;
        let scaled_value_second_fixed_surface = nommer.read_u32()?;
        Ok(Self {
            parameter_category,
            parameter_number,
            type_of_generating_process,
            background_generating_process_id,
            type_of_generating_process_identified,
            obs_data_cutoff_hours,
            obs_data_cutoff_minutes,
            indicator_of_unit_in_time_range,
            forecast_time,
            first_fixed_surface_type,
            scale_factor_first_fixed_surface,
            scaled_value_first_fixed_surface,
            second_fixed_surface_type,
            scale_factor_second_fixed_surface,
            scaled_value_second_fixed_surface,
        })
    }
}

#[derive(Debug)]
pub enum GridDefinition {
    LatitudeLongitude(LatLongGrid)
}
impl GridDefinition {
    fn image_coordinates(&self, at: LatLong) -> Option<(u64, u64)> {
        match self {
            GridDefinition::LatitudeLongitude(grid) => grid.image_coordinates(at)
        }
    }
}

#[derive(Debug)]
pub struct LatLongGrid {
    pub shape_of_the_earth: u8,
    pub scale_factor_of_radius_of_spherical_earth: u8,
    pub scale_value_of_radius_of_spherical_earth: u32,
    pub scale_factor_of_major_axis_of_oblate_spheroid_earth: u8,
    pub scaled_value_of_major_axis_of_oblate_spheroid_earth: u32,
    pub scale_factor_of_minor_axis_of_oblate_spheroid_earth: u8,
    pub scaled_value_of_minor_axis_of_oblate_spheroid_earth: u32,
    pub ni: u32, // number of points along a parallel
    pub nj: u32, // number of points along a meridian
    pub basic_angle_of_the_initial_production_domain: u32,
    pub subdivisions_of_basic_angle: u32,
    pub la1: f64,
    pub lo1: f64,
    pub i_direction_increments_given: bool,
    pub j_direction_increments_given: bool,
    pub vector_relativity: LatLongVectorRelativity,
    pub la2: f64,
    pub lo2: f64,
    pub di: f64,
    pub dj: f64,
    pub scanning_mode_flags: u8,
}
#[derive(Debug)]
pub enum LatLongVectorRelativity {
    EasterlyAndNortherly,
    IncreasingXY
}
impl LatLongGrid {
    fn parse<R: Read>(length: u32, nommer: &mut NomReader<R>) -> Result<Self, GribError> {
        if length != 72 {
            return Err(GribError::ListOfNumbersNotSupported);
        }
        let shape_of_the_earth= nommer.read_u8()?;
        let scale_factor_of_radius_of_spherical_earth= nommer.read_u8()?;
        let scale_value_of_radius_of_spherical_earth= nommer.read_u32()?;
        let scale_factor_of_major_axis_of_oblate_spheroid_earth= nommer.read_u8()?;
        let scaled_value_of_major_axis_of_oblate_spheroid_earth= nommer.read_u32()?;
        let scale_factor_of_minor_axis_of_oblate_spheroid_earth= nommer.read_u8()?;
        let scaled_value_of_minor_axis_of_oblate_spheroid_earth= nommer.read_u32()?;
        let ni= nommer.read_u32()?; // number of points along a parallel
        let nj= nommer.read_u32()?; // number of points along a meridian
        let basic_angle_of_the_initial_production_domain= nommer.read_u32()?;
        let subdivisions_of_basic_angle= nommer.read_u32()?;
        let la1= nommer.read_u32()? as f64 * 10.0_f64.powi(-6);
        let lo1= convert_longitude(nommer.read_u32()? as f64 * 10.0_f64.powi(-6));
        let resolution_and_component_flags= nommer.read_u8()?;

        let i_direction_increments_given = resolution_and_component_flags & (1 << 5) != 0;
        let j_direction_increments_given = resolution_and_component_flags & (1 << 4) != 0;
        let vector_relativity = if resolution_and_component_flags & (1 << 3) != 0 { IncreasingXY } else { EasterlyAndNortherly };

        let la2= nommer.read_u32()? as f64 * 10.0_f64.powi(-6);
        let lo2= convert_longitude(nommer.read_u32()? as f64 * 10.0_f64.powi(-6));
        let di= nommer.read_u32()? as f64 * 10.0_f64.powi(-6);
        let dj= nommer.read_u32()? as f64 * 10.0_f64.powi(-6);
        let scanning_mode_flags= nommer.read_u8()?;

        Ok(Self {
            shape_of_the_earth,
            scale_factor_of_radius_of_spherical_earth,
            scale_value_of_radius_of_spherical_earth,
            scale_factor_of_major_axis_of_oblate_spheroid_earth,
            scaled_value_of_major_axis_of_oblate_spheroid_earth,
            scale_factor_of_minor_axis_of_oblate_spheroid_earth,
            scaled_value_of_minor_axis_of_oblate_spheroid_earth,
            ni,
            nj,
            basic_angle_of_the_initial_production_domain,
            subdivisions_of_basic_angle,
            la1,
            lo1,
            i_direction_increments_given,
            j_direction_increments_given,
            vector_relativity,
            la2,
            lo2,
            di,
            dj,
            scanning_mode_flags,
        })
    }

    fn image_coordinates(&self, at: LatLong) -> Option<(u64, u64)> {
        if at.lat > self.la1 || at.lat < self.la2 { return None; }
        if at.long < self.lo1 || at.long > self.lo2 { return None; }

        let lat = ((at.lat - self.la1) / self.di ).round() * self.di + self.la1;
        let long = ((at.long - self.lo1) / self.dj).round() * self.dj + self.lo1;

        if (lat - at.lat).abs() > self.di || (long - at.long).abs() > self.dj {
            return None;
        }

        let diff_lat = at.lat - self.la1;
        let num_steps_lat = diff_lat / -self.di;

        let diff_long = at.long - self.lo1;
        let num_steps_long = diff_long / self.dj;

        Some((num_steps_long as u64, num_steps_lat as u64))
    }
}




fn convert_longitude(longitude: f64) -> f64 {
    if longitude > 180.0 {
        longitude - 360.0
    } else {
        longitude
    }
}