use super::{
    intern::{Index as Interned, PointerFlags, StructFlags},
    visitor::AstExt,
    Ast, Index, ParseError, Tag,
};

pub trait AstNodeExt {
    fn get_ast_node(&self, idx: Index) -> AstNode;
}

impl AstNodeExt for &mut Ast {
    fn get_ast_node(&self, idx: Index) -> AstNode {
        <Ast as AstNodeExt>::get_ast_node(self, idx)
    }
}

impl AstNodeExt for &Ast {
    fn get_ast_node(&self, idx: Index) -> AstNode {
        <Ast as AstNodeExt>::get_ast_node(self, idx)
    }
}

impl AstNodeExt for Ast {
    fn get_ast_node(&self, idx: Index) -> AstNode {
        let (tag, data) = self.get_node_tag_and_data(idx);

        match tag {
            Tag::Root => {
                unreachable!()
            }
            Tag::File => {
                let (a, b) = data.as_extra_range();
                let decls = unsafe { Index::from_slice_unchecked(&self.extra[a..b]).to_vec() };

                AstNode::File { decls }
            }
            Tag::FunctionProto => {
                let (name, extra) = data.as_intern_and_extra_offset();

                let (return_type, parameter_list) = (
                    Index::from_u32(self.extra[extra]).unwrap(),
                    Index::from_u32(self.extra[extra + 1]).unwrap(),
                );

                AstNode::FunctionProto {
                    name,
                    return_type,
                    parameter_list,
                }
            }
            Tag::FunctionProtoInterned => {
                let (name, ty) = data.as_two_interns();
                AstNode::FunctionProtoInterned { name, ty }
            }
            Tag::FunctionDecl => {
                let (proto, body) = data.as_two_indices();

                AstNode::FunctionDecl { proto, body }
            }
            Tag::ParameterList => {
                let (a, b) = data.as_extra_range();
                let params = unsafe { Index::from_slice_unchecked(&self.extra[a..b]).to_vec() };

                AstNode::ParameterList { params }
            }
            Tag::Parameter => {
                let (ty, name) = data.as_index_intern();

                AstNode::Parameter { ty, name }
            }
            Tag::Block => {
                let (a, b) = data.as_extra_range();
                let statements = unsafe { Index::from_slice_unchecked(&self.extra[a..b]).to_vec() };

                AstNode::Block {
                    statements,
                    expr: None,
                }
            }
            Tag::BlockTrailingExpr => {
                let (a, b) = data.as_extra_range();
                let (expr, statements) = unsafe {
                    Index::from_slice_unchecked(&self.extra[a..b])
                        .split_last()
                        .unwrap()
                };

                AstNode::Block {
                    statements: statements.to_vec(),
                    expr: Some(*expr),
                }
            }
            Tag::Constant => {
                let (ty, value) = data.as_index_intern();
                AstNode::Constant { ty, value }
            }
            Tag::ExprStmt => AstNode::ExprStmt {
                expr: data.as_index(),
            },
            Tag::ReturnStmt => AstNode::ReturnStmt,
            Tag::ReturnExprStmt => AstNode::ReturnExprStmt {
                expr: data.as_index(),
            },
            Tag::VarDecl => {
                let (a, _) = data.as_extra_range();
                let name = Interned::from_u32(self.extra[a]);
                let ty = Index::from_u32(self.extra[a + 1]).unwrap();

                AstNode::VarDecl { name, ty }
            }
            Tag::MutVarDecl => {
                let (a, _) = data.as_extra_range();
                let name = Interned::from_u32(self.extra[a]);
                let ty = Index::from_u32(self.extra[a + 1]).unwrap();

                AstNode::MutVarDecl { name, ty }
            }
            Tag::VarDeclAssignment => {
                let (a, b) = data.as_extra_range();
                let extra = &self.extra[a..b];
                let name = Interned::from_u32(*extra.get(0).unwrap());
                let expr = Index::from_u32(*extra.get(1).unwrap()).unwrap();
                let ty = extra.get(2).map(|&inner| Index::from_u32(inner).unwrap());

                AstNode::MutVarDeclAssignment { name, expr, ty }
            }
            Tag::MutVarDeclAssignment => {
                let (a, b) = data.as_extra_range();
                let extra = &self.extra[a..b];
                let name = Interned::from_u32(*extra.get(0).unwrap());
                let expr = Index::from_u32(*extra.get(1).unwrap()).unwrap();
                let ty = extra.get(2).map(|&inner| Index::from_u32(inner).unwrap());

                AstNode::MutVarDeclAssignment { name, expr, ty }
            }
            Tag::GlobalDecl => {
                let (name, offset) = data.as_intern_and_extra_offset();
                let ty = Index::from_u32(self.extra[offset]).unwrap();
                let expr = Index::from_u32(self.extra[offset + 1]).unwrap();

                AstNode::GlobalDecl { name, expr, ty }
            }
            Tag::StructDecl => {
                let (name, offset) = data.as_intern_and_extra_offset();
                let flags = StructFlags::unpack(self.extra[offset]);

                let types = (offset + 1)..(offset + 1 + flags.num_fields as usize);
                let names = (offset + 1 + flags.num_fields as usize)
                    ..(offset + 1 + flags.num_fields as usize * 2);

                let field_types =
                    unsafe { Index::from_slice_unchecked(&self.extra[types]).to_vec() };

                let field_names = self.extra[names]
                    .iter()
                    .map(|&i| Interned::from_u32(i))
                    .collect();

                AstNode::StructDecl {
                    name,
                    flags,
                    field_names,
                    field_types,
                }
            }
            Tag::StructDeclInterned => {
                let (name, ty) = data.as_two_interns();

                AstNode::StructDeclInterned { name, ty }
            }
            Tag::FieldDecl => {
                let (ty, name) = data.as_index_intern();

                AstNode::FieldDecl { name, ty }
            }
            Tag::DeclRef => AstNode::DeclRef {
                decl: data.as_index(),
            },
            Tag::DeclRefUnresolved => {
                let (scope, name) = data.as_index_intern();
                AstNode::DeclRefUnresolved { scope, name }
            }
            Tag::InternedType => AstNode::InternedType {
                intern: data.as_intern(),
            },
            Tag::TypeDeclRef => AstNode::TypeDeclRef {
                decl: data.as_index(),
            },
            Tag::TypeDeclRefUnresolved => {
                let (scope, name) = data.as_index_intern();
                AstNode::TypeDeclRefUnresolved { scope, name }
            }
            Tag::PointerType => {
                let (ty, flags) = data.as_index_and_opaque();
                let flags = PointerFlags::unpack(flags as u8);
                AstNode::PointerType { ty, flags }
            }
            Tag::ArrayType => {
                let (length, pointer) = data.as_two_indices();

                AstNode::ArrayType { length, pointer }
            }
            Tag::CallExpr => {
                let (func, argument_list) = data.as_two_indices();
                AstNode::CallExpr {
                    func,
                    argument_list,
                }
            }
            Tag::FieldAccess => {
                let (expr, field_name) = data.as_index_intern();
                AstNode::FieldAccess { field_name, expr }
            }
            Tag::ArgumentList => {
                let (a, b) = data.as_extra_range();
                let arguments = unsafe { Index::from_slice_unchecked(&self.extra[a..b]).to_vec() };

                AstNode::ArgumentList { arguments }
            }
            Tag::Argument => AstNode::Argument {
                expr: data.as_index(),
                name: None,
            },
            Tag::NamedArgument => {
                let (expr, name) = data.as_index_intern();
                AstNode::Argument {
                    expr,
                    name: Some(name),
                }
            }
            Tag::ExplicitCast => {
                let (expr, ty) = data.as_two_indices();

                AstNode::ExplicitCast { expr, ty }
            }
            Tag::Deref => AstNode::Deref {
                expr: data.as_index(),
            },
            Tag::AddressOf => AstNode::AddressOf {
                expr: data.as_index(),
            },
            Tag::Not => AstNode::Not {
                expr: data.as_index(),
            },
            Tag::Negate => AstNode::Negate {
                expr: data.as_index(),
            },
            Tag::PlaceToValueConversion => AstNode::PlaceToValueConversion {
                expr: data.as_index(),
            },
            Tag::ValueToPlaceConversion => AstNode::ValueToPlaceConversion {
                expr: data.as_index(),
            },
            Tag::Or => {
                let (lhs, rhs) = data.as_two_indices();
                AstNode::Or { lhs, rhs }
            }
            Tag::And => {
                let (lhs, rhs) = data.as_two_indices();
                AstNode::And { lhs, rhs }
            }
            Tag::BitOr => {
                let (lhs, rhs) = data.as_two_indices();
                AstNode::BitOr { lhs, rhs }
            }
            Tag::BitXOr => {
                let (lhs, rhs) = data.as_two_indices();
                AstNode::BitXOr { lhs, rhs }
            }
            Tag::BitAnd => {
                let (lhs, rhs) = data.as_two_indices();
                AstNode::BitAnd { lhs, rhs }
            }
            Tag::Eq => {
                let (lhs, rhs) = data.as_two_indices();
                AstNode::Eq { lhs, rhs }
            }
            Tag::NEq => {
                let (lhs, rhs) = data.as_two_indices();
                AstNode::NEq { lhs, rhs }
            }
            Tag::Lt => {
                let (lhs, rhs) = data.as_two_indices();
                AstNode::Lt { lhs, rhs }
            }
            Tag::Gt => {
                let (lhs, rhs) = data.as_two_indices();
                AstNode::Gt { lhs, rhs }
            }
            Tag::Le => {
                let (lhs, rhs) = data.as_two_indices();
                AstNode::Le { lhs, rhs }
            }
            Tag::Ge => {
                let (lhs, rhs) = data.as_two_indices();
                AstNode::Ge { lhs, rhs }
            }
            Tag::Shl => {
                let (lhs, rhs) = data.as_two_indices();
                AstNode::Shl { lhs, rhs }
            }
            Tag::Shr => {
                let (lhs, rhs) = data.as_two_indices();
                AstNode::Shr { lhs, rhs }
            }
            Tag::Add => {
                let (lhs, rhs) = data.as_two_indices();
                AstNode::Add { lhs, rhs }
            }
            Tag::Sub => {
                let (lhs, rhs) = data.as_two_indices();
                AstNode::Sub { lhs, rhs }
            }
            Tag::Mul => {
                let (lhs, rhs) = data.as_two_indices();
                AstNode::Mul { lhs, rhs }
            }
            Tag::Div => {
                let (lhs, rhs) = data.as_two_indices();
                AstNode::Div { lhs, rhs }
            }
            Tag::Rem => {
                let (lhs, rhs) = data.as_two_indices();
                AstNode::Rem { lhs, rhs }
            }
            Tag::Assign => {
                let (lhs, rhs) = data.as_two_indices();
                AstNode::Assign { lhs, rhs }
            }
            Tag::SubscriptExpr => {
                let (lhs, index) = data.as_two_indices();
                AstNode::SubscriptExpr { lhs, index }
            }
            Tag::IfExpr => {
                let (cond, body) = data.as_two_indices();
                AstNode::IfExpr { cond, body }
            }
            Tag::IfElseExpr => {
                let (cond, extra) = data.as_index_and_extra_offset();
                let [a, b] = self.extra[extra..][..2] else {
                    unreachable!()
                };

                AstNode::IfElseExpr {
                    cond,
                    a: Index::from_u32(a).unwrap(),
                    b: Index::from_u32(b).unwrap(),
                }
            }
            Tag::Error => AstNode::Error {
                err: data.as_error(),
            },
            Tag::Undefined => AstNode::Undefined,
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum AstNode {
    /// pseudo tag, contains a range from a..b into extra of all files.
    Root {
        files: Vec<Index>,
    },
    /// `data` is a range from a..b into extra of all global nodes.
    File {
        decls: Vec<Index>,
    },
    /// `data` is an intern to a name, and an index into extra of [index: return_type, index: ParameterList]
    FunctionProto {
        name: Interned,
        return_type: Index,
        parameter_list: Index,
    },
    /// `data` is an intern to a name, and an intern to the function type
    FunctionProtoInterned {
        name: Interned,
        ty: Interned,
    },
    /// `data` is an index to a FunctionProto and an index to a Block
    FunctionDecl {
        proto: Index,
        body: Index,
    },
    /// `data` is a range from a..b into extra of indices to parameters
    ParameterList {
        params: Vec<Index>,
    },
    /// `data` is an index to a type, and an intern to a name
    Parameter {
        ty: Index,
        name: Interned,
    },
    /// `data` is range from a..b into `extra` of indices to statements
    Block {
        statements: Vec<Index>,
        expr: Option<Index>,
    },
    /// `data` is an index to a type, and an intern to a value
    Constant {
        ty: Index,
        value: Interned,
    },
    /// `data` is an index to an expression
    ExprStmt {
        expr: Index,
    },
    /// `data` is none
    ReturnStmt,
    /// `data` is an index to an expr
    ReturnExprStmt {
        expr: Index,
    },
    /// `data` is a range from a..b into `extra` of `[name: intern, type: index]`
    VarDecl {
        name: Interned,
        ty: Index,
    },
    /// `data` is a range from a..b into `extra` of `[name: intern, type: index]`
    MutVarDecl {
        name: Interned,
        ty: Index,
    },
    /// `data` is a range from a..b into `extra` of `[name: intern, expr: index, type?: index]`
    VarDeclAssignment {
        name: Interned,
        expr: Index,
        ty: Option<Index>,
    },
    /// `data` is a range from a..b into `extra` of `[name: intern, expr: index, type?: index]`
    MutVarDeclAssignment {
        name: Interned,
        expr: Index,
        ty: Option<Index>,
    },
    /// `data` is an intern to a name, and an offset into `extra` of `[type: index, expr: index]`
    GlobalDecl {
        name: Interned,
        expr: Index,
        ty: Index,
    },
    /// `data` is an intern to a name, and an offset into extra of `[flags, type0 ,..., typeN ,name0 ,..., nameN]`
    StructDecl {
        name: Interned,
        flags: StructFlags,
        field_names: Vec<Interned>,
        field_types: Vec<Index>,
    },
    /// `data` is an intern to a name, and an intern to the type of the struct
    StructDeclInterned {
        name: Interned,
        ty: Interned,
    },
    /// `data` is an index to a type, and an intern to a name
    FieldDecl {
        name: Interned,
        ty: Index,
    },
    /// `data` is an index to a Parameter, VarDecl, GlobalDecl or FunctionDecl, and an opaque DeclKind
    DeclRef {
        decl: Index,
    },
    /// `data` is an inlined key into the symbol table (scope: index, name: intern)
    DeclRefUnresolved {
        scope: Index,
        name: Interned,
    },
    /// `data` is an intern of a type
    InternedType {
        intern: Interned,
    },
    /// `data` is an index to a StructDecl
    TypeDeclRef {
        decl: Index,
    },
    /// `data` is an inlined key into the symbol table (scope: index, name: intern)
    TypeDeclRefUnresolved {
        scope: Index,
        name: Interned,
    },
    /// `data` is an index to a Type and u32 PointerFlags (extra offset)
    PointerType {
        ty: Index,
        flags: PointerFlags,
    },
    /// `data` is an index to a length expression, and an underlying pointer type
    ArrayType {
        length: Index,
        pointer: Index,
    },
    /// `data` is an index to an expr and an index to an ArgumentList
    CallExpr {
        func: Index,
        argument_list: Index,
    },
    /// `data` is an index to an expr and an intern to a field name
    FieldAccess {
        field_name: Interned,
        expr: Index,
    },
    /// `data` is a range from a..b into extra of indices to arguments
    ArgumentList {
        arguments: Vec<Index>,
    },
    /// `data` is an index to an expression
    Argument {
        expr: Index,
        name: Option<Interned>,
    },
    /// `data` is an index to lhs, and an index to the type
    ExplicitCast {
        expr: Index,
        ty: Index,
    },
    /// `data` is a single index to an expr
    Deref {
        expr: Index,
    },
    AddressOf {
        expr: Index,
    },
    Not {
        expr: Index,
    },
    Negate {
        expr: Index,
    },
    PlaceToValueConversion {
        expr: Index,
    },
    ValueToPlaceConversion {
        expr: Index,
    },
    /// data is two indices for `lhs` and `rhs`
    Or {
        lhs: Index,
        rhs: Index,
    },
    And {
        lhs: Index,
        rhs: Index,
    },
    BitOr {
        lhs: Index,
        rhs: Index,
    },
    BitXOr {
        lhs: Index,
        rhs: Index,
    },
    BitAnd {
        lhs: Index,
        rhs: Index,
    },
    Eq {
        lhs: Index,
        rhs: Index,
    },
    NEq {
        lhs: Index,
        rhs: Index,
    },
    Lt {
        lhs: Index,
        rhs: Index,
    },
    Gt {
        lhs: Index,
        rhs: Index,
    },
    Le {
        lhs: Index,
        rhs: Index,
    },
    Ge {
        lhs: Index,
        rhs: Index,
    },
    Shl {
        lhs: Index,
        rhs: Index,
    },
    Shr {
        lhs: Index,
        rhs: Index,
    },
    Add {
        lhs: Index,
        rhs: Index,
    },
    Sub {
        lhs: Index,
        rhs: Index,
    },
    Mul {
        lhs: Index,
        rhs: Index,
    },
    Div {
        lhs: Index,
        rhs: Index,
    },
    Rem {
        lhs: Index,
        rhs: Index,
    },
    Assign {
        lhs: Index,
        rhs: Index,
    },
    SubscriptExpr {
        lhs: Index,
        index: Index,
    },
    IfExpr {
        cond: Index,
        body: Index,
    },
    /// `data` is an index to an expression and an index into extra for [if, else]
    IfElseExpr {
        cond: Index,
        a: Index,
        b: Index,
    },
    // TODO:
    /// `data` is a ParseError
    Error {
        err: ParseError,
    },
    /// placeholder tag for reserved indices/nodes, `data` is none
    Undefined,
}