use std::ops::Range;

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

#[derive(Debug)]
pub struct GenericParams<'a> {
    pub span: Range<usize>,
    pub params: Vec<Ident<'a>>,
}

#[derive(Debug)]
pub struct GenericArgs<'a> {
    pub span: Range<usize>,
    pub args: Vec<TypeExpr<'a>>,
}

#[derive(Debug)]
pub struct TypeExpr<'a> {
    pub span: Range<usize>,
    pub path: Path<'a>,
    pub generic_args: GenericArgs<'a>,
}

#[derive(Debug)]
pub struct NamedField<'a> {
    pub span: Range<usize>,
    pub name: Ident<'a>,
    pub ty: TypeExpr<'a>,
}

#[derive(Debug)]
pub struct EnumVariant<'a> {
    pub span: std::ops::Range<usize>,
    pub name: Ident<'a>,
    pub fields: Fields<'a>,
}

#[derive(Debug)]
pub enum Fields<'a> {
    Unit,
    Unnamed(Vec<TypeExpr<'a>>),
    Named(Vec<NamedField<'a>>),
    Enum(Vec<EnumVariant<'a>>),
}

#[derive(Debug)]
pub struct CreateTypeStmt<'a> {
    pub span: Range<usize>,
    pub path: Path<'a>,
    pub generic_params: GenericParams<'a>,
    pub fields: Fields<'a>,
}

impl<'a> Parser<'a> {
    fn parse_generic_args(&mut self) -> Result<GenericArgs<'a>, ParseError> {
        let tok = self.lexer.peek();

        let mut args = Vec::new();
        let span_start = if tok.token == Token::LBracket {
            tok.span.start
        } else {
            return Ok(GenericArgs {
                span: self.current_span.end..self.current_span.end,
                args,
            });
        };
        self.lexer.advance();

        loop {
            let tok = self.lexer.peek();
            match tok.token {
                Token::Identifier(_) => {
                    let arg = self.parse_type_expr()?;
                    args.push(arg);
                }
                Token::Comma => self.lexer.advance(),
                Token::RBracket => break,
                _ => {
                    let err = ParseError {
                        span: tok.span.clone(),
                        kind: ParserErrorKind::unexpected_token(
                            &[Token::empty_ident(), Token::Comma, Token::RBracket],
                            &tok.token,
                        ),
                    };
                    if tok.token == Token::Eof {
                        return Err(err);
                    }
                    self.errors.push(err);
                    self.sync();
                }
            }
        }

        let span_end = self.consume(&[Token::RBracket])?.span.end;

        Ok(GenericArgs {
            span: span_start..span_end,
            args,
        })
    }

    pub(crate) fn parse_type_expr(&mut self) -> Result<TypeExpr<'a>, ParseError> {
        let path = self.parse_path()?;
        let generic_args = self.parse_generic_args()?;
        Ok(TypeExpr {
            span: path.span.start..generic_args.span.end,
            path,
            generic_args,
        })
    }

    fn parse_struct_decl_field(&mut self) -> Result<NamedField<'a>, ParseError> {
        let name = self.parse_ident()?;
        self.consume(&[Token::Colon])?;
        let type_expr = self.parse_type_expr()?;

        Ok(NamedField {
            span: name.span.start..type_expr.span.end,
            name,
            ty: type_expr,
        })
    }

    fn parse_fields_named(&mut self) -> Result<Fields<'a>, ParseError> {
        // TODO: do we want to collect the span here?
        self.consume(&[Token::LBrace])?.span.start;
        let mut fields = Vec::new();
        loop {
            let tok = self.lexer.peek();
            match tok.token {
                Token::Identifier(_) => {
                    match self.parse_struct_decl_field() {
                        Ok(field) => fields.push(field),
                        Err(err) => {
                            self.errors.push(err);
                            self.sync();
                            continue;
                        }
                    }
                    self.maybe_skip_comma();
                }
                Token::RBrace => break,
                _ => {
                    let err = ParseError {
                        span: tok.span.clone(),
                        kind: ParserErrorKind::unexpected_token(
                            &[Token::empty_ident(), Token::RBrace],
                            &tok.token,
                        ),
                    };
                    if tok.token == Token::Eof {
                        return Err(err);
                    }
                    self.errors.push(err);
                    self.lexer.advance();
                    self.sync();
                }
            }
        }
        self.consume(&[Token::RBrace])?.span.end;
        Ok(Fields::Named(fields))
    }

    fn parse_fields_unnamed(&mut self) -> Result<Fields<'a>, ParseError> {
        // TODO: do we want to collect the span here?
        self.consume(&[Token::LParen])?.span.start;
        let mut exprs = Vec::new();
        loop {
            let tok = self.lexer.peek();
            match tok.token {
                Token::Identifier(_) => {
                    match self.parse_type_expr() {
                        Ok(expr) => exprs.push(expr),
                        Err(err) => {
                            self.errors.push(err);
                            self.sync();
                            continue;
                        }
                    }
                    self.maybe_skip_comma();
                }
                Token::RParen => break,
                _ => {
                    let err = ParseError {
                        span: tok.span.clone(),
                        kind: ParserErrorKind::unexpected_token(
                            &[Token::empty_ident(), Token::RBrace],
                            &tok.token,
                        ),
                    };
                    if tok.token == Token::Eof {
                        return Err(err);
                    }
                    self.errors.push(err);
                    self.lexer.advance();
                    self.sync();
                }
            }
        }
        self.consume(&[Token::RParen])?.span.end;
        Ok(Fields::Unnamed(exprs))
    }

    fn parse_fields(&mut self) -> Result<Fields<'a>, ParseError> {
        let tok = self.lexer.peek();
        match tok.token {
            Token::LBrace => self.parse_fields_named(),
            Token::LParen => self.parse_fields_unnamed(),
            Token::Semicolon => Ok(Fields::Unit),
            _ => Err(ParseError {
                span: tok.span.clone(),
                kind: ParserErrorKind::unexpected_token(
                    &[Token::LBrace, Token::LParen, Token::Semicolon],
                    &tok.token,
                ),
            }),
        }
    }

    fn parse_enum_body(&mut self) -> Result<Fields<'a>, ParseError> {
        self.consume(&[Token::LBrace])?;
        let mut variants = Vec::new();
        loop {
            let tok = self.lexer.peek();
            match tok.token {
                Token::Identifier(_) => {
                    let name = self.parse_ident()?;
                    let span_start = name.span.start;
                    let fields = match self.lexer.peek().token {
                        Token::LBrace => self.parse_fields_named()?,
                        Token::LParen => self.parse_fields_unnamed()?,
                        _ => Fields::Unit,
                    };
                    let span_end = self.current_span.end;
                    variants.push(EnumVariant { span: span_start..span_end, name, fields });
                    self.maybe_skip_comma();
                }
                Token::RBrace => break,
                _ => {
                    let err = ParseError {
                        span: tok.span.clone(),
                        kind: ParserErrorKind::unexpected_token(
                            &[Token::empty_ident(), Token::RBrace],
                            &tok.token,
                        ),
                    };
                    if tok.token == Token::Eof {
                        return Err(err);
                    }
                    self.errors.push(err);
                    self.sync();
                }
            }
        }
        self.consume(&[Token::RBrace])?;
        Ok(Fields::Enum(variants))
    }

    fn parse_generic_params(&mut self) -> Result<GenericParams<'a>, ParseError> {
        let tok = self.lexer.peek();
        let mut params = Vec::new();
        let span_start = if tok.token == Token::LBracket {
            tok.span.start
        } else {
            return Ok(GenericParams {
                span: self.current_span.end..self.current_span.end,
                params,
            });
        };
        self.lexer.advance();

        loop {
            let tok = self.lexer.peek();
            match &tok.token {
                Token::Identifier(name) => {
                    params.push(Ident {
                        span: tok.span.clone(),
                        name: name.clone(),
                    });
                    self.lexer.advance();
                }
                Token::Comma => self.lexer.advance(),
                Token::RBracket => break,
                _ => {
                    let err = ParseError {
                        span: tok.span.clone(),
                        kind: ParserErrorKind::unexpected_token(
                            &[Token::empty_ident(), Token::Comma, Token::RBracket],
                            &tok.token,
                        ),
                    };
                    if tok.token == Token::Eof {
                        return Err(err);
                    }
                    self.errors.push(err);
                    self.sync();
                }
            }
        }

        let span_end = self.consume(&[Token::RBracket])?.span.end;

        Ok(GenericParams {
            span: span_start..span_end,
            params,
        })
    }

    /// Assumes that `create` has already been parsed and the span is set.
    #[instrument(skip_all, level = "trace")]
    pub(crate) fn parse_create_type_stmt(&mut self) -> Result<CreateTypeStmt<'a>, ParseError> {
        self.consume(&[Token::Type])?;
        let path = self.parse_path()?;
        let generic_params = self.parse_generic_params()?;
        let tok = self.lexer.peek();
        match tok.token {
            Token::Struct => {
                self.lexer.advance();
            }
            Token::Enum => {
                self.lexer.advance();
                let fields = self.parse_enum_body()?;
                self.consume(&[Token::Semicolon])?;
                return Ok(CreateTypeStmt {
                    span: self.current_span.clone(),
                    path,
                    generic_params,
                    fields,
                });
            }
            _ => {
                return Err(ParseError {
                    span: tok.span.clone(),
                    kind: ParserErrorKind::unexpected_token(
                        &[Token::Struct, Token::Enum],
                        &tok.token,
                    ),
                });
            }
        }
        let fields = self.parse_fields()?;
        self.consume(&[Token::Semicolon])?;
        Ok(CreateTypeStmt {
            span: self.current_span.clone(),
            path,
            generic_params,
            fields,
        })
    }
}

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

    #[test]
    fn test_create_type_struct() {
        let source = "create type User struct {
            id       : Int64,
            username : String,
        };";

        let (statements, errors) = parse(source);
        assert_no_errors(&errors);

        let Stmt::CreateType(CreateTypeStmt { path, fields, .. }) = &statements[0] else {
            panic!("invalid statement type: {:?}", &statements[0]);
        };

        let Fields::Named(fields) = fields else {
            panic!("expected named fields, got {:?}", fields);
        };

        assert_eq!(path.name().unwrap().name, "User".into());
        assert_eq!(fields[0].name.name, "id".into());
        assert_eq!(fields[0].ty.path.name().unwrap().name, "Int64".into());
        assert_eq!(fields[1].name.name, "username".into());
        assert_eq!(fields[1].ty.path.name().unwrap().name, "String".into());
        assert_eq!(fields.len(), 2);
        assert_eq!(statements.len(), 1);
    }

    #[test]
    fn test_create_type_qualified_path() {
        let source = "create type main.User struct { id: Int64 };";

        let (statements, errors) = parse(source);
        assert_no_errors(&errors);

        let Stmt::CreateType(CreateTypeStmt { path, .. }) = &statements[0] else {
            panic!("invalid statement type: {:?}", &statements[0]);
        };

        assert_eq!(path.database().unwrap().name, "main".into());
        assert_eq!(path.name().unwrap().name, "User".into());
    }

    #[test]
    fn test_create_type_no_params() {
        let source = "create type User struct { id: Int64 };";
        let (stmts, errors) = parse(source);
        assert_no_errors(&errors);
        let Stmt::CreateType(stmt) = &stmts[0] else {
            panic!()
        };
        assert!(stmt.generic_params.params.is_empty());
    }

    #[test]
    fn test_create_type_single_param() {
        let source = "create type Option[T] struct { value: T };";
        let (stmts, errors) = parse(source);
        assert_no_errors(&errors);
        let Stmt::CreateType(stmt) = &stmts[0] else {
            panic!("invalid statement type: {:?}", &stmts[0]);
        };

        let Fields::Named(fields) = &stmt.fields else {
            panic!("expected named fields, got {:?}", stmt.fields);
        };

        let params = &stmt.generic_params.params;
        assert_eq!(params.len(), 1);
        assert_eq!(params[0].name, "T".into());
        assert_eq!(fields[0].ty.path.name().unwrap().name, "T".into());
    }

    #[test]
    fn test_create_type_multiple_params() {
        let source = "create type Map[K, V] struct { k: K, v: V };";
        let (stmts, errors) = parse(source);
        assert_no_errors(&errors);
        let Stmt::CreateType(stmt) = &stmts[0] else {
            panic!()
        };
        let params = &stmt.generic_params.params;
        assert_eq!(params.len(), 2);
        assert_eq!(params[0].name, "K".into());
        assert_eq!(params[1].name, "V".into());
    }

    #[test]
    fn test_create_type_param_with_qualified_field_type() {
        // field types can still be qualified paths even when params are present
        let source = "create type Wrapper[T] struct { inner: main.SomeType };";
        let (stmts, errors) = parse(source);
        assert_no_errors(&errors);
        let Stmt::CreateType(stmt) = &stmts[0] else {
            panic!()
        };
        let Fields::Named(fields) = &stmt.fields else {
            panic!("expected named fields, got {:?}", stmt.fields);
        };

        assert!(!stmt.generic_params.params.is_empty());
        assert_eq!(
            fields[0].ty.path.database().unwrap().name,
            "main".into()
        );
        assert_eq!(fields[0].ty.path.name().unwrap().name, "SomeType".into());
    }

    #[test]
    fn test_create_type_unit_struct() {
        let source = "create type Empty struct;";
        let (stmts, errors) = parse(source);
        assert_no_errors(&errors);
        let Stmt::CreateType(stmt) = &stmts[0] else {
            panic!("invalid statement type: {:?}", &stmts[0]);
        };
        assert!(matches!(stmt.fields, Fields::Unit));
        assert_eq!(stmt.path.name().unwrap().name, "Empty".into());
    }

    #[test]
    fn test_create_type_tuple_struct() {
        let source = "create type Point struct(Int64, String);";
        let (stmts, errors) = parse(source);
        assert_no_errors(&errors);
        let Stmt::CreateType(stmt) = &stmts[0] else {
            panic!("invalid statement type: {:?}", &stmts[0]);
        };
        let Fields::Unnamed(fields) = &stmt.fields else {
            panic!("expected unnamed fields, got {:?}", stmt.fields);
        };
        assert_eq!(fields.len(), 2);
        assert_eq!(fields[0].path.name().unwrap().name, "Int64".into());
        assert_eq!(fields[1].path.name().unwrap().name, "String".into());
    }

    #[test]
    fn test_create_type_tuple_struct_single_field() {
        let source = "create type Wrapper struct(Int64);";
        let (stmts, errors) = parse(source);
        assert_no_errors(&errors);
        let Stmt::CreateType(stmt) = &stmts[0] else {
            panic!("invalid statement type: {:?}", &stmts[0]);
        };
        let Fields::Unnamed(fields) = &stmt.fields else {
            panic!("expected unnamed fields, got {:?}", stmt.fields);
        };
        assert_eq!(fields.len(), 1);
        assert_eq!(fields[0].path.name().unwrap().name, "Int64".into());
    }

    #[test]
    fn test_create_type_tuple_struct_with_generic() {
        let source = "create type Box[T] struct(T);";
        let (stmts, errors) = parse(source);
        assert_no_errors(&errors);
        let Stmt::CreateType(stmt) = &stmts[0] else {
            panic!()
        };
        let Fields::Unnamed(fields) = &stmt.fields else {
            panic!("expected unnamed fields, got {:?}", stmt.fields);
        };
        assert_eq!(stmt.generic_params.params.len(), 1);
        assert_eq!(stmt.generic_params.params[0].name, "T".into());
        assert_eq!(fields.len(), 1);
        assert_eq!(fields[0].path.name().unwrap().name, "T".into());
    }

    #[test]
    fn test_type_expr_simple() {
        let mut parser = Parser::new("Int64");
        let expr = parser.parse_type_expr().unwrap();
        assert_eq!(expr.path.name().unwrap().name, "Int64".into());
        assert!(expr.path.database().is_none());
        assert!(expr.generic_args.args.is_empty());
        parser.assert_finished();
    }

    #[test]
    fn test_type_expr_qualified() {
        let mut parser = Parser::new("main.User");
        let expr = parser.parse_type_expr().unwrap();
        assert_eq!(expr.path.database().unwrap().name, "main".into());
        assert_eq!(expr.path.name().unwrap().name, "User".into());
        assert!(expr.generic_args.args.is_empty());
        parser.assert_finished();
    }

    #[test]
    fn test_type_expr_single_arg() {
        let mut parser = Parser::new("Option[Int64]");
        let expr = parser.parse_type_expr().unwrap();
        assert_eq!(expr.path.name().unwrap().name, "Option".into());
        assert_eq!(expr.generic_args.args.len(), 1);
        assert_eq!(expr.generic_args.args[0].path.name().unwrap().name, "Int64".into());
        parser.assert_finished();
    }

    #[test]
    fn test_type_expr_multiple_args() {
        let mut parser = Parser::new("Map[String, Int64]");
        let expr = parser.parse_type_expr().unwrap();
        assert_eq!(expr.path.name().unwrap().name, "Map".into());
        assert_eq!(expr.generic_args.args.len(), 2);
        assert_eq!(expr.generic_args.args[0].path.name().unwrap().name, "String".into());
        assert_eq!(expr.generic_args.args[1].path.name().unwrap().name, "Int64".into());
        parser.assert_finished();
    }

    #[test]
    fn test_type_expr_qualified_arg() {
        let mut parser = Parser::new("Option[main.User]");
        let expr = parser.parse_type_expr().unwrap();
        assert_eq!(expr.path.name().unwrap().name, "Option".into());
        assert_eq!(expr.generic_args.args.len(), 1);
        let arg = &expr.generic_args.args[0];
        assert_eq!(arg.path.database().unwrap().name, "main".into());
        assert_eq!(arg.path.name().unwrap().name, "User".into());
        parser.assert_finished();
    }

    #[test]
    fn test_type_expr_nested() {
        let mut parser = Parser::new("Option[Map[String, Int64]]");
        let expr = parser.parse_type_expr().unwrap();
        assert_eq!(expr.path.name().unwrap().name, "Option".into());
        assert_eq!(expr.generic_args.args.len(), 1);
        let inner = &expr.generic_args.args[0];
        assert_eq!(inner.path.name().unwrap().name, "Map".into());
        assert_eq!(inner.generic_args.args.len(), 2);
        assert_eq!(inner.generic_args.args[0].path.name().unwrap().name, "String".into());
        assert_eq!(inner.generic_args.args[1].path.name().unwrap().name, "Int64".into());
        parser.assert_finished();
    }

    #[test]
    fn test_type_expr_args_span() {
        let src = "Option[Int64]";
        let mut parser = Parser::new(src);
        let expr = parser.parse_type_expr().unwrap();
        assert_eq!(&src[expr.span.clone()], "Option[Int64]");
        assert_eq!(&src[expr.generic_args.span.clone()], "[Int64]");
        parser.assert_finished();
    }

    // -- enum type declarations --

    #[test]
    fn test_create_unit_enum() {
        let source = "create type main.Color enum { Red, Green, Blue };";
        let (stmts, errors) = parse(source);
        assert_no_errors(&errors);
        let Stmt::CreateType(stmt) = &stmts[0] else { panic!() };
        let Fields::Enum(variants) = &stmt.fields else { panic!("expected Enum fields, got {:?}", stmt.fields) };
        assert_eq!(variants.len(), 3);
        assert_eq!(variants[0].name.name, "Red".into());
        assert!(matches!(variants[0].fields, Fields::Unit));
        assert_eq!(variants[1].name.name, "Green".into());
        assert!(matches!(variants[1].fields, Fields::Unit));
        assert_eq!(variants[2].name.name, "Blue".into());
        assert!(matches!(variants[2].fields, Fields::Unit));
    }

    #[test]
    fn test_create_enum_tuple_variant() {
        let source = "create type main.Option[T] enum { None, Some(T) };";
        let (stmts, errors) = parse(source);
        assert_no_errors(&errors);
        let Stmt::CreateType(stmt) = &stmts[0] else { panic!() };
        let Fields::Enum(variants) = &stmt.fields else { panic!() };
        assert_eq!(variants.len(), 2);
        assert_eq!(variants[0].name.name, "None".into());
        assert!(matches!(variants[0].fields, Fields::Unit));
        assert_eq!(variants[1].name.name, "Some".into());
        let Fields::Unnamed(types) = &variants[1].fields else { panic!() };
        assert_eq!(types.len(), 1);
        assert_eq!(types[0].path.name().unwrap().name, "T".into());
    }

    #[test]
    fn test_create_enum_named_variant() {
        let source = "create type main.Shape enum { Circle { radius: Int64 }, Rect { w: Int64, h: Int64 } };";
        let (stmts, errors) = parse(source);
        assert_no_errors(&errors);
        let Stmt::CreateType(stmt) = &stmts[0] else { panic!() };
        let Fields::Enum(variants) = &stmt.fields else { panic!() };
        assert_eq!(variants.len(), 2);
        assert_eq!(variants[0].name.name, "Circle".into());
        let Fields::Named(circle_fields) = &variants[0].fields else { panic!() };
        assert_eq!(circle_fields[0].name.name, "radius".into());
        assert_eq!(variants[1].name.name, "Rect".into());
        let Fields::Named(rect_fields) = &variants[1].fields else { panic!() };
        assert_eq!(rect_fields.len(), 2);
    }

    #[test]
    fn test_create_enum_mixed_variants() {
        let source = "create type main.Result[T, E] enum { Ok(T), Err(E), Unknown };";
        let (stmts, errors) = parse(source);
        assert_no_errors(&errors);
        let Stmt::CreateType(stmt) = &stmts[0] else { panic!() };
        let Fields::Enum(variants) = &stmt.fields else { panic!() };
        assert_eq!(variants.len(), 3);
        assert_eq!(variants[0].name.name, "Ok".into());
        let Fields::Unnamed(_) = &variants[0].fields else { panic!() };
        assert_eq!(variants[1].name.name, "Err".into());
        let Fields::Unnamed(_) = &variants[1].fields else { panic!() };
        assert_eq!(variants[2].name.name, "Unknown".into());
        assert!(matches!(variants[2].fields, Fields::Unit));
    }
}
