//! # nexrad2
//! A parser for the NEXRAD II raw datafile format.

use std::collections::HashMap;
use std::error::Error;
use std::fmt::{Display, Formatter};
use std::io;
use std::io::{Cursor, Read, Seek, SeekFrom};
use std::str::Utf8Error;
use log::{info, trace, warn};
use crate::bzip::BzipDecoder;
use crate::message::{FromBody, LEGACY_CTM_HEADER_LEN, Message, MESSAGE_BODY_SIZE, MESSAGE_HEADER_SIZE, MessageHeader};
use crate::message2::Msg02RDAStatusData;
use crate::message31::{DataMoment, ELEVATION_DATA_LENGTH, ElevationData, GenericDataMoment, Message31, Message31Header, MSG31_HEADER_LENGTH, RADIAL_DATA_LENGTH, RadialData, VOLUME_DATA_LENGTH, VolumeData};

pub mod message;
pub mod bzip;
pub mod message2;
pub mod message31;

#[derive(Debug, Clone, PartialEq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct Nexrad2Chunk {
    pub volume_header_record: VolumeHeaderRecord,
    pub chunks: Vec<Vec<Message>>,
    pub meta_rda_status_data: Msg02RDAStatusData,
    pub elevations: HashMap<usize, Vec<Message31>>
}

#[derive(Debug, Clone, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
/// The Volume Header Record as defined by 2620010H (Build 19.0) section 7.3.3 Volume Header Record
pub struct VolumeHeaderRecord {
    /// A character constant of which the last 2 characters identify the version.
    /// Versions defined by 262010H Build 19.0:
    /// - `AR2V0002.`: Super Resolution disabled at the RDA (pre RDA Build 12.0)
    /// - `AR2V0003.`: Super Resolution (pre RDA Build 12.0)
    /// - `AR2V0004.`: Recombined Super Resolution
    /// - `AR2V0005.`: Super Resolution disabled at the RDA (RDA Build 12.0 and later)
    /// - `AR2V0006.`: Super Resolution (RDA Build 12.0 and later)
    /// - `AR2V0007.`: Recombined Super Resolution (RDA Build 12.0 and later)
    pub tape_filename: String,
    /// Increases by 1 for each volume of radar data in the queue to a maximum of 999, whereupon it rolls back to 001
    pub extension_number: String,
    /// Date - NEXRAD-modified Julian - days since 1/1/1970 where 1/1/1970 is 1.
    pub date: u32,
    /// Time - milliseconds past midnight
    pub time: u32,
    /// ICAO of the radar
    pub icao: String
}


#[derive(Debug)]
pub enum NexradParseError {
    InvalidHeaderTapeFilename(Utf8Error),
    HeaderReadError(io::Error),
    InvalidHeaderLength(usize),
    InvalidHeaderExtensionNumber(Utf8Error),
    InvalidICAO(Utf8Error),
    FileTooShort,
    DecompressError(io::Error),
    InvalidMetaChunkLength(usize),
    TcmHeaderReadFailed(io::Error),
    TcmChunkReadFailed(io::Error),
    FailedToReadFile(io::Error),
    LdmReadFailed(io::Error),
    MissingMsg02,
    Msg31HeaderReadFailed(io::Error),
    Msg31HeaderInvalidRadarID(Utf8Error),
    Msg31HeaderPointerReadFailed(io::Error),
    Msg31DatablockSeekFailed(io::Error),
    Msg31MomentReadFailed(io::Error),
    Msg31MomentStringInvalid(Utf8Error)
}
impl Display for NexradParseError {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::InvalidHeaderTapeFilename(e) => write!(f, "invalid header tape filename: {}", e),
            Self::HeaderReadError(e) => write!(f, "error reading header from cursor: {}", e),
            Self::InvalidHeaderLength(read) => write!(f, "invalid header length, expected 24 bytes but only read {}", read),
            Self::InvalidHeaderExtensionNumber(e) => write!(f, "invalid header extension number: {}", e),
            Self::InvalidICAO(e) => write!(f, "invalid station icao: {}", e),
            Self::FileTooShort => write!(f, "file too short (failed to fill whole buffer)"),
            Self::DecompressError(e) => write!(f, "decompression error: {}", e),
            Self::InvalidMetaChunkLength(real) => write!(f, "invalid meta chunk length, spec defined as {NEXRAD2_META_CHUNK_FIXED_LENGTH} but got {real}"),
            Self::TcmHeaderReadFailed(e) => write!(f, "tcm message header read failed: {}", e),
            Self::TcmChunkReadFailed(e) => write!(f, "tcm chunk read failed: {}", e),
            Self::FailedToReadFile(e) => write!(f, "failed to read file: {}", e),
            Self::LdmReadFailed(e) => write!(f, "ldm read failed: {}", e),
            Self::MissingMsg02 => write!(f, "missing RDA status data (message type 02) in metadata header"),
            Self::Msg31HeaderReadFailed(e) => write!(f, "type 31 header read failed: {}", e),
            Self::Msg31HeaderInvalidRadarID(e) => write!(f, "type 31 header invalid: invalid radar ID: {}", e),
            Self::Msg31HeaderPointerReadFailed(e) => write!(f, "type 31 header invalid: pointer read failed: {}", e),
            Self::Msg31DatablockSeekFailed(e) => write!(f, "type 31 datablock invalid: failed to seek: {}", e),
            Self::Msg31MomentReadFailed(e) => write!(f, "type 31 datablock invalid: moment read failed: {}", e),
            Self::Msg31MomentStringInvalid(e) => write!(f, "type 31 moment invalid: invalid string: {}", e)
        }
    }
}
impl Error for NexradParseError {}

