// Copyright 2023 RisingWave Labs
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::convert::Infallible;

use auto_enums::auto_enum;

use super::StreamChunk;
use crate::array::Op;
use crate::row::Row;
use crate::types::DataType;

/// Type of a row change, without row data.
#[derive(Debug, Copy, Clone)]
pub enum RecordType {
    Insert,
    Delete,
    Update,
}

impl RecordType {
    /// Get the corresponding `Op`s for this record type.
    pub fn ops(self) -> &'static [Op] {
        match self {
            RecordType::Insert => &[Op::Insert],
            RecordType::Delete => &[Op::Delete],
            RecordType::Update => &[Op::UpdateDelete, Op::UpdateInsert],
        }
    }
}

/// Generic type to represent a row change.
#[derive(Debug, Clone, Copy)]
pub enum Record<R> {
    Insert { new_row: R },
    Delete { old_row: R },
    Update { old_row: R, new_row: R },
}

impl<R> Record<R> {
    /// Convert this stream record to one or two rows with corresponding ops.
    #[auto_enum(Iterator)]
    pub fn into_rows(self) -> impl Iterator<Item = (Op, R)> {
        match self {
            Record::Insert { new_row } => std::iter::once((Op::Insert, new_row)),
            Record::Delete { old_row } => std::iter::once((Op::Delete, old_row)),
            Record::Update { old_row, new_row } => {
                [(Op::UpdateDelete, old_row), (Op::UpdateInsert, new_row)].into_iter()
            }
        }
    }

    /// Get record type of this record.
    pub fn to_record_type(&self) -> RecordType {
        match self {
            Record::Insert { .. } => RecordType::Insert,
            Record::Delete { .. } => RecordType::Delete,
            Record::Update { .. } => RecordType::Update,
        }
    }

    /// Convert from `&Record<R>` to `Record<&R>`.
    pub fn as_ref(&self) -> Record<&R> {
        match self {
            Record::Insert { new_row } => Record::Insert { new_row },
            Record::Delete { old_row } => Record::Delete { old_row },
            Record::Update { old_row, new_row } => Record::Update { old_row, new_row },
        }
    }

    /// Convert this record to upsert format, by rewriting `Update` to `Insert` and only keeping
    /// the new row.
    pub fn into_upsert(self) -> Self {
        match self {
            Record::Insert { new_row } => Record::Insert { new_row },
            Record::Delete { old_row } => Record::Delete { old_row },
            Record::Update { new_row, .. } => Record::Insert { new_row },
        }
    }

    /// Try mapping the row in the record to another row, returning error if any of the mapping fails.
    pub fn try_map<R2, E>(self, f: impl Fn(R) -> Result<R2, E>) -> Result<Record<R2>, E> {
        Ok(match self {
            Record::Insert { new_row } => Record::Insert {
                new_row: f(new_row)?,
            },
            Record::Delete { old_row } => Record::Delete {
                old_row: f(old_row)?,
            },
            Record::Update { old_row, new_row } => Record::Update {
                old_row: f(old_row)?,
                new_row: f(new_row)?,
            },
        })
    }

    /// Map the row in the record to another row.
    pub fn map<R2>(self, f: impl Fn(R) -> R2) -> Record<R2> {
        let Ok(record) = self.try_map::<R2, Infallible>(|row| Ok(f(row)));
        record
    }
}

impl<R: Row> Record<R> {
    /// Convert this stream record to a stream chunk containing only 1 or 2 rows.
    pub fn to_stream_chunk(&self, data_types: &[DataType]) -> StreamChunk {
        match self {
            Record::Insert { new_row } => {
                StreamChunk::from_rows(&[(Op::Insert, new_row)], data_types)
            }
            Record::Delete { old_row } => {
                StreamChunk::from_rows(&[(Op::Delete, old_row)], data_types)
            }
            Record::Update { old_row, new_row } => StreamChunk::from_rows(
                &[(Op::UpdateDelete, old_row), (Op::UpdateInsert, new_row)],
                data_types,
            ),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::row::OwnedRow;
    use crate::test_prelude::StreamChunkTestExt;

    #[test]
    fn test_into_rows() {
        let record = Record::Insert {
            new_row: OwnedRow::new(vec![Some(1.into())]),
        };
        let rows: Vec<_> = record.into_rows().collect();
        assert_eq!(rows.len(), 1);
        assert_eq!(rows[0].0, Op::Insert);
        assert_eq!(rows[0].1, OwnedRow::new(vec![Some(1.into())]));

        let record = Record::Delete {
            old_row: OwnedRow::new(vec![Some(1.into())]),
        };
        let rows: Vec<_> = record.into_rows().collect();
        assert_eq!(rows.len(), 1);
        assert_eq!(rows[0].0, Op::Delete);
        assert_eq!(rows[0].1, OwnedRow::new(vec![Some(1.into())]));

        let record = Record::Update {
            old_row: OwnedRow::new(vec![Some(1.into())]),
            new_row: OwnedRow::new(vec![Some(2.into())]),
        };
        let rows: Vec<_> = record.into_rows().collect();
        assert_eq!(rows.len(), 2);
        assert_eq!(rows[0].0, Op::UpdateDelete);
        assert_eq!(rows[0].1, OwnedRow::new(vec![Some(1.into())]));
        assert_eq!(rows[1].0, Op::UpdateInsert);
        assert_eq!(rows[1].1, OwnedRow::new(vec![Some(2.into())]));
    }

    #[test]
    fn test_to_stream_chunk() {
        let record = Record::Insert {
            new_row: OwnedRow::new(vec![Some(1i64.into())]),
        };
        let chunk = record.to_stream_chunk(&[DataType::Int64]);
        assert_eq!(
            chunk,
            StreamChunk::from_pretty(
                " I
                + 1"
            )
        );

        let record = Record::Delete {
            old_row: OwnedRow::new(vec![Some(1i64.into())]),
        };
        let chunk = record.to_stream_chunk(&[DataType::Int64]);
        assert_eq!(
            chunk,
            StreamChunk::from_pretty(
                " I
                - 1"
            )
        );

        let record = Record::Update {
            old_row: OwnedRow::new(vec![Some(1i64.into())]),
            new_row: OwnedRow::new(vec![Some(2i64.into())]),
        };
        let chunk = record.to_stream_chunk(&[DataType::Int64]);
        assert_eq!(
            chunk,
            StreamChunk::from_pretty(
                "  I
                U- 1
                U+ 2"
            )
        );
    }
}
