use async_trait::async_trait;
use minisql::operation::ColumnSelection;
use minisql::response_writer::{CompleteStatus, ResponseWriter};
use minisql::restricted_row::RestrictedRow;
use minisql::schema::{Column, TableSchema};
use proto::message::backend::{
    BackendMessage, ColumnDescription, CommandCompleteData, DataRowData, ErrorResponseData,
    ReadyForQueryData, RowDescriptionData,
};
use proto::message::primitive::pglist::PgList;
use proto::message::primitive::pgoid::PgOid;
use proto::writer::backend::BackendProtoWriter;
use proto::writer::protowriter::ProtoFlush;
use std::io::Error;
use std::time::Duration;

pub struct ServerProtoWrapper<W>(W, Option<Duration>);

impl<W> ServerProtoWrapper<W>
where
    W: BackendProtoWriter + ProtoFlush + Send,
{
    pub fn new(writer: W, throttle: Option<Duration>) -> Self {
        Self(writer, throttle)
    }

    pub async fn write_error_message(&mut self, error_message: &str) -> anyhow::Result<()> {
        self.0
            .write_proto(
                ErrorResponseData {
                    code: b'M',
                    message: format!("{error_message}\0").into(),
                }
                .into(),
            )
            .await?;

        Ok(())
    }

    pub async fn write_ready_for_query(&mut self) -> anyhow::Result<()> {
        self.0
            .write_proto(ReadyForQueryData { status: b'I' }.into())
            .await?;
        Ok(())
    }
}

#[async_trait]
impl<W> ProtoFlush for ServerProtoWrapper<W>
where
    W: ProtoFlush + Send,
{
    async fn flush(&mut self) -> Result<(), Error> {
        self.0.flush().await
    }
}

#[async_trait]
impl<W> ResponseWriter for ServerProtoWrapper<W>
where
    W: BackendProtoWriter + ProtoFlush + Send,
{
    async fn write_table_header(
        &mut self,
        table_schema: &TableSchema,
        columns: &ColumnSelection,
    ) -> anyhow::Result<()> {
        let columns = columns
            .iter()
            .map(|column| column_to_description(table_schema, *column))
            .collect::<anyhow::Result<Vec<ColumnDescription>>>()?;

        self.0
            .write_proto(
                RowDescriptionData {
                    columns: columns.into(),
                }
                .into(),
            )
            .await?;
        Ok(())
    }

    async fn write_table_row(&mut self, row: &RestrictedRow) -> anyhow::Result<()> {
        let values = row
            .iter()
            .map(|(_, value)| value.as_text_bytes().into())
            .collect::<Vec<PgList<u8, i32>>>();

        self.0
            .write_proto(BackendMessage::DataRow(DataRowData {
                columns: values.into(),
            }))
            .await?;

        if let Some(throttle) = self.1 {
            self.0.flush().await?;
            tokio::time::sleep(throttle).await;
        }

        Ok(())
    }

    async fn write_command_complete(&mut self, status: CompleteStatus) -> anyhow::Result<()> {
        self.0
            .write_proto(BackendMessage::CommandComplete(CommandCompleteData {
                tag: status.to_string().into(),
            }))
            .await?;
        Ok(())
    }
}

fn column_to_description(
    schema: &TableSchema,
    column: Column,
) -> anyhow::Result<ColumnDescription> {
    let table_name = schema.table_name();
    let table_oid = PgOid::from_unique_name(table_name);

    let column_type = schema.column_type(column);
    let name = schema.column_name_from_column(column);

    let column_index = column.try_into()?;
    let type_oid = column_type.type_oid();
    let type_size = column_type.type_size();

    Ok(ColumnDescription {
        name: name.to_string().into(),
        table_oid,
        column_index,
        type_oid,
        type_size,
        type_modifier: -1,
        format_code: 0, // text format
    })
}