// Copyright (c) 2026 Cyprien Avico avicocyprien@yahoo.com
//
// Licensed under the MIT License <LICENSE-MIT or http://opensource.org/licenses/MIT>.
// This file may not be copied, modified, or distributed except according to those terms.

use core::convert::TryFrom;

use crate::{
    checks::application::opcua::{
        OPCUA_TCP_HEADER_LEN, validate_body_len, validate_message_size, validate_tcp_header_length,
        validate_ua_string_available, validate_ua_string_len,
    },
    errors::application::opcua::OpcuaParseError,
};

#[cfg_attr(doc, aquamarine::aquamarine)]
/// OPC UA TCP Packet
///
/// ```mermaid
/// ---
/// title: OpcuaPacket
/// ---
/// packet-beta
/// 0-23: "Message Type bytes[3]"
/// 24-31: "Chunk Type u8"
/// 32-63: "Message Size u32"
/// 64-127: "Payload variable"
/// ```
#[derive(Debug)]
pub struct OpcuaPacket<'a> {
    pub chunks: Vec<OpcuaChunk<'a>>,
}

#[derive(Debug)]
pub struct OpcuaChunk<'a> {
    pub header: OpcuaTcpHeader,
    pub payload: OpcuaPayload<'a>,
}

#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub struct OpcuaTcpHeader {
    pub message_type: OpcuaMessageType,
    pub chunk_type: OpcuaChunkType,
    pub message_size: u32,
}

#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum OpcuaMessageType {
    Hello,
    Acknowledge,
    ErrorMessage,
    ReverseHello,
    OpenSecureChannel,
    Message,
    CloseSecureChannel,
}

#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum OpcuaChunkType {
    Final,
    Intermediate,
    Abort,
}

#[derive(Debug)]
pub enum OpcuaPayload<'a> {
    Hello(OpcuaHello<'a>),
    Acknowledge(OpcuaAcknowledge),
    Error(OpcuaError<'a>),
    ReverseHello(OpcuaReverseHello<'a>),
    SecureConversation(OpcuaSecureConversation<'a>),
    Partial(&'a [u8]),
}

#[derive(Debug)]
pub struct OpcuaHello<'a> {
    pub protocol_version: u32,
    pub receive_buffer_size: u32,
    pub send_buffer_size: u32,
    pub max_message_size: u32,
    pub max_chunk_count: u32,
    pub endpoint_url: Option<&'a str>,
}

#[derive(Debug)]
pub struct OpcuaAcknowledge {
    pub protocol_version: u32,
    pub receive_buffer_size: u32,
    pub send_buffer_size: u32,
    pub max_message_size: u32,
    pub max_chunk_count: u32,
}

#[derive(Debug)]
pub struct OpcuaError<'a> {
    pub status_code: u32,
    pub reason: Option<&'a str>,
}

#[derive(Debug)]
pub struct OpcuaReverseHello<'a> {
    pub server_uri: Option<&'a str>,
    pub endpoint_url: Option<&'a str>,
}

#[derive(Debug)]
pub struct OpcuaSecureConversation<'a> {
    pub secure_channel_id: u32,
    pub data: &'a [u8],
}

impl TryFrom<[u8; 3]> for OpcuaMessageType {
    type Error = OpcuaParseError;

    fn try_from(value: [u8; 3]) -> Result<Self, Self::Error> {
        match &value {
            b"HEL" => Ok(OpcuaMessageType::Hello),
            b"ACK" => Ok(OpcuaMessageType::Acknowledge),
            b"ERR" => Ok(OpcuaMessageType::ErrorMessage),
            b"RHE" => Ok(OpcuaMessageType::ReverseHello),
            b"OPN" => Ok(OpcuaMessageType::OpenSecureChannel),
            b"MSG" => Ok(OpcuaMessageType::Message),
            b"CLO" => Ok(OpcuaMessageType::CloseSecureChannel),
            _ => Err(OpcuaParseError::UnknownMessageType(value)),
        }
    }
}

impl TryFrom<u8> for OpcuaChunkType {
    type Error = OpcuaParseError;

    fn try_from(value: u8) -> Result<Self, Self::Error> {
        match value {
            b'F' => Ok(OpcuaChunkType::Final),
            b'C' => Ok(OpcuaChunkType::Intermediate),
            b'A' => Ok(OpcuaChunkType::Abort),
            _ => Err(OpcuaParseError::UnknownChunkType(value)),
        }
    }
}

impl TryFrom<&[u8]> for OpcuaTcpHeader {
    type Error = OpcuaParseError;

