hamelin_lib 0.15.4

use crate::sql::expression::identifier::SimpleIdentifier;
use crate::sql::expression::{OrderByExpression, SQLExpression};
use crate::write_utils::{maybe_newline, maybe_pad, newline_or_space, Indent};
use derive_more::From;
use ordermap::OrderMap;
use std::fmt;
use std::fmt::{Display, Formatter};

#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct Windows {
    pub windows: OrderMap<SimpleIdentifier, WindowSpecification>,
}

impl Windows {
    pub fn with(mut self, identifier: SimpleIdentifier, window: WindowSpecification) -> Self {
        self.windows.insert(identifier, window);
        self
    }

    pub fn fmt_indented(&self, f: &mut Formatter<'_>, indentation: Indent) -> fmt::Result {
        if self.windows.is_empty() {
            Ok(())
        } else {
            newline_or_space(f, indentation)?;
            maybe_pad(f, indentation)?;
            write!(f, "WINDOW")?;
            for (i, (si, window_spec)) in self.windows.iter().enumerate() {
                if i != 0 {
                    write!(f, ",")?;
                }
                write!(f, " {} AS ", si.to_string())?;
                window_spec.fmt_indented(f, indentation.nested())?;
            }
            Ok(())
        }
    }
}

impl Display for Windows {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        self.fmt_indented(f, Indent::default())
    }
}

#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub struct WindowSpecification {
    pub partition_by: Vec<SQLExpression>,
    pub order_by: Vec<OrderByExpression>,
    pub frame: Option<WindowFrame>,
}

impl WindowSpecification {
    pub fn fmt_indented(&self, f: &mut Formatter<'_>, indentation: Indent) -> fmt::Result {
        write!(f, "(")?;
        let spec_indentation = indentation.nested();
        if !self.partition_by.is_empty() {
            newline_or_space(f, indentation)?;
            maybe_pad(f, spec_indentation)?;
            write!(f, "PARTITION BY ")?;
            let partition_count = self.partition_by.len();
            for (i, expr) in self.partition_by.iter().enumerate() {
                write!(f, "{}", expr)?;
                if i != partition_count - 1 {
                    write!(f, ", ")?;
                }
            }
        };

        if !self.order_by.is_empty() {
            newline_or_space(f, indentation)?;
            maybe_pad(f, spec_indentation)?;
            write!(f, "ORDER BY ")?;
            let order_count = self.order_by.len();
            for (i, expr) in self.order_by.iter().enumerate() {
                write!(f, "{}", expr)?;
                if i != order_count - 1 {
                    write!(f, ", ")?;
                }
            }
        };
        if let Some(frame) = &self.frame {
            newline_or_space(f, indentation)?;
            frame.fmt_indented(f, spec_indentation)?;
        }
        maybe_newline(f, indentation)?;
        maybe_pad(f, indentation)?;
        write!(f, ")")
    }

    fn subelements(&self) -> usize {
        self.partition_by
            .iter()
            .map(|e| e.subelements() + 1)
            .sum::<usize>()
            + self
                .order_by
                .iter()
                .map(|e| e.expression.subelements() + 1)
                .sum::<usize>()
            + 1 // For the frame
    }
}

impl Display for WindowSpecification {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        self.fmt_indented(f, Indent::default())
    }
}

#[derive(Debug, Clone)]
pub enum Direction {
    Asc,
    Desc,
}

impl Direction {
    pub fn fmt_indented(&self, f: &mut Formatter<'_>, indentation: Indent) -> fmt::Result {
        maybe_pad(f, indentation)?;
        match self {
            Direction::Asc => write!(f, "ASC"),
            Direction::Desc => write!(f, "DESC"),
        }
    }
}
impl Display for Direction {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        self.fmt_indented(f, Indent::default())
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct WindowFrame {
    pub frame_type: FrameType,
    /// Start bound in `RANGE` / `ROWS BETWEEN <start> AND <end>`.
    pub start: FrameEndpoint,
    /// End bound.
    pub end: FrameEndpoint,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum FrameType {
    RANGE,
    ROWS,
}

impl FrameType {
    pub fn fmt_indented(&self, f: &mut Formatter<'_>, indentation: Indent) -> fmt::Result {
        maybe_pad(f, indentation)?;
        match self {
            FrameType::RANGE => write!(f, "RANGE"),
            FrameType::ROWS => write!(f, "ROWS"),
        }
    }
}

impl Display for FrameType {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        self.fmt_indented(f, Indent::default())
    }
}

/// One side of a SQL window frame (`ROWS` / `RANGE BETWEEN … AND …`).
///
/// SQL allows both bounds to use `PRECEDING` (e.g. `RANGE BETWEEN 50 MINUTE PRECEDING AND 5 MINUTE PRECEDING`).
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum FrameEndpoint {
    UnboundedPreceding,
    UnboundedFollowing,
    CurrentRow,
    Preceding(Box<SQLExpression>),
    Following(Box<SQLExpression>),
}

impl FrameEndpoint {
    pub fn fmt_indented(&self, f: &mut Formatter<'_>, indentation: Indent) -> fmt::Result {
        maybe_pad(f, indentation)?;
        match self {
            FrameEndpoint::UnboundedPreceding => write!(f, "UNBOUNDED PRECEDING"),
            FrameEndpoint::UnboundedFollowing => write!(f, "UNBOUNDED FOLLOWING"),
            FrameEndpoint::CurrentRow => write!(f, "CURRENT ROW"),
            FrameEndpoint::Preceding(expr) => write!(f, "{} PRECEDING", expr),
            FrameEndpoint::Following(expr) => write!(f, "{} FOLLOWING", expr),
        }
    }
}

impl Display for FrameEndpoint {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        self.fmt_indented(f, Indent::default())
    }
}

