use std::ops::Range;

use crate::{
    ParseError, Parser, ParserErrorKind,
    ast::{Expr, Ident, Path},
    lexer::Token,
};

#[derive(Debug)]
pub enum ProjectionKind<'a> {
    All,
    Columns(Vec<Path<'a>>),
}

#[derive(Debug)]
pub struct Projection<'a> {
    pub span: Range<usize>,
    pub kind: ProjectionKind<'a>,
}

#[derive(Debug)]
pub struct SelectFromStmt<'a> {
    pub span: Range<usize>,
    pub projection: Projection<'a>,
    pub table: Path<'a>,
    pub alias: Option<Ident<'a>>,
    pub selection: Option<Expr<'a>>,
}

impl<'a> Parser<'a> {
    pub(crate) fn parse_projection(&mut self) -> Result<Projection<'a>, ParseError> {
        let tok = self.lexer.peek();
        match tok.token {
            Token::Asterisk => {
                let span = tok.span.clone();
                self.lexer.advance();
                Ok(Projection {
                    span,
                    kind: ProjectionKind::All,
                })
            }
            Token::Identifier(_) => {
                let span_start = tok.span.start;
                let mut columns = Vec::new();
                loop {
                    let tok = self.lexer.peek();
                    match tok.token {
                        Token::Identifier(_) => {
                            let col = self.parse_path()?;
                            columns.push(col);
                        }
                        Token::Comma => self.lexer.advance(),
                        Token::From => break,
                        _ => {
                            let err = ParseError {
                                span: tok.span.clone(),
                                kind: ParserErrorKind::unexpected_token(
                                    &[Token::empty_ident(), Token::Comma, Token::From],
                                    &tok.token,
                                ),
                            };
                            if tok.token == Token::Eof {
                                return Err(err);
                            }
                            self.errors.push(err);
                            self.sync();
                        }
                    }
                }
                Ok(Projection {
                    span: span_start..self.current_span.end,
                    kind: ProjectionKind::Columns(columns),
                })
            }
            _ => Err(ParseError {
                span: tok.span.clone(),
                kind: ParserErrorKind::unexpected_token(
                    &[Token::Asterisk, Token::empty_ident()],
                    &tok.token,
                ),
            }),
        }
    }

    pub(crate) fn parse_select_stmt(&mut self) -> Result<SelectFromStmt<'a>, ParseError> {
        self.current_span.start = self.consume(&[Token::Select])?.span.start;
        let projection = self.parse_projection()?;
        self.consume(&[Token::From])?;
        let table = self.parse_path()?;

        let alias = if self.lexer.peek().token == Token::As {
            self.lexer.advance();
            let ident = self.parse_ident()?;
            Some(ident)
        } else {
            None
        };

        let selection = if self.lexer.peek().token == Token::Where {
            self.lexer.advance();
            let expr = self.parse_expr()?;
            Some(expr)
        } else {
            None
        };

        // TODO: proper checks in between sections, for better error messages on unexpected tokens
        self.consume(&[Token::Semicolon])?;

        Ok(SelectFromStmt {
            span: self.current_span.clone(),
            projection,
            table,
            alias,
            selection,
        })
    }
}

#[cfg(test)]
mod tests {
    use crate::{ast::*, parse, tests::assert_no_errors};

    #[test]
    fn test_select_all() {
        let source = "select * from mydb.users;";
        let (stmts, errors) = parse(source);
        assert_no_errors(&errors);
        let Stmt::SelectFrom(stmt) = &stmts[0] else {
            panic!("expected select")
        };
        assert!(matches!(stmt.projection.kind, ProjectionKind::All));
        assert_eq!(stmt.table.name().unwrap().name, "users".into());
    }

    #[test]
    fn test_select_columns() {
        let source = "select id, username from mydb.users;";
        let (stmts, errors) = parse(source);
        assert_no_errors(&errors);
        let Stmt::SelectFrom(stmt) = &stmts[0] else {
            panic!("expected select")
        };
        let ProjectionKind::Columns(cols) = &stmt.projection.kind else {
            panic!("expected columns")
        };
        assert_eq!(cols.len(), 2);
        assert!(matches!(&cols[0].segments[0], PathSegment::Named(i) if i.name == "id".into()));
        assert!(matches!(&cols[1].segments[0], PathSegment::Named(i) if i.name == "username".into()));
    }

    #[test]
    fn test_select_where_col_eq() {
        let source = "select * from mydb.users where id = 5;";
        let (stmts, errors) = parse(source);
        assert_no_errors(&errors);
        let Stmt::SelectFrom(stmt) = &stmts[0] else {
            panic!("expected select")
        };
        assert!(
            matches!(&stmt.selection, Some(Expr { kind: ExprKind::BinaryOp { left, op, right }, .. })
                if matches!(left.kind, ExprKind::Path(_)) &&
                matches!(right.kind, ExprKind::IntegerLiteral(_)) &&
                *op == BinaryOpKind::Eq
            )
        );
    }

    #[test]
    fn test_select_where_col_gt() {
        let source = "select * from mydb.users where id > 10;";
        let (stmts, errors) = parse(source);
        assert_no_errors(&errors);
        let Stmt::SelectFrom(stmt) = &stmts[0] else {
            panic!("expected select")
        };
        assert!(matches!(stmt.selection, Some(Expr { .. })))
    }

    #[test]
    fn test_select_where_col_gte_and_col_lt() {
        let source = "select * from mydb.users where age >= 18 and age < 30;";
        let (stmts, errors) = parse(source);
        assert_no_errors(&errors);
        let Stmt::SelectFrom(stmt) = &stmts[0] else {
            panic!("expected select")
        };
        assert!(matches!(stmt.selection, Some(Expr { .. })));
    }

    #[test]
    fn test_select_where_col_eq_or() {
        let source = "select * from mydb.users where age >= 18 or id = 42;";
        let (stmts, errors) = parse(source);
        assert_no_errors(&errors);
        let Stmt::SelectFrom(stmt) = &stmts[0] else {
            panic!("expected select")
        };
        assert!(matches!(stmt.selection, Some(Expr { .. })));
    }
}
