/*-------------------------------------------------------------------------
 * Copyright (c) Microsoft Corporation.  All rights reserved.
 *
 * documentdb_gateway_core/src/responses/writer.rs
 *
 *-------------------------------------------------------------------------
 */

use crate::{
    context::ConnectionContext,
    error::{DocumentDBError, Result},
    protocol::{header::Header, opcode::OpCode},
    responses::{CommandError, Response},
};
use bson::RawDocument;
use tokio::io::{AsyncWrite, AsyncWriteExt};

/// Write a server response to the client stream
/// # Errors
/// Returns error if the operation fails.
pub async fn write<S>(header: &Header, response: &Response, stream: &mut S) -> Result<()>
where
    S: AsyncWrite + Unpin,
{
    write_and_flush(header, response.as_raw_document()?, stream).await
}

/// Write a raw BSON object to the client stream
/// # Errors
/// Returns error if the operation fails.
#[expect(clippy::cast_possible_truncation, reason = "message size fits in i32")]
#[expect(clippy::cast_possible_wrap, reason = "message size is always positive")]
pub async fn write_and_flush<S>(
    header: &Header,
    response: &RawDocument,
    stream: &mut S,
) -> Result<()>
where
    S: AsyncWrite + Unpin,
{
    // The format of the response will depend on the OP which the client sent
    match header.op_code() {
        OpCode::Command => unimplemented!(),

        // Messages are always responded to with messages
        OpCode::Msg => write_message(header, response, stream).await,

        // Query is responded to with Reply
        #[expect(
            deprecated,
            reason = "OP_QUERY is still supported for legacy clients and testing"
        )]
        OpCode::Query => {
            // Write the header
            let header = Header::new(
                // Total size of the response is the bytes + standard header + reply header
                (response.as_bytes().len() + Header::LENGTH + 20) as i32,
                header.request_id(),
                header.request_id(),
                OpCode::Reply,
            )?;
            header.write_to(stream).await?;

            stream.write_i32_le(0).await?; // Response flags
            stream.write_i64_le(0).await?; // Cursor Id
            stream.write_i32_le(0).await?; // startingFrom
            stream.write_i32_le(1).await?; // numberReturned

            stream.write_all(response.as_bytes()).await?;
            Ok(())
        }

        // Insert has no response
        #[expect(
            deprecated,
            reason = "OP_INSERT is still supported for legacy clients and testing"
        )]
        OpCode::Insert => Ok(()),
        _ => Err(DocumentDBError::internal_error(format!(
            "Unexpected response opcode: {:?}",
            header.op_code()
        ))),
    }?;

    stream.flush().await?;

    Ok(())
}

/// Serializes the Message to bytes and writes them to `writer`.
/// # Errors
/// Returns error if the operation fails.
#[expect(clippy::cast_possible_truncation, reason = "message size fits in i32")]
#[expect(clippy::cast_possible_wrap, reason = "message size is always positive")]
pub async fn write_message<S>(header: &Header, response: &RawDocument, writer: &mut S) -> Result<()>
where
    S: AsyncWrite + Unpin,
{
    let total_length = Header::LENGTH
        + std::mem::size_of::<u32>()
        + std::mem::size_of::<u8>()
        + response.as_bytes().len(); // Message payload + status code

    let header = Header::new(
        total_length as i32,
        header.request_id(),
        header.request_id(),
        OpCode::Msg,
    )?;
    header.write_to(writer).await?;

    // Write Flags
    writer.write_u32_le(0).await?;

    // Write payload type + section
    writer.write_u8(0).await?;

    writer.write_all(response.as_bytes()).await?;

    Ok(())
}

/// # Errors
/// Returns error if the operation fails.
#[expect(clippy::cast_possible_truncation, reason = "message size fits in i32")]
#[expect(clippy::cast_possible_wrap, reason = "message size is always positive")]
pub async fn write_error_without_header<S>(
    connection_context: &ConnectionContext,
    err: DocumentDBError,
    stream: &mut S,
    activity_id: &str,
) -> Result<()>
where
    S: AsyncWrite + Unpin,
{
    let response =
        CommandError::from_error(connection_context, &err, activity_id).to_raw_document_buf();

    let header = Header::new((response.as_bytes().len() + 1) as i32, 0, 0, OpCode::Msg)?;

    write_and_flush(&header, &response, stream).await?;

    Ok(())
}
