//-----------------------------------------------------------------------------
// AST
//-----------------------------------------------------------------------------
enum Ast_Kind: U32{
  AST_NONE,

  AST_PACKAGE,

  AST_VALUE,
  AST_CAST,
  AST_IDENT,
  AST_INDEX,
  AST_UNARY,
  AST_BINARY,
  AST_CALL_ITEM,
  AST_CALL,

  AST_POINTER,
  AST_ARRAY,
  AST_FOR,
  AST_IF,
  AST_IF_NODE,
  AST_RETURN,
  AST_BLOCK,
  AST_PASS,
  AST_LAMBDA,
  AST_LAMBDA_ARG,
  AST_ENUM,
  AST_ENUM_MEMBER,
  AST_STRUCT,
  AST_CONST,
  AST_VAR,
};

typedef U32 Ast_Flag;
enum{
  AST_EXPR = 1,
  AST_STMT = 2,
  AST_BINDING = 4,
  AST_AGGREGATE = 8,
  AST_AGGREGATE_CHILD = 16,
  AST_ITEM_INCLUDED = 32,
  AST_ATOM = 64,
  AST_FOREIGN = 128,
};

struct Ast{
  U64 id;
  Token *pos;

  Ast_Kind kind;
  Ast *parent;
  Ast_Flag flags;
};

struct Ast_Resolved_Type;
struct Ast_Expr:Ast{};

#define VALUE_FIELDS             \
Ast_Resolved_Type *type;         \
union{                           \
  bool               bool_val;   \
  F64                f64_val;    \
  Intern_String      intern_val; \
  BigInt             big_int_val;\
  Ast_Resolved_Type *type_val;   \
};
#define INLINE_VALUE_FIELDS union{Value value; struct{VALUE_FIELDS};}
struct Value{VALUE_FIELDS};
  // BigInt             big_int_val;

struct Ast_Atom: Ast_Expr{
  INLINE_VALUE_FIELDS;
};

struct Ast_Call_Item: Ast_Expr{
  Ast_Atom *name; // index | name
  Ast_Expr *index;
  Ast_Expr *item;
};

struct Ast_Call: Ast_Expr{
  Ast_Resolved_Type *type; // @todo: to map
  Ast_Expr *name;
  Array<Ast_Call_Item *> exprs;
};

struct Ast_Unary: Ast_Expr{
  Token_Kind op;
  Ast_Expr *expr;
  U64 padding[3]; // For folding constants into atoms
};

struct Ast_Cast: Ast_Expr{
  Ast_Expr *expr;
  Ast_Expr *typespec;
};

struct Ast_Index: Ast_Expr{
  Ast_Expr *expr;
  Ast_Expr *index;
};

struct Ast_Binary: Ast_Expr{
  Token_Kind op;
  Ast_Expr *left;
  Ast_Expr *right;
};

struct Ast_Block : Ast {
  Array<Ast *> stmts;
};

struct Ast_Return: Ast{
  Ast_Expr *expr;
};

struct Ast_If_Node: Ast{
  Ast_Expr  *expr ;
  Ast_Block *block;
  Ast_Binary*init;
};

struct Ast_If: Ast{
  Array<Ast_If_Node *> ifs;
};

struct Ast_Pass: Ast{}; // @todo

struct Ast_For: Ast{
  Ast_Expr *init;
  Ast_Expr *cond;
  Ast_Expr *iter;
  Ast_Block *block;
};

struct Ast_Lambda_Arg: Ast_Expr{
  Intern_String name;
  Ast_Expr *typespec;
  Ast_Expr *default_value;
};

struct Ast_Lambda : Ast_Expr {
  Array<Ast_Lambda_Arg *> args;
  Ast_Expr               *ret;
  Ast_Block              *block;
};

struct Ast_Array: Ast_Expr{
  Ast_Expr *base;
  Ast_Expr *expr;
};

struct Ast_Named:Ast{
  Intern_String name;
};

struct Ast_Var: Ast_Named{
  Ast_Expr *typespec;
  Ast_Expr *expr;
};

struct Ast_Const;
struct Ast_Resolved_Type;
struct Ast_Struct: Ast{
  // Required to be Ast_Struct or Ast_Var or Ast_Const
  Array<Ast_Var *> members;
  Array<Ast_Const *> const_members;
  Ast_Resolved_Type *type;
};

struct Ast_Enum_Member: Ast{
  Intern_String name;
  Ast_Expr *value;
};