pub const NEXRAD2_META_CHUNK_FIXED_LENGTH: usize = 325888;

#[profiling::function]
pub fn parse_nx2_chunk(cursor: &mut (impl Read + Seek)) -> Result<Nexrad2Chunk, NexradParseError> {
    let mut volume_header = [0u8; 24];
    let read = cursor.read(&mut volume_header).map_err(|e| NexradParseError::HeaderReadError(e))?;
    if read != 24 { return Err(NexradParseError::InvalidHeaderLength(read)); }

    let tape_filename_bytes = &volume_header[0..9];
    let tape_filename = std::str::from_utf8(tape_filename_bytes).map_err(|e| NexradParseError::InvalidHeaderTapeFilename(e))?;

    let extension_number_bytes = &volume_header[9..12];
    let extension_number = std::str::from_utf8(extension_number_bytes).map_err(|e| NexradParseError::InvalidHeaderTapeFilename(e))?;

    let days_since_epoch = u32::from_be_bytes(volume_header[12..16].try_into().unwrap());
    let millis_after_midnight = u32::from_be_bytes(volume_header[16..20].try_into().unwrap());

    let icao_bytes = &volume_header[20..24];
    let icao = std::str::from_utf8(icao_bytes).map_err(|e| NexradParseError::InvalidHeaderTapeFilename(e))?;

    let header = VolumeHeaderRecord {
        tape_filename: tape_filename.to_string(),
        extension_number: extension_number.to_string(),
        date: days_since_epoch,
        time: millis_after_midnight,
        icao: icao.to_string(),
    };

    trace!("Loaded - {:#?}", header);

    let mut records: Vec<Vec<Message>> = vec![];
    let mut elevations: HashMap<usize, Vec<Message31>> = HashMap::new();

    loop {
        // LDM records
        let mut ldm_size_bytes = [0u8; 4];
        match cursor.read_exact(&mut ldm_size_bytes) {
            Ok(_) => (),
            Err(e) if matches!(e.kind(), io::ErrorKind::UnexpectedEof) => {
                break;
            },
            Err(e) => return Err(NexradParseError::LdmReadFailed(e))
        }
        let ldm_size = i32::from_be_bytes(ldm_size_bytes).abs() as usize;

        trace!("Reading LDM record - {} bytes compressed", ldm_size);

        if ldm_size == 0 {
            trace!("Missing LDM record, unseeking size");
            cursor.seek(SeekFrom::Current(-4)).map_err(|e| NexradParseError::LdmReadFailed(e))?;
        }

        let mut compressed_buf = vec![0u8; ldm_size];
        cursor.read_exact(&mut compressed_buf).map_err(|e| NexradParseError::LdmReadFailed(e))?;

        let mut bz_decoder = BzipDecoder::new(Cursor::new(compressed_buf));

        let mut decompressed_buf = vec![];

        io::copy(&mut bz_decoder, &mut decompressed_buf).map_err(|e| NexradParseError::DecompressError(e))?;

        trace!("LDM record decompressed to {} bytes", decompressed_buf.len());

        let mut decompressed = Cursor::new(decompressed_buf);

        let mut messages = vec![];

        loop {
            decompressed.seek(SeekFrom::Current(LEGACY_CTM_HEADER_LEN as i64)).map_err(|e| NexradParseError::TcmChunkReadFailed(e))?;

            let mut message_header = [0u8; MESSAGE_HEADER_SIZE];
            match decompressed.read_exact(&mut message_header) {
                Ok(_) => (),
                Err(e) if matches!(e.kind(), io::ErrorKind::UnexpectedEof) => {
                    break;
                },
                Err(e) => return Err(NexradParseError::TcmChunkReadFailed(e))
            }

            let message_header = MessageHeader {
                message_size: u16::from_be_bytes(message_header[0..2].try_into().unwrap()),
                rda_redundant_channel: message_header[2],
                message_type: message_header[3],
                message_sequence_number: u16::from_be_bytes(message_header[4..6].try_into().unwrap()),
                julian_date: u16::from_be_bytes(message_header[6..8].try_into().unwrap()),
                millis_after_midnight: u32::from_be_bytes(message_header[8..12].try_into().unwrap()),
                num_of_message_segments: u16::from_be_bytes(message_header[12..14].try_into().unwrap()),
                message_segment_num: u16::from_be_bytes(message_header[14..].try_into().unwrap()),
            };

            if message_header.message_type == 31 {
                // msg31 fun
                let offset_pos = decompressed.position();

                let mut header = [0u8; MSG31_HEADER_LENGTH];
                decompressed.read_exact(&mut header).map_err(|e| NexradParseError::Msg31HeaderReadFailed(e))?;

                let radar_identifier_bytes = &header[0..4];
                let radar_identifier = std::str::from_utf8(radar_identifier_bytes).map_err(|e| NexradParseError::Msg31HeaderInvalidRadarID(e))?;

                let header_struct = Message31Header {
                    radar_identifier: radar_identifier.to_string(),
                    collection_time: u32::from_be_bytes(header[4..8].try_into().unwrap()),
                    date: u16::from_be_bytes(header[8..10].try_into().unwrap()),
                    azimuth_number: u16::from_be_bytes(header[10..12].try_into().unwrap()),
                    azimuth_angle: f32::from_be_bytes(header[12..16].try_into().unwrap()),
                    compression_indicator: header[16],
                    spare: header[17],
                    radial_length: u16::from_be_bytes(header[18..20].try_into().unwrap()),
                    azimuth_resolution_spacing: header[20],
                    radial_status: header[21],
                    elevation_number: header[22],
                    cut_sector_number: header[23],
                    elevation_angle: f32::from_be_bytes(header[24..28].try_into().unwrap()),
                    radial_spot_blanking_status: header[28],
                    azimuth_indexing_mode: header[29],
                    data_block_count: u16::from_be_bytes(header[30..32].try_into().unwrap()),
                };

                trace!("{:#?}", header_struct);
                trace!("loading {} datablocks", header_struct.data_block_count);

                let mut data_block_pointers = vec![];
                let mut pointer_buf = [0u8; 4];

                for _ in 0..header_struct.data_block_count {
                    decompressed.read_exact(&mut pointer_buf).map_err(|e| NexradParseError::Msg31HeaderPointerReadFailed(e))?;
                    data_block_pointers.push(u32::from_be_bytes(pointer_buf) as u64);
                }

                let mut volume_data: Option<VolumeData> = None;
                let mut elevation_data: Option<ElevationData> = None;
                let mut radial_data: Option<RadialData> = None;
                let mut data_map = HashMap::new();

                for db_pointer in &data_block_pointers {
                    decompressed.seek(SeekFrom::Start(offset_pos + db_pointer)).map_err(|e| NexradParseError::Msg31DatablockSeekFailed(e))?;

                    // get the type and name and use it to determine what to read
                    let mut type_name_buf = [0u8; 4];
                    decompressed.read_exact(&mut type_name_buf).map_err(|e| NexradParseError::Msg31MomentReadFailed(e))?;

                    decompressed.seek(SeekFrom::Current(-4)).map_err(|e| NexradParseError::Msg31DatablockSeekFailed(e))?;

                    let block_name = std::str::from_utf8(&type_name_buf[1..]).map_err(|e| NexradParseError::Msg31MomentStringInvalid(e))?;

                    match block_name {
                        "VOL" => {
                            trace!("reading volume data");
                            let mut volume_buf = [0u8; VOLUME_DATA_LENGTH];
                            decompressed.read_exact(&mut volume_buf).map_err(|e| NexradParseError::Msg31MomentReadFailed(e))?;
                            volume_data = Some(VolumeData {
                                datablock_type: std::str::from_utf8(&volume_buf[0..1]).map_err(|e| NexradParseError::Msg31MomentStringInvalid(e))?.to_string(),
                                data_name: std::str::from_utf8(&volume_buf[1..4]).map_err(|e| NexradParseError::Msg31MomentStringInvalid(e))?.to_string(),
                                lrtup: u16::from_be_bytes(volume_buf[4..6].try_into().unwrap()),
                                version_major: volume_buf[6],
                                version_minor: volume_buf[7],
                                lat: f32::from_be_bytes(volume_buf[8..12].try_into().unwrap()),
                                long: f32::from_be_bytes(volume_buf[12..16].try_into().unwrap()),
                                site_height: u16::from_be_bytes(volume_buf[16..18].try_into().unwrap()),
                                feedhorn_height: u16::from_be_bytes(volume_buf[18..20].try_into().unwrap()),
                                calibration_constant: f32::from_be_bytes(volume_buf[20..24].try_into().unwrap()),
                                shvtx_power_hor: f32::from_be_bytes(volume_buf[24..28].try_into().unwrap()),
                                shvtx_power_vert: f32::from_be_bytes(volume_buf[28..32].try_into().unwrap()),
                                system_diff_ref: f32::from_be_bytes(volume_buf[32..36].try_into().unwrap()),
                                initial_system_diff_phase: f32::from_be_bytes(volume_buf[36..40].try_into().unwrap()),
                                vcp: u16::from_be_bytes(volume_buf[40..42].try_into().unwrap()),
                                processing_status: u16::from_be_bytes(volume_buf[42..44].try_into().unwrap()),
                            });
                            // read in any extra bytes we are currently ignoring
                            if volume_data.as_ref().expect("unreachable").lrtup != VOLUME_DATA_LENGTH as u16 {
                                let mut extra_buf = vec![0u8; (volume_data.as_ref().expect("unreachable").lrtup as usize) - VOLUME_DATA_LENGTH];
                                decompressed.read_exact(&mut extra_buf).map_err(|e| NexradParseError::Msg31MomentReadFailed(e))?;
                            }
                        },
                        "ELV" => {
                            trace!("reading elevation data");
                            let mut elev_buf = [0u8; ELEVATION_DATA_LENGTH];
                            decompressed.read_exact(&mut elev_buf).map_err(|e| NexradParseError::Msg31MomentReadFailed(e))?;
                            elevation_data = Some(ElevationData {
                                datablock_type: std::str::from_utf8(&elev_buf[0..1]).map_err(|e| NexradParseError::Msg31MomentStringInvalid(e))?.to_string(),
                                data_name: std::str::from_utf8(&elev_buf[1..4]).map_err(|e| NexradParseError::Msg31MomentStringInvalid(e))?.to_string(),
                                lrtup: u16::from_be_bytes(elev_buf[4..6].try_into().unwrap()),
                                atmos: elev_buf[6..8].try_into().unwrap(),
                                calib_const: f32::from_be_bytes(elev_buf[8..12].try_into().unwrap()),
                            });
                            if elevation_data.as_ref().expect("unreachable").lrtup != ELEVATION_DATA_LENGTH as u16 {
                                let mut extra_buf = vec![0u8; (elevation_data.as_ref().expect("unreachable").lrtup as usize) - ELEVATION_DATA_LENGTH];
                                decompressed.read_exact(&mut extra_buf).map_err(|e| NexradParseError::Msg31MomentReadFailed(e))?;
                            }
                        },
                        "RAD" => {
                            trace!("reading radial data");
                            let mut rad_buf = [0u8; RADIAL_DATA_LENGTH];
                            decompressed.read_exact(&mut rad_buf).map_err(|e| NexradParseError::Msg31MomentReadFailed(e))?;
                            radial_data = Some(RadialData {
                                datablock_type: std::str::from_utf8(&rad_buf[0..1]).map_err(|e| NexradParseError::Msg31MomentStringInvalid(e))?.to_string(),
                                data_name: std::str::from_utf8(&rad_buf[1..4]).map_err(|e| NexradParseError::Msg31MomentStringInvalid(e))?.to_string(),
                                lrtup: u16::from_be_bytes(rad_buf[4..6].try_into().unwrap()),
                                unambiguous_range: u16::from_be_bytes(rad_buf[6..8].try_into().unwrap()),
                                noise_level_horz: f32::from_be_bytes(rad_buf[8..12].try_into().unwrap()),
                                noise_level_vert: f32::from_be_bytes(rad_buf[12..16].try_into().unwrap()),
                                nyquist_velocity: u16::from_be_bytes(rad_buf[16..18].try_into().unwrap()),
                                radial_flags: u16::from_be_bytes(rad_buf[18..20].try_into().unwrap()),
                                calib_const_horz_channel: f32::from_be_bytes(rad_buf[20..24].try_into().unwrap()),
                                calib_const_vert_channel: f32::from_be_bytes(rad_buf[24..28].try_into().unwrap()),
                            });
                            if radial_data.as_ref().expect("unreachable").lrtup != RADIAL_DATA_LENGTH as u16 {
                                let mut extra_buf = vec![0u8; (radial_data.as_ref().expect("unreachable").lrtup as usize) - RADIAL_DATA_LENGTH];
                                decompressed.read_exact(&mut extra_buf).map_err(|e| NexradParseError::Msg31MomentReadFailed(e))?;
                            }
                        },
                        "REF" | "VEL" | "CFP" | "SW " | "ZDR" | "PHI" | "RHO" => {
                            trace!("reading {}", block_name);
                            let mut gdm_buf = [0u8; RADIAL_DATA_LENGTH];
                            decompressed.read_exact(&mut gdm_buf).map_err(|e| NexradParseError::Msg31MomentReadFailed(e))?;
                            let gdm = GenericDataMoment {
                                datablock_type: std::str::from_utf8(&gdm_buf[0..1]).map_err(|e| NexradParseError::Msg31MomentStringInvalid(e))?.to_string(),
                                data_name: std::str::from_utf8(&gdm_buf[1..4]).map_err(|e| NexradParseError::Msg31MomentStringInvalid(e))?.to_string(),
                                reserved: u32::from_be_bytes(gdm_buf[4..8].try_into().unwrap()),
                                data_moment_gate_count: u16::from_be_bytes(gdm_buf[8..10].try_into().unwrap()),
                                data_moment_range: u16::from_be_bytes(gdm_buf[10..12].try_into().unwrap()),
                                data_moment_range_sample_interval: u16::from_be_bytes(gdm_buf[12..14].try_into().unwrap()),
                                tover: u16::from_be_bytes(gdm_buf[14..16].try_into().unwrap()),
                                snr_threshold: u16::from_be_bytes(gdm_buf[16..18].try_into().unwrap()),
                                control_flags: gdm_buf[18],
                                data_word_size: gdm_buf[19],
                                scale: f32::from_be_bytes(gdm_buf[20..24].try_into().unwrap()),
                                offset: f32::from_be_bytes(gdm_buf[24..28].try_into().unwrap()),
                            };

                            let ldm = ((gdm.data_moment_gate_count as usize) * (gdm.data_word_size as usize))/8;

                            let mut buf = vec![0u8; ldm];

                            decompressed.read_exact(&mut buf).map_err(|e| NexradParseError::Msg31MomentReadFailed(e))?;

                            data_map.insert(gdm.data_name.clone(), DataMoment {
                                gdm,
                                data: buf
                            });
                        },
                        unknown => warn!("unknown block type {}", unknown)
                    }
                }

                let message31 = Message31 {
                    header: header_struct,
                    volume_info: volume_data,
                    elevation_info: elevation_data,
                    radial_info: radial_data,
                    available_data: data_map
                };

                if let Some(elev_mut) = elevations.get_mut(&(message31.header.elevation_number as usize)) {
                    elev_mut.push(message31.clone());
                } else {
                    elevations.insert(message31.header.elevation_number as usize, vec![message31.clone()]);
                }

                let message = Message::Msg31(message31);

                trace!("{:#?}", message);

                messages.push(message);

                continue;
            }

            let body_size = MESSAGE_BODY_SIZE.max(message_header.message_size as usize);

            let mut body_buf = vec![0u8; body_size];

            decompressed.read_exact(&mut body_buf).map_err(|e| NexradParseError::TcmChunkReadFailed(e))?;

            if message_header.message_type == 0 { trace!("extra segment message"); continue; }

            trace!("Message: {:#?}", message_header);

            let message = match message_header.message_type {
                2 => {
                    Message::Msg02(Msg02RDAStatusData::from_body(body_buf[0..68].try_into().unwrap())?)
                },
                unknown => {
                    trace!("ignoring message type {}", unknown);
                    continue;
                }
            };

            messages.push(message);
        }

        trace!("{} messages loaded from chunk", messages.len());

        records.push(messages);
    }

    info!("Extracting meta records");

    let mut msg2: Option<Msg02RDAStatusData> = None;

    let mut messages = 0;
    for chunk in &records {
        messages += chunk.len();

        for message in chunk {
            match message {
                Message::Msg02(m) if msg2.is_none() => { msg2 = Some(m.clone()); },
                _ => ()
            }
        }
    }

    let msg2 = match msg2 {
        Some(m) => m,
        None => return Err(NexradParseError::MissingMsg02)
    };

    info!("File loaded successfully! {} messages loaded in {} chunks", messages, records.len());

    Ok(Nexrad2Chunk {
        volume_header_record: header,
        chunks: records,
        meta_rda_status_data: msg2,
        elevations
    })
}

