//! Second Life terrain patch codec and layer packet construction.

#![allow(clippy::missing_errors_doc)]
#![allow(clippy::must_use_candidate)]
#![allow(clippy::needless_pass_by_value)]
#![allow(clippy::cast_possible_truncation)]
#![allow(clippy::cast_precision_loss)]
#![allow(clippy::cast_possible_wrap)] // Bit widths and bounded patch coordinates are range-checked.
#![allow(clippy::cast_sign_loss)] // Patch coordinates are validated before indexing.

use crate::{BitPack, Error, TerrainPatchLayerType};

const PATCH_SIZE: usize = 16;
const PATCH_VALUES: usize = PATCH_SIZE * PATCH_SIZE;
const STRIDE: i32 = 264;
const OO_SQRT2: f32 = std::f32::consts::FRAC_1_SQRT_2;

#[derive(Clone, Debug, PartialEq)]
pub struct TerrainPatch {
    pub data: Vec<f32>,
    pub x: i32,
    pub y: i32,
}
impl TerrainPatch {
    pub fn new() -> Result<Self, Error> {
        Ok(Self {
            data: Vec::new(),
            x: 0,
            y: 0,
        })
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct TerrainPatchGroupHeader {
    pub stride: i32,
    pub patch_size: i32,
    pub type_: TerrainPatchLayerType,
}

#[derive(Clone, Copy, Debug, PartialEq)]
pub struct TerrainPatchHeader {
    pub dc_offset: f32,
    pub range: i32,
    pub quant_w_bits: i32,
    pub patch_i_ds: i32,
    pub word_bits: u32,
    large_region: bool,
}
impl Default for TerrainPatchHeader {
    fn default() -> Self {
        Self {
            dc_offset: 0.0,
            range: 0,
            quant_w_bits: 0,
            patch_i_ds: 0,
            word_bits: 0,
            large_region: false,
        }
    }
}
impl TerrainPatchHeader {
    pub fn set_patch_i_ds(&mut self, xx: i32, yy: i32) -> Result<i32, Error> {
        let valid = if self.large_region {
            (0..=65_535).contains(&xx) && (0..=65_535).contains(&yy)
        } else {
            (0..=31).contains(&xx) && (0..=31).contains(&yy)
        };
        if !valid {
            return Err(Error::Argument);
        }
        self.patch_i_ds = if self.large_region {
            (xx << 16) | (yy & 0xffff)
        } else {
            (xx << 5) | (yy & 0x1f)
        };
        Ok(self.patch_i_ds)
    }
    pub fn large_region(&self) -> bool {
        self.large_region
    }
    pub fn set_large_region(&mut self, value: bool) {
        self.large_region = value;
    }
    pub fn x(&self) -> i32 {
        self.patch_i_ds >> if self.large_region { 16 } else { 5 }
    }
    pub fn set_x(&mut self, value: i32) {
        let _ = self.set_patch_i_ds(value, self.y());
    }
    pub fn y(&self) -> i32 {
        self.patch_i_ds & if self.large_region { 0xffff } else { 0x1f }
    }
    pub fn set_y(&mut self, value: i32) {
        let _ = self.set_patch_i_ds(self.x(), value);
    }
}

pub struct TerrainCompressor;
impl TerrainCompressor {
    pub const PATCHES_PER_EDGE: i32 = 16;
    pub const END_OF_PATCHES: i32 = 97;

    fn copy_matrix() -> [usize; PATCH_VALUES] {
        let mut result = [0; PATCH_VALUES];
        let (mut diag, mut right, mut i, mut j, mut count) = (false, true, 0usize, 0usize, 0usize);
        while i < PATCH_SIZE && j < PATCH_SIZE {
            result[j * PATCH_SIZE + i] = count;
            count += 1;
            if !diag {
                if right {
                    if i < PATCH_SIZE - 1 {
                        i += 1;
                    } else {
                        j += 1;
                    }
                } else if j < PATCH_SIZE - 1 {
                    j += 1;
                } else {
                    i += 1;
                }
                right = !right;
                diag = true;
            } else if right {
                i += 1;
                j -= 1;
                if i == PATCH_SIZE - 1 || j == 0 {
                    diag = false;
                }
            } else {
                i -= 1;
                j += 1;
                if j == PATCH_SIZE - 1 || i == 0 {
                    diag = false;
                }
            }
        }
        result
    }

    fn cosine(u: usize, n: usize) -> f32 {
        (((2 * n + 1) as f32) * (u as f32) * std::f32::consts::PI * 0.5 / PATCH_SIZE as f32).cos()
    }

    fn prescan(data: &[f32]) -> Result<TerrainPatchHeader, Error> {
        if data.len() != PATCH_VALUES || data.iter().any(|value| !value.is_finite()) {
            return Err(Error::Argument);
        }
        let mut min = f32::INFINITY;
        let mut max = f32::NEG_INFINITY;
        for value in data {
            min = min.min(*value);
            max = max.max(*value);
        }
        Ok(TerrainPatchHeader {
            dc_offset: min,
            range: ((max - min) + 1.0) as i32,
            ..TerrainPatchHeader::default()
        })
    }

    fn compress(data: &[f32], header: &TerrainPatchHeader) -> [i32; PATCH_VALUES] {
        let premult = 1024.0 / header.range.max(1) as f32;
        let sub = 512.0 + header.dc_offset * premult;
        let mut block = [0.0; PATCH_VALUES];
        for (target, value) in block.iter_mut().zip(data) {
            *target = *value * premult - sub;
        }
        let mut rows = [0.0; PATCH_VALUES];
        for line in 0..PATCH_SIZE {
            let base = line * PATCH_SIZE;
            let total: f32 = block[base..base + PATCH_SIZE].iter().sum();
            rows[base] = OO_SQRT2 * total;
            for u in 1..PATCH_SIZE {
                rows[base + u] = (0..PATCH_SIZE)
                    .map(|n| block[base + n] * Self::cosine(u, n))
                    .sum();
            }
        }
        let copy = Self::copy_matrix();
        let mut output = [0; PATCH_VALUES];
        for column in 0..PATCH_SIZE {
            let total: f32 = (0..PATCH_SIZE).map(|n| rows[n * PATCH_SIZE + column]).sum();
            output[copy[column]] = (OO_SQRT2 * total * 0.125 / (1.0 + 2.0 * column as f32)) as i32;
            for u in 1..PATCH_SIZE {
                let total: f32 = (0..PATCH_SIZE)
                    .map(|n| rows[n * PATCH_SIZE + column] * Self::cosine(u, n))
                    .sum();
                output[copy[u * PATCH_SIZE + column]] =
                    (total * 0.125 / (1.0 + 2.0 * (u + column) as f32)) as i32;
            }
        }
        output
    }

    fn encode_header(
        output: &mut BitPack,
        mut header: TerrainPatchHeader,
        patch: &[i32],
    ) -> Result<i32, Error> {
        let base = (header.quant_w_bits & 0x0f) + 2;
        let mut word_bits = base / 2;
        for value in patch.iter().copied().filter(|value| *value != 0) {
            let magnitude = value.unsigned_abs();
            let bits = (u32::BITS - magnitude.leading_zeros()) as i32;
            word_bits = word_bits.max(bits);
        }
        word_bits = (word_bits + 1).clamp(2, 17);
        header.quant_w_bits = (header.quant_w_bits & 0xf0) | (word_bits - 2);
        output.pack_bits_with_int32_int32(header.quant_w_bits, 8)?;
        output.pack_float(header.dc_offset)?;
        output.pack_bits_with_int32_int32(header.range, 16)?;
        output.pack_bits_with_int32_int32(
            header.patch_i_ds,
            if header.large_region { 32 } else { 10 },
        )?;
        Ok(word_bits)
    }

    fn encode_patch_values(
        output: &mut BitPack,
        patch: &[i32],
        word_bits: i32,
    ) -> Result<(), Error> {
        for (index, value) in patch.iter().copied().enumerate() {
            if value == 0 {
                if patch[index..].iter().all(|candidate| *candidate == 0) {
                    output.pack_bits_with_int32_int32(2, 2)?;
                    return Ok(());
                }
                output.pack_bits_with_int32_int32(0, 1)?;
            } else {
                output.pack_bits_with_int32_int32(if value < 0 { 7 } else { 6 }, 3)?;
                let maximum = (1_i32 << word_bits) - 1;
                output.pack_bits_with_int32_int32(
                    i32::try_from(value.unsigned_abs())
                        .unwrap_or(maximum)
                        .min(maximum),
                    word_bits,
                )?;
            }
        }
        Ok(())
    }

    fn encode_one(
        output: &mut BitPack,
        data: &[f32],
        x: i32,
        y: i32,
        large_region: bool,
    ) -> Result<(), Error> {
        let mut header = Self::prescan(data)?;
        header.large_region = large_region;
        header.quant_w_bits = 136;
        header.set_patch_i_ds(x, y)?;
        let patch = Self::compress(data, &header);
        let word_bits = Self::encode_header(output, header, &patch)?;
        Self::encode_patch_values(output, &patch, word_bits)
    }

    pub fn create_layer_data_packet(
        patches: Vec<TerrainPatch>,
        type_: TerrainPatchLayerType,
    ) -> Result<crate::packets::LayerDataPacket, Error> {
        if patches.len() > 4096 {
            return Err(Error::Argument);
        }
        let capacity = patches
            .len()
            .checked_mul(1024)
            .and_then(|value| value.checked_add(32))
            .ok_or(Error::Argument)?;
        let mut bits = BitPack::new(vec![0; capacity], 0)?;
        bits.pack_bits_with_int32_int32(STRIDE, 16)?;
        bits.pack_bits_with_int32_int32(PATCH_SIZE as i32, 8)?;
        bits.pack_bits_with_int32_int32(i32::from(type_.0), 8)?;
        let large = matches!(type_.0, 57 | 58 | 77);
        for patch in patches {
            Self::encode_one(&mut bits, &patch.data, patch.x, patch.y, large)?;
        }
        bits.pack_bits_with_int32_int32(Self::END_OF_PATCHES, 8)?;
        let used = usize::try_from(bits.byte_pos())
            .map_err(|_| Error::Argument)?
            .checked_add(1)
            .ok_or(Error::Argument)?;
        let mut packet = crate::packets::LayerDataPacket::new_with_constructor()?;
        packet.layer_id.type_ = type_.0;
        packet.layer_data.data = bits.data[..used.min(bits.data.len())].to_vec();
        Ok(packet)
    }

    pub fn create_land_packet_with_single_array_int32_int32(
        patch_data: Vec<f32>,
        x: i32,
        y: i32,
    ) -> Result<crate::packets::LayerDataPacket, Error> {
        Self::create_layer_data_packet(
            vec![TerrainPatch {
                data: patch_data,
                x,
                y,
            }],
            TerrainPatchLayerType::LAND,
        )
    }
    pub fn create_land_packet_with_single_array_int32_int32_80090bdf(
        patch_data: Vec<f32>,
        x: i32,
        y: i32,
    ) -> Result<crate::packets::LayerDataPacket, Error> {
        Self::create_land_packet_with_single_array_int32_int32(patch_data, x, y)
    }
    pub fn create_land_packet_with_single_array_int32_array(
        heightmap: Vec<f32>,
        patches: Vec<i32>,
    ) -> Result<crate::packets::LayerDataPacket, Error> {
        let edge = (heightmap.len() as f64).sqrt() as usize;
        if edge * edge != heightmap.len()
            || edge < PATCH_SIZE
            || !edge.is_multiple_of(PATCH_SIZE)
            || patches.len() > 4096
        {
            return Err(Error::Argument);
        }
        let per_edge = edge / PATCH_SIZE;
        let mut values = Vec::with_capacity(patches.len());
        for id in patches {
            let id = usize::try_from(id).map_err(|_| Error::Argument)?;
            let (x, y) = (id % per_edge, id / per_edge);
            if y >= per_edge {
                return Err(Error::Argument);
            }
            let mut data = Vec::with_capacity(PATCH_VALUES);
            for row in 0..PATCH_SIZE {
                let start = (y * PATCH_SIZE + row) * edge + x * PATCH_SIZE;
                data.extend_from_slice(&heightmap[start..start + PATCH_SIZE]);
            }
            values.push(TerrainPatch {
                data,
                x: x as i32,
                y: y as i32,
            });
        }
        let layer = if edge > 256 {
            TerrainPatchLayerType::LAND_EXTENDED
        } else {
            TerrainPatchLayerType::LAND
        };
        Self::create_layer_data_packet(values, layer)
    }

    pub fn create_patch_with_bit_pack_single_array_int32_int32(
        mut output: BitPack,
        patch_data: Vec<f32>,
        x: i32,
        y: i32,
    ) -> Result<(), Error> {
        Self::encode_one(&mut output, &patch_data, x, y, false)
    }
    pub fn create_patch_with_bit_pack_single_array_int32_int32_boolean(
        mut output: BitPack,
        patch_data: Vec<f32>,
        x: i32,
        y: i32,
        large_region: bool,
    ) -> Result<(), Error> {
        Self::encode_one(&mut output, &patch_data, x, y, large_region)
    }
    pub fn create_patch_from_heightmap_with_bit_pack_single_array_int32_int32(
        output: BitPack,
        heightmap: Vec<f32>,
        x: i32,
        y: i32,
    ) -> Result<(), Error> {
        Self::create_patch_from_heightmap_with_bit_pack_single_array_int32_int32_boolean(
            output, heightmap, x, y, false,
        )
    }
    pub fn create_patch_from_heightmap_with_bit_pack_single_array_int32_int32_boolean(
        mut output: BitPack,
        heightmap: Vec<f32>,
        x: i32,
        y: i32,
        large_region: bool,
    ) -> Result<(), Error> {
        if heightmap.len() != 65_536 || !(0..16).contains(&x) || !(0..16).contains(&y) {
            return Err(Error::Argument);
        }
        let mut data = Vec::with_capacity(PATCH_VALUES);
        for row in 0..PATCH_SIZE {
            let start = (y as usize * PATCH_SIZE + row) * 256 + x as usize * PATCH_SIZE;
            data.extend_from_slice(&heightmap[start..start + PATCH_SIZE]);
        }
        Self::encode_one(&mut output, &data, x, y, large_region)
    }

    pub fn decode_patch_header_with_bit_pack(
        bitpack: BitPack,
    ) -> Result<TerrainPatchHeader, Error> {
        Self::decode_patch_header_with_bit_pack_boolean(bitpack, false)
    }
    pub fn decode_patch_header_with_bit_pack_boolean(
        mut bitpack: BitPack,
        large_region: bool,
    ) -> Result<TerrainPatchHeader, Error> {
        Self::decode_header(&mut bitpack, large_region)
    }
    pub(crate) fn decode_header(
        bitpack: &mut BitPack,
        large_region: bool,
    ) -> Result<TerrainPatchHeader, Error> {
        let quant_w_bits = bitpack.unpack_bits(8)?;
        if quant_w_bits == Self::END_OF_PATCHES {
            return Ok(TerrainPatchHeader {
                quant_w_bits,
                large_region,
                ..TerrainPatchHeader::default()
            });
        }
        let dc_offset = bitpack.unpack_float()?;
        let range = bitpack.unpack_bits(16)?;
        let patch_i_ds = bitpack.unpack_bits(if large_region { 32 } else { 10 })?;
        Ok(TerrainPatchHeader {
            dc_offset,
            range,
            quant_w_bits,
            patch_i_ds,
            word_bits: ((quant_w_bits & 0x0f) + 2) as u32,
            large_region,
        })
    }
    pub fn decode_patch(
        mut patches: Vec<i32>,
        mut bitpack: BitPack,
        header: TerrainPatchHeader,
        size: i32,
    ) -> Result<(), Error> {
        Self::decode_values(&mut patches, &mut bitpack, header, size)
    }
    pub(crate) fn decode_values(
        patches: &mut [i32],
        bitpack: &mut BitPack,
        header: TerrainPatchHeader,
        size: i32,
    ) -> Result<(), Error> {
        let count = usize::try_from(size.checked_mul(size).ok_or(Error::Argument)?)
            .map_err(|_| Error::Argument)?;
        if patches.len() < count
            || !(1..=32).contains(&size)
            || !(2..=17).contains(&header.word_bits)
        {
            return Err(Error::Argument);
        }
        for index in 0..count {
            if bitpack.unpack_bits(1)? == 0 {
                patches[index] = 0;
                continue;
            }
            if bitpack.unpack_bits(1)? == 0 {
                patches[index..count].fill(0);
                break;
            }
            let negative = bitpack.unpack_bits(1)? != 0;
            let value = bitpack.unpack_bits(header.word_bits as i32)?;
            patches[index] = if negative { -value } else { value };
        }
        Ok(())
    }
    pub fn decompress_patch(
        patches: Vec<i32>,
        header: TerrainPatchHeader,
        group: TerrainPatchGroupHeader,
    ) -> Result<Vec<f32>, Error> {
        if group.patch_size != 16 || patches.len() < PATCH_VALUES {
            return Err(Error::Argument);
        }
        let copy = Self::copy_matrix();
        let mut block = [0.0; PATCH_VALUES];
        for index in 0..PATCH_VALUES {
            block[index] = patches[copy[index]] as f32
                * (1.0 + 2.0 * ((index / PATCH_SIZE) + (index % PATCH_SIZE)) as f32);
        }
        let mut columns = [0.0; PATCH_VALUES];
        for column in 0..PATCH_SIZE {
            for n in 0..PATCH_SIZE {
                columns[n * PATCH_SIZE + column] = OO_SQRT2 * block[column]
                    + (1..PATCH_SIZE)
                        .map(|u| block[u * PATCH_SIZE + column] * Self::cosine(u, n))
                        .sum::<f32>();
            }
        }
        let mut transformed = [0.0; PATCH_VALUES];
        for line in 0..PATCH_SIZE {
            let base = line * PATCH_SIZE;
            for n in 0..PATCH_SIZE {
                transformed[base + n] = (OO_SQRT2 * columns[base]
                    + (1..PATCH_SIZE)
                        .map(|u| columns[base + u] * Self::cosine(u, n))
                        .sum::<f32>())
                    * 0.125;
            }
        }
        let prequant = (header.quant_w_bits >> 4) + 2;
        if !(1..31).contains(&prequant) {
            return Err(Error::Argument);
        }
        let mult = header.range as f32 / (1_i32 << prequant) as f32;
        let add = mult * (1_i32 << (prequant - 1)) as f32 + header.dc_offset;
        Ok(transformed
            .into_iter()
            .map(|value| value * mult + add)
            .collect())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn terrain_codec_round_trip_stays_within_quantization_error() {
        let source: Vec<f32> = (0..PATCH_VALUES)
            .map(|index| 20.0 + index as f32 * 0.03125)
            .collect();
        let packet = TerrainCompressor::create_land_packet_with_single_array_int32_int32(
            source.clone(),
            3,
            7,
        )
        .unwrap();
        let mut bits = BitPack::new(packet.layer_data.data, 0).unwrap();
        assert_eq!(bits.unpack_bits(16).unwrap(), STRIDE);
        let size = bits.unpack_bits(8).unwrap();
        assert_eq!(
            bits.unpack_bits(8).unwrap(),
            i32::from(TerrainPatchLayerType::LAND.0)
        );
        let header = TerrainCompressor::decode_header(&mut bits, false).unwrap();
        let mut coefficients = vec![0; PATCH_VALUES];
        TerrainCompressor::decode_values(&mut coefficients, &mut bits, header, size).unwrap();
        let decoded = TerrainCompressor::decompress_patch(
            coefficients,
            header,
            TerrainPatchGroupHeader {
                stride: STRIDE,
                patch_size: size,
                type_: TerrainPatchLayerType::LAND,
            },
        )
        .unwrap();
        assert_eq!((header.x(), header.y()), (3, 7));
        assert!(
            source
                .iter()
                .zip(decoded)
                .all(|(expected, actual)| (expected - actual).abs() < 0.2)
        );
    }
    #[test]
    fn terrain_payloads_are_bounded_and_coordinates_validated() {
        assert!(
            TerrainCompressor::create_land_packet_with_single_array_int32_int32(
                vec![0.0; 255],
                0,
                0
            )
            .is_err()
        );
        assert!(
            TerrainCompressor::create_land_packet_with_single_array_int32_int32(
                vec![0.0; 256],
                32,
                0
            )
            .is_err()
        );
    }
}