struct Ast_Enum: Ast{
  Ast_Expr *typespec;
  Array<Ast_Enum_Member *> members;
};

struct Ast_Const: Ast_Named{
  union{
    Ast *ast;
    Ast_Expr *value;
    Ast_Struct *agg;
    Ast_Enum *enu;
  };
};

struct Ast_Package:Ast{
  Intern_String name;
  Array<Ast_Named *> decls;
  Array<Ast_Named *> ordered;
};

//-----------------------------------------------------------------------------
// AST Constructors beginning with expressions
//-----------------------------------------------------------------------------
#define AST_NEW(T,ikind,ipos,iflags)                                   \
  Ast_##T *result = exp_alloc_type(pctx->perm, Ast_##T, AF_ZeroMemory);\
  result->flags   = iflags;                                            \
  result->kind    = AST_##ikind;                                       \
  result->pos     = ipos;                                              \
  result->id      = ++pctx->unique_ids

function Ast_Atom *
ast_str(Token *pos, Intern_String string){
  AST_NEW(Atom, VALUE, pos, AST_EXPR | AST_ATOM);
  result->type = untyped_string;
  result->intern_val = string;
  return result;
}

function Ast_Atom *
ast_ident(Token *pos, Intern_String string){
  AST_NEW(Atom, IDENT, pos, AST_EXPR | AST_ATOM);
  result->intern_val = string;
  return result;
}

function Ast_Atom *
ast_bool(Token *pos, B32 bool_val){
  AST_NEW(Atom, VALUE, pos, AST_EXPR | AST_ATOM);
  result->bool_val = bool_val;
  result->type     = untyped_bool;
  return result;
}

function Ast_Atom *
ast_float(Token *pos, F64 value){
  AST_NEW(Atom, VALUE, pos, AST_EXPR | AST_ATOM);
  result->type = untyped_float;
  result->f64_val = value;
  return result;
}

function Ast_Atom *
ast_int(Token *pos, BigInt val){
  AST_NEW(Atom, VALUE, pos, AST_EXPR | AST_ATOM);
  result->type = untyped_int;
  result->big_int_val = bigint_copy(pctx->perm, &val);
  return result;
}

function Ast_Atom *
ast_int(Token *pos, U64 value){
  return ast_int(pos, bigint_u64(value));
}

function Ast_Expr *
ast_expr_binary(Ast_Expr *left, Ast_Expr *right, Token *op){
  AST_NEW(Binary, BINARY, op, AST_EXPR);
  result->op    = op->kind;
  result->left  = left;
  result->right = right;
  result->left->parent = result;
  if(result->right) result->right->parent = result;
  return result;
}

function Ast_Call *
ast_call(Token *pos, Ast_Expr *name, Array<Ast_Call_Item *> exprs){
  AST_NEW(Call, CALL, pos, AST_EXPR);
  result->name = name;
  result->exprs = exprs.tight_copy(pctx->perm);
  if(result->name) result->name->parent = result;
  For(result->exprs) it->parent = result;
  return result;
}

function Ast_Call_Item *
ast_call_item(Token *pos, Ast_Expr *index, Ast_Atom *name, Ast_Expr *item){
  AST_NEW(Call_Item, CALL_ITEM, pos, AST_EXPR);
  result->name  = name;
  result->index = index;
  result->item  = item;
  if(result->name)  result->name->parent = result;
  if(result->index) result->index->parent = result;
  item->parent = result;
  return result;
}

function Ast_Expr *
ast_expr_cast(Token *pos, Ast_Expr *expr, Ast_Expr *typespec){
  AST_NEW(Cast, CAST, pos, AST_EXPR);
  result->flags = AST_EXPR;
  result->expr     = expr;
  result->typespec = typespec;
  expr->parent = result;
  typespec->parent = result;
  return result;
}

function Ast_Expr *
ast_expr_unary(Token *pos, Token_Kind op, Ast_Expr *expr){
  AST_NEW(Unary, UNARY, pos, AST_EXPR);
  result->flags = AST_EXPR;
  result->expr = expr;
  result->op   = op;
  expr->parent = result;
  return result;
}

function Ast_Expr *
ast_expr_index(Token *pos, Ast_Expr *expr, Ast_Expr *index){
  AST_NEW(Index, INDEX, pos, AST_EXPR);
  result->flags = AST_EXPR;
  result->expr = expr;
  result->index = index;
  expr->parent = result;
  index->parent = result;
  return result;
}

