use deku::prelude::*;
use frontend_protocol::{
    DekuBytes,
    read_pascal_string,
    read_vec_u32,
    write_pascal_string,
    write_vec_u32,
};

#[derive(Debug, Clone, PartialEq, DekuRead, DekuWrite)]
pub struct MessageInChat {
    #[deku(
        reader = "read_pascal_string(deku::reader)",
        writer = "write_pascal_string(deku::writer, &self.role)"
    )]
    pub role: String,

    #[deku(
        reader = "read_pascal_string(deku::reader)",
        writer = "write_pascal_string(deku::writer, &self.content)"
    )]
    pub content: String,
}

#[deku::deku_derive(DekuRead, DekuWrite)]
#[derive(Debug, Clone, PartialEq)]
pub struct ChatCompletionRequest {
    #[deku(
        reader = "read_vec_u32(deku::reader)",
        writer = "write_vec_u32(deku::writer, &self.messages)"
    )]
    pub messages: Vec<MessageInChat>,
}


impl ChatCompletionRequest {
    pub fn new(messages: Vec<MessageInChat>) -> Self {
        Self { messages }
    }
}


#[derive(Debug, Clone, PartialEq, DekuRead, DekuWrite)]
#[deku(id_type = "u8")]
#[repr(u8)]
pub enum RequestPayload {
    #[deku(id = "0")]
    ChatCompletion(ChatCompletionRequest),
    #[deku(id = "1")]
    ChatCompletionCancellation,
}

#[derive(Debug, Clone, PartialEq, DekuRead, DekuWrite)]
pub struct Request {
    #[deku(endian = "little")]
    pub request_id: u64,
    pub payload: RequestPayload,
}

#[derive(Debug, Clone, PartialEq, DekuRead, DekuWrite)]
pub struct ResponseChatCompletion {
    #[deku(
        reader = "read_pascal_string(deku::reader)",
        writer = "write_pascal_string(deku::writer, &self.piece)"
    )]
    pub piece: String,
}

#[derive(Debug, Clone, PartialEq, DekuRead, DekuWrite)]
#[deku(id_type = "u8")]
#[repr(u8)]
pub enum ResponsePayload {
    #[deku(id = "0")]
    ChatCompletion(ResponseChatCompletion),
    #[deku(id = "1")]
    ChatCompletionCancellation,
    #[deku(id = "2")]
    ChatCompletionEnd
}

#[derive(Debug, Clone, PartialEq, DekuRead, DekuWrite)]
pub struct Response {
    #[deku(endian = "little")]
    pub request_id: u64,
    pub payload: ResponsePayload,
}

impl DekuBytes for MessageInChat {}

impl DekuBytes for ChatCompletionRequest {}

impl DekuBytes for RequestPayload {}
impl DekuBytes for Request {}

impl DekuBytes for ResponseChatCompletion {}
impl DekuBytes for ResponsePayload {}
impl DekuBytes for Response {}
