use std::ops::Range;

use crate::{
    ParseError, Parser, ParserErrorKind,
    ast::{Path, PrimaryKey, TypeExpr},
    lexer::Token,
};

#[derive(Debug)]
pub struct TableDeclBody<'a> {
    pub span: Range<usize>,
    pub primary_key: PrimaryKey<'a>,
}

#[derive(Debug)]
pub struct CreateTableStmt<'a> {
    pub span: Range<usize>,
    pub path: Path<'a>,
    pub ty: TypeExpr<'a>,
    pub body: TableDeclBody<'a>,
}

impl<'a> Parser<'a> {
    #[instrument(skip_all, level = "trace")]
    pub(crate) fn parse_table_decl_body(&mut self) -> Result<TableDeclBody<'a>, ParseError> {
        let body_span_start = self.consume(&[Token::LBrace])?.span.start;
        let mut primary_key = None;
        loop {
            let tok = self.lexer.peek();
            match tok.token {
                Token::Primary => {
                    if primary_key.is_some() {
                        self.errors.push(ParseError {
                            span: tok.span.clone(),
                            kind: ParserErrorKind::DuplicatePrimaryKey,
                        });
                        self.sync();
                        continue;
                    }
                    let pk = match self.parse_primary_key() {
                        Ok(pk) => pk,
                        Err(err) => {
                            self.errors.push(err);
                            self.sync();
                            continue;
                        }
                    };
                    primary_key = Some(pk);
                }
                Token::Comma => self.lexer.advance(),
                Token::RBrace => break,
                _ => {
                    let err = ParseError {
                        span: tok.span.clone(),
                        kind: ParserErrorKind::unexpected_token(
                            &[Token::Primary, Token::Comma, Token::RBrace],
                            &tok.token,
                        ),
                    };
                    if tok.token == Token::Eof {
                        return Err(err);
                    }
                    self.errors.push(err);
                    self.sync();
                }
            }
        }
        let body_span_end = self.consume(&[Token::RBrace])?.span.end;

        Ok(TableDeclBody {
            span: body_span_start..body_span_end,
            primary_key: primary_key.ok_or_else(|| ParseError {
                // TODO: use the whole span, or just the end of the body?
                span: self.current_span.clone(),
                kind: ParserErrorKind::MissingPrimaryKey,
            })?,
        })
    }

    /// Assumes that `create` has already been parsed and the span is set.
    #[instrument(skip_all, level = "trace")]
    pub(crate) fn parse_create_table_stmt(&mut self) -> Result<CreateTableStmt<'a>, ParseError> {
        self.consume(&[Token::Table])?;
        let path = self.parse_path()?;
        let _colon = self.consume(&[Token::Colon])?;
        let ty = self.parse_type_expr()?;

        // -- parse table body --
        let body = self.parse_table_decl_body()?;

        self.consume(&[Token::Semicolon])?;

        Ok(CreateTableStmt {
            span: self.current_span.clone(),
            path,
            body,
            ty,
        })
    }
}

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

    #[test]
    fn test_create_table() {
        let source = "create table mydb.users : User {
            primary key (id),
        };";

        let (statements, errors) = parse(source);
        assert_no_errors(&errors);
        let Stmt::CreateTable(CreateTableStmt { path, ty, body, .. }) = &statements[0] else {
            panic!("invalid statement type: {:?}", &statements[0]);
        };
        assert_eq!(path.database().unwrap().name, "mydb".into());
        assert_eq!(path.name().unwrap().name, "users".into());
        assert_eq!(ty.path.name().unwrap().name, "User".into());
        assert!(
            matches!(&body.primary_key.columns[0].segments[0], PathSegment::Named(id) if id.name.as_ref() == "id")
        );
        assert_eq!(body.primary_key.columns.len(), 1);
        assert_eq!(statements.len(), 1);
    }

    #[test]
    fn test_create_table_index_primary_key() {
        let source = "create table main.points : Point { primary key (0) };";
        let (stmts, errors) = parse(source);
        assert_no_errors(&errors);
        let Stmt::CreateTable(stmt) = &stmts[0] else {
            panic!("invalid statement type: {:?}", &stmts[0]);
        };
        assert!(matches!(
            stmt.body.primary_key.columns[0].segments[0],
            PathSegment::Indexed(0)
        ));
        assert_eq!(stmt.body.primary_key.columns.len(), 1);
    }

    #[test]
    fn test_create_table_composite_index_primary_key() {
        let source = "create table main.edges : Edge { primary key (0, 1) };";
        let (stmts, errors) = parse(source);
        assert_no_errors(&errors);
        let Stmt::CreateTable(stmt) = &stmts[0] else {
            panic!()
        };
        assert!(matches!(
            stmt.body.primary_key.columns[0].segments[0],
            PathSegment::Indexed(0)
        ));
        assert!(matches!(
            stmt.body.primary_key.columns[1].segments[0],
            PathSegment::Indexed(1)
        ));
        assert_eq!(stmt.body.primary_key.columns.len(), 2);
    }
}