function Ast_Lambda *
ast_lambda(Token *pos, Array<Ast_Lambda_Arg *> params, Ast_Expr *ret, Ast_Block *block){
  AST_NEW(Lambda, LAMBDA, pos, AST_EXPR);
  result->flags = AST_EXPR;
  result->args  = params.tight_copy(pctx->perm);
  result->block = block;
  result->ret   = ret;
  if(!ret) result->ret = ast_ident(result->pos, intern_void);

  if(result->block) result->block->parent = result;
  result->ret->parent = result;
  For(result->args) it->parent = result;
  return result;
}

function Ast_Lambda_Arg *
ast_expr_lambda_arg(Token *pos, Intern_String name, Ast_Expr *typespec, Ast_Expr *default_value){
  AST_NEW(Lambda_Arg, LAMBDA_ARG, pos, AST_EXPR);
  result->flags            = AST_EXPR;
  result->name             = name;
  result->typespec         = typespec;
  result->default_value    = default_value;
  result->typespec->parent = result;
  if(result->default_value) result->default_value->parent = result;
  return result;
}

function Ast_Block *
ast_block(Token *pos, Array<Ast *> stmts){
  AST_NEW(Block, BLOCK, pos, AST_STMT);
  result->stmts = stmts.tight_copy(pctx->perm);
  For(result->stmts) it->parent = result;
  return result;
}

function Ast_If *
ast_if(Token *pos, Array<Ast_If_Node *> ifs){
  AST_NEW(If, IF, pos, AST_STMT);
  result->ifs = ifs.tight_copy(pctx->perm);
  For(result->ifs) it->parent = result;
  return result;
}

function Ast_For *
ast_for(Token *pos, Ast_Expr *init, Ast_Expr *cond, Ast_Expr *iter, Ast_Block *block){
  AST_NEW(For, FOR, pos, AST_STMT);
  result->init = init;
  result->cond = cond;
  result->iter = iter;
  result->block = block;
  if(result->init) result->init->parent = result;
  if(result->cond) result->cond->parent = result;
  if(result->iter) result->iter->parent = result;
  result->block->parent = result;
  return result;
}

function Ast_Pass *
ast_pass(Token *pos){
  AST_NEW(Pass, PASS, pos, AST_STMT);
  return result;
}

function Ast_Return *
ast_return(Token *pos, Ast_Expr *expr){
  AST_NEW(Return, RETURN, pos, AST_STMT);
  if(expr){
    assert(is_flag_set(expr->flags, AST_EXPR));
    result->expr = expr;
    result->expr->parent = result;
  }
  return result;
}

function Ast_If_Node *
ast_if_node(Token *pos, Ast_Expr *init, Ast_Expr *expr, Ast_Block *block){
  AST_NEW(If_Node, IF_NODE, pos, AST_STMT);
  result->block    = block;
  result->expr     = expr;
  result->init     = (Ast_Binary *)init;
  if(result->block) result->block->parent = result;
  if(result->expr) result->expr->parent = result;
  if(result->init) {
    assert(init->kind == AST_BINARY);
    result->init->parent = result;
  }
  return result;
}

function Ast_Array *
ast_array(Token *pos, Ast_Expr *expr){
  AST_NEW(Array, ARRAY, pos, AST_EXPR);
  result->expr = expr;
  result->expr->parent = result;
  return result;
}

function Ast_Enum_Member *
ast_enum_member(Token *pos, Intern_String name, Ast_Expr *default_value){
  AST_NEW(Enum_Member, ENUM_MEMBER, pos, AST_AGGREGATE_CHILD);
  result->name = name;
  result->value = default_value;
  if(result->value) result->value->parent = result;
  return result;
}

function Ast_Enum *
ast_enum(Token *pos, Ast_Expr *typespec, Array<Ast_Enum_Member *> members){
  AST_NEW(Enum, ENUM, pos, AST_AGGREGATE);
  result->members = members.tight_copy(pctx->perm);
  result->typespec = typespec;
  if(result->typespec) result->typespec->parent = result;
  For(result->members){
    it->parent = result;
  }
  return result;
}

function Ast_Struct *
ast_struct(Token *pos, Array<Ast_Var *> members, Array<Ast_Const *> const_members){
  AST_NEW(Struct, STRUCT, pos, AST_AGGREGATE);
  result->members = members.tight_copy(pctx->perm);
  result->const_members = const_members.tight_copy(pctx->perm);
  For(result->members) {
    assert(is_flag_set(it->flags, AST_BINDING));
    assert(it->kind == AST_VAR);
    it->parent = result;
  }
  For(result->const_members) {
    assert(is_flag_set(it->flags, AST_BINDING));
    assert(it->kind == AST_CONST);
    it->parent = result;
  }
  return result;
}