impl WindowFrame {
    pub fn fmt_indented(&self, f: &mut Formatter<'_>, indentation: Indent) -> fmt::Result {
        self.frame_type.fmt_indented(f, indentation)?;
        write!(f, " BETWEEN ")?;
        self.start.fmt_indented(f, indentation)?;
        write!(f, " AND ")?;
        self.end.fmt_indented(f, indentation)?;
        Ok(())
    }
}

impl Display for WindowFrame {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        self.fmt_indented(f, Indent::default())
    }
}

#[derive(Debug, Clone, From, PartialEq, Eq)]
pub enum WindowReference {
    NamedWindowReference(NamedWindowReference),
    WindowSpecification(WindowSpecification),
}

impl WindowReference {
    pub fn fmt_indented(&self, f: &mut Formatter<'_>, indentation: Indent) -> fmt::Result {
        match self {
            WindowReference::NamedWindowReference(l) => l.fmt_indented(f, indentation),
            WindowReference::WindowSpecification(l) => l.fmt_indented(f, indentation),
        }
    }

    pub fn subelements(&self) -> usize {
        match self {
            WindowReference::NamedWindowReference(_) => 0,
            WindowReference::WindowSpecification(l) => l.subelements(),
        }
    }
}

impl Display for WindowReference {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        self.fmt_indented(f, Indent::default())
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct NamedWindowReference {
    pub name: SimpleIdentifier,
}

impl NamedWindowReference {
    pub fn new(name: SimpleIdentifier) -> Self {
        Self { name }
    }

    pub fn fmt_indented(&self, f: &mut Formatter<'_>, indentation: Indent) -> fmt::Result {
        self.name.fmt_indented(f, indentation)
    }
}

impl Display for NamedWindowReference {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        self.fmt_indented(f, Indent::default())
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct WindowExpression {
    pub expression: Box<SQLExpression>,
    pub window_reference: WindowReference,
}

impl WindowExpression {
    pub fn new(expression: SQLExpression, window_reference: WindowReference) -> Self {
        Self {
            expression: Box::new(expression),
            window_reference,
        }
    }

    pub fn fmt_indented(&self, f: &mut Formatter<'_>, indentation: Indent) -> fmt::Result {
        self.expression.fmt_indented(f, indentation)?;
        write!(f, " OVER ")?;
        self.window_reference.fmt_indented(f, indentation)
    }

    pub(crate) fn subelements(&self) -> usize {
        self.expression.subelements() + self.window_reference.subelements() + 1
    }
}

impl Display for WindowExpression {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        self.fmt_indented(f, Indent::default())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::sql::expression::literal::IntegerLiteral;

    #[test]
    fn range_between_two_preceding_renders_sql_correctly() {
        let wf = WindowFrame {
            frame_type: FrameType::RANGE,
            start: FrameEndpoint::Preceding(Box::new(IntegerLiteral::from_int(50).into())),
            end: FrameEndpoint::Preceding(Box::new(IntegerLiteral::from_int(5).into())),
        };
        assert_eq!(wf.to_string(), "RANGE BETWEEN 50 PRECEDING AND 5 PRECEDING");
    }

    #[test]
    fn rows_unbounded_preceding_to_current_row() {
        let wf = WindowFrame {
            frame_type: FrameType::ROWS,
            start: FrameEndpoint::UnboundedPreceding,
            end: FrameEndpoint::CurrentRow,
        };
        assert_eq!(
            wf.to_string(),
            "ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW"
        );
    }
}