    fn try_from(bytes: &[u8]) -> Result<Self, Self::Error> {
        validate_tcp_header_length(bytes.len())?;

        let message_type = OpcuaMessageType::try_from([bytes[0], bytes[1], bytes[2]])?;
        let chunk_type = OpcuaChunkType::try_from(bytes[3])?;
        let message_size = u32::from_le_bytes([bytes[4], bytes[5], bytes[6], bytes[7]]);

        validate_message_size(message_size)?;

        Ok(OpcuaTcpHeader {
            message_type,
            chunk_type,
            message_size,
        })
    }
}

impl<'a> TryFrom<&'a [u8]> for OpcuaPacket<'a> {
    type Error = OpcuaParseError;

    fn try_from(bytes: &'a [u8]) -> Result<Self, Self::Error> {
        validate_tcp_header_length(bytes.len())?;

        let mut chunks = Vec::new();
        let mut offset = 0usize;

        while offset < bytes.len() {
            let remaining = &bytes[offset..];
            validate_tcp_header_length(remaining.len())?;

            let header = OpcuaTcpHeader::try_from(remaining)?;
            let message_size = header.message_size as usize;

            let (consumed, payload) = if remaining.len() < message_size {
                let body = &remaining[OPCUA_TCP_HEADER_LEN..];
                (remaining.len(), OpcuaPayload::Partial(body))
            } else {
                let body = &remaining[OPCUA_TCP_HEADER_LEN..message_size];
                (message_size, parse_payload(header.message_type, body)?)
            };

            chunks.push(OpcuaChunk { header, payload });
            offset += consumed;
        }

        Ok(OpcuaPacket { chunks })
    }
}

fn parse_payload<'a>(
    message_type: OpcuaMessageType,
    body: &'a [u8],
) -> Result<OpcuaPayload<'a>, OpcuaParseError> {
    match message_type {
        OpcuaMessageType::Hello => parse_hello(body).map(OpcuaPayload::Hello),
        OpcuaMessageType::Acknowledge => parse_acknowledge(body).map(OpcuaPayload::Acknowledge),
        OpcuaMessageType::ErrorMessage => parse_error(body).map(OpcuaPayload::Error),
        OpcuaMessageType::ReverseHello => parse_reverse_hello(body).map(OpcuaPayload::ReverseHello),
        OpcuaMessageType::OpenSecureChannel
        | OpcuaMessageType::Message
        | OpcuaMessageType::CloseSecureChannel => {
            parse_secure_conversation(body).map(OpcuaPayload::SecureConversation)
        }
    }
}

fn parse_hello(bytes: &[u8]) -> Result<OpcuaHello<'_>, OpcuaParseError> {
    const FIXED_LEN: usize = 20;
    validate_body_len(bytes.len(), FIXED_LEN)?;

    let (endpoint_url, _) = parse_ua_string(&bytes[FIXED_LEN..])?;

    Ok(OpcuaHello {
        protocol_version: read_u32_le(bytes, 0),
        receive_buffer_size: read_u32_le(bytes, 4),
        send_buffer_size: read_u32_le(bytes, 8),
        max_message_size: read_u32_le(bytes, 12),
        max_chunk_count: read_u32_le(bytes, 16),
        endpoint_url,
    })
}

fn parse_acknowledge(bytes: &[u8]) -> Result<OpcuaAcknowledge, OpcuaParseError> {
    const ACK_LEN: usize = 20;
    validate_body_len(bytes.len(), ACK_LEN)?;

    Ok(OpcuaAcknowledge {
        protocol_version: read_u32_le(bytes, 0),
        receive_buffer_size: read_u32_le(bytes, 4),
        send_buffer_size: read_u32_le(bytes, 8),
        max_message_size: read_u32_le(bytes, 12),
        max_chunk_count: read_u32_le(bytes, 16),
    })
}

fn parse_error(bytes: &[u8]) -> Result<OpcuaError<'_>, OpcuaParseError> {
    const STATUS_LEN: usize = 4;
    validate_body_len(bytes.len(), STATUS_LEN)?;

    let (reason, _) = parse_ua_string(&bytes[STATUS_LEN..])?;

    Ok(OpcuaError {
        status_code: read_u32_le(bytes, 0),
        reason,
    })
}

fn parse_reverse_hello(bytes: &[u8]) -> Result<OpcuaReverseHello<'_>, OpcuaParseError> {
    let (server_uri, consumed) = parse_ua_string(bytes)?;
    let (endpoint_url, _) = parse_ua_string(&bytes[consumed..])?;

    Ok(OpcuaReverseHello {
        server_uri,
        endpoint_url,
    })
}

fn parse_secure_conversation(bytes: &[u8]) -> Result<OpcuaSecureConversation<'_>, OpcuaParseError> {
    const SECURE_CHANNEL_ID_LEN: usize = 4;
    validate_body_len(bytes.len(), SECURE_CHANNEL_ID_LEN)?;

    Ok(OpcuaSecureConversation {
        secure_channel_id: read_u32_le(bytes, 0),
        data: &bytes[SECURE_CHANNEL_ID_LEN..],
    })
}

fn parse_ua_string(bytes: &[u8]) -> Result<(Option<&str>, usize), OpcuaParseError> {
    const LEN_FIELD: usize = 4;
    validate_body_len(bytes.len(), LEN_FIELD)?;

    let length = i32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]);
    if length == -1 {
        return Ok((None, LEN_FIELD));
    }
    validate_ua_string_len(length)?;

    let length = length as usize;
    let end = LEN_FIELD + length;
    validate_ua_string_available(bytes.len(), end)?;

    let value =
        core::str::from_utf8(&bytes[LEN_FIELD..end]).map_err(|_| OpcuaParseError::InvalidUtf8)?;

    Ok((Some(value), end))
}

fn read_u32_le(bytes: &[u8], offset: usize) -> u32 {
    u32::from_le_bytes([
        bytes[offset],
        bytes[offset + 1],
        bytes[offset + 2],
        bytes[offset + 3],
    ])
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::convert::hex_stream_to_bytes;

    #[test]
    fn parse_hello_chunk() {
        let mut bytes = b"HELF".to_vec();
        bytes.extend_from_slice(&40u32.to_le_bytes());
        bytes.extend_from_slice(&0u32.to_le_bytes());
        bytes.extend_from_slice(&65536u32.to_le_bytes());
        bytes.extend_from_slice(&65536u32.to_le_bytes());
        bytes.extend_from_slice(&0u32.to_le_bytes());
        bytes.extend_from_slice(&0u32.to_le_bytes());
        bytes.extend_from_slice(&8i32.to_le_bytes());
        bytes.extend_from_slice(b"opc.tcp:");

        let packet = OpcuaPacket::try_from(bytes.as_slice()).unwrap();
        assert_eq!(packet.chunks.len(), 1);
        assert_eq!(
            packet.chunks[0].header.message_type,
            OpcuaMessageType::Hello
        );

        match &packet.chunks[0].payload {
            OpcuaPayload::Hello(hello) => {
                assert_eq!(hello.receive_buffer_size, 65536);
                assert_eq!(hello.endpoint_url, Some("opc.tcp:"));
            }
            _ => panic!("expected hello payload"),
        }
    }

    #[test]
    fn parse_acknowledge_chunk() {
        let mut bytes = b"ACKF".to_vec();
        bytes.extend_from_slice(&28u32.to_le_bytes());
        bytes.extend_from_slice(&0u32.to_le_bytes());
        bytes.extend_from_slice(&65536u32.to_le_bytes());
        bytes.extend_from_slice(&65536u32.to_le_bytes());
        bytes.extend_from_slice(&0u32.to_le_bytes());
        bytes.extend_from_slice(&0u32.to_le_bytes());

        let packet = OpcuaPacket::try_from(bytes.as_slice()).unwrap();
        assert_eq!(packet.chunks.len(), 1);

        match &packet.chunks[0].payload {
            OpcuaPayload::Acknowledge(ack) => {
                assert_eq!(ack.send_buffer_size, 65536);
                assert_eq!(ack.max_chunk_count, 0);
            }
            _ => panic!("expected acknowledge payload"),
        }
    }

    #[test]
    fn reject_unknown_message_type() {
        let mut bytes = b"BADF".to_vec();
        bytes.extend_from_slice(&8u32.to_le_bytes());

        let err = OpcuaPacket::try_from(bytes.as_slice()).unwrap_err();
        assert_eq!(err, OpcuaParseError::UnknownMessageType(*b"BAD"));
    }

    #[test]
    fn parse_truncated_chunk_as_partial() {
        let mut bytes = b"ACKF".to_vec();
        bytes.extend_from_slice(&28u32.to_le_bytes());

        let packet = OpcuaPacket::try_from(bytes.as_slice()).unwrap();
        assert_eq!(packet.chunks.len(), 1);
        assert_eq!(packet.chunks[0].header.message_size, 28);
        match packet.chunks[0].payload {
            OpcuaPayload::Partial(body) => assert!(body.is_empty()),
            _ => panic!("expected partial OPC UA payload"),
        }
    }

    #[test]
    fn parse_opcua_conversation_message_creatsessionresponse() {
        let hex_stream = "[encoded data, 16172 characters]";
        let packet = hex_stream_to_bytes(hex_stream);
        let opcua_packet = OpcuaPacket::try_from(packet.as_slice());
        assert!(
            opcua_packet.is_ok(),
            "failed to parse OPC UA CreateSessionResponse payload: {:?}",
            opcua_packet.unwrap_err()
        );
    }
}