//-----------------------------------------------------------------------------
// Declarations
//-----------------------------------------------------------------------------
function Ast_Var *
ast_var(Token *pos, Ast_Expr *typespec, Intern_String name, Ast_Expr *expr){
  AST_NEW(Var, VAR, pos, AST_BINDING);
  result->expr     = expr;
  result->typespec = typespec;
  result->name     = name;
  if(result->expr) result->expr->parent = result;
  if(result->typespec) result->typespec->parent = result;
  return result;
}

function Ast_Const *
ast_const(Token *pos, Intern_String name, Ast_Expr *value){
  assert(is_flag_set(value->flags, AST_AGGREGATE) || is_flag_set(value->flags, AST_EXPR) );
  AST_NEW(Const, CONST, pos, AST_BINDING);
  result->value = value;
  result->name = name;
  result->value->parent = result;
  return result;
}

function Ast_Package *
ast_package(Token *pos, String name, Array<Ast_Named *> decls){
  AST_NEW(Package, PACKAGE, pos, 0);
  result->decls   = decls.tight_copy(pctx->perm);
  result->ordered = array_make<Ast_Named *>(pctx->perm, decls.len);
  result->name    = intern_string(&pctx->interns, name);
  For(result->decls) it->parent = result;
  return result;
}

//-----------------------------------------------------------------------------
// Value
//-----------------------------------------------------------------------------
function Value
value_bool(B32 v){
  Value value;
  value.bool_val = v;
  value.type = untyped_bool;
  return value;
}

function Value
value_int(BigInt b){
  Value value;
  value.big_int_val = b;
  value.type = untyped_int;
  return value;
}

function Value
value_int(S64 s64){
  Value value;
  value.type = untyped_int;
  bigint_init_signed(&value.big_int_val, s64);
  return value;
}

function Value
value_float(F64 b){
  Value value;
  value.f64_val = b;
  value.type = untyped_float;
  return value;
}

function Value
value_float(BigInt a){
  Value value;
  value.f64_val = bigint_as_float(&a);
  value.type = untyped_float;
  return value;
}


//-----------------------------------------------------------------------------
// Utillities
//-----------------------------------------------------------------------------
function Ast_Struct *
const_try_getting_struct(Ast *ast){
  assert(ast->kind == AST_CONST);
  Ast_Const *constant = (Ast_Const *)ast;
  if(constant->value->kind == AST_STRUCT){
    return (Ast_Struct *)constant->value;
  }
  return 0;
}

function Ast_Struct *
const_get_struct(Ast *ast){
  auto result = const_try_getting_struct(ast);
  assert(result);
  return result;
}

function Ast_Lambda *
const_try_getting_lambda(Ast *ast){
  assert(ast->kind == AST_CONST);
  Ast_Const *constant = (Ast_Const *)ast;
  if(constant->value->kind == AST_LAMBDA){
    return (Ast_Lambda *)constant->value;
  }
  return 0;
}

function Ast_Lambda *
const_get_lambda(Ast *ast){
  auto result = const_try_getting_lambda(ast);
  assert(result);
  return result;
}

function Intern_String
ast_get_name(Ast *ast){
  assert(is_flag_set(ast->flags, AST_BINDING));
  auto constant = (Ast_Named *)ast;
  return constant->name;
}

function B32
ast_is_struct(Ast *ast){
  if(ast->kind == AST_CONST){
    auto a = (Ast_Const *)ast;
    B32 result = a->agg->kind == AST_STRUCT;
    return result;
  }
  return false;
}

function B32
is_ident(Ast *ast){
  B32 result = ast->kind == AST_IDENT;
  return result;
}

function B32
is_binary(Ast *ast){
  B32 result = ast->kind == AST_BINARY;
  return result;
}

function B32
is_atom(Ast *ast){
  B32 result = is_flag_set(ast->flags, AST_ATOM);
  return result;
}

function Ast *
query_struct(Ast_Struct *agg, Intern_String string){
  For(agg->members){
    if(it->name == string){
      return it;
    }
  }
  For(agg->const_members){
    if(it->name == string){
      return it;
    }
  }
  return 0;
}

function Ast_Enum_Member *
query_enum(Ast_Enum *enu, Intern_String string){
  For(enu->members){
    if(it->name == string){
      return it;
    }
  }
  return 0;
}