package typechecker

import (
  "fmt"
  "strings"

  "github.com/marzeq/quokka/parser"
  "github.com/marzeq/quokka/shared"
)

type (
  Node = parser.Node
  Type = shared.Type
)

type VarSig struct {
  Type    Type
  Mutable bool
}

type FunctionSig struct {
  Name           string
  ArgTypes       []shared.Pair[string, Type]
  HasVariadic    bool
  RetType        Type
  ImplicitReturn bool
}

type (
  VarTable  = *shared.SymbolTable[*VarSig]
  FuncTable = *shared.SymbolTable[*FunctionSig]
  TypeTable = shared.TypeTable
)

type TypeChecker struct {
  VarTable  VarTable
  FuncTable FuncTable
  TypeTable TypeTable
}

func NewTypeChecker() *TypeChecker {
  return &TypeChecker{
    VarTable:  shared.NewSymbolTable[*VarSig](),
    FuncTable: shared.NewSymbolTable[*FunctionSig](),
    TypeTable: shared.NewTypeTable(),
  }
}

func (tc *TypeChecker) TypeCheck(root *Node) (*Node, map[string]*FunctionSig, TypeTable, error) {
  fmt.Print()

  for _, node := range root.Children {
    switch node.Type {
    case parser.NODE_TYPE_FUNCTION_DEF:
      fsig, err := tc.ExtractFunctionSig(node)
      if err != nil {
        return nil, nil, nil, err
      }

      if ok := tc.FuncTable.Define(fsig.Name, fsig); !ok {
        return nil, nil, nil, shared.NewError(node.Loc, "function '%s' is already defined", fsig.Name)
      }

      if fsig.Name == "main" && fsig.RetType != shared.PRIMITIVE_I32 {
        return nil, nil, nil, shared.NewError(node.Loc, "'main' function must be of '%s' return type", shared.PRIMITIVE_I32)
      }
    case parser.NODE_TYPE_STRUCT_DEF:
      val := node.Value.(*parser.StructValue)
      sname := IdentToStr(node.Left)

      if _, ok := tc.TypeTable.Lookup(sname); ok {
        return nil, nil, nil, shared.NewError(node.Loc, "struct '%s' is already defined", sname)
      }

      fields := make([]shared.Pair[string, Type], len(val.Fields))

      for i, field := range val.Fields {
        tpe := IdentToStr(field.R)
        resolved, ok := tc.TypeTable.Lookup(tpe)
        if !ok {
          return nil, nil, nil, shared.NewError(field.R.Loc, "field '%s' has undefined type '%s'", IdentToStr(field.L), tpe)
        }
        fields[i] = shared.Pair[string, Type]{L: IdentToStr(field.L), R: resolved}
      }

      tc.TypeTable.Define(sname, shared.Struct{
        Fields: fields,
      })
    }
  }

  for _, node := range root.Children {
    if node.Type != parser.NODE_TYPE_FUNCTION_DEF {
      continue
    }

    name := IdentToStr(node.Value.(*parser.FunctionValue).Name)
    sig, _ := tc.FuncTable.Lookup(name)

    if len(node.Children) != 0 {
      tc.enterScope()
      if err := tc.typeCheckFunction(node, sig); err != nil {
        return nil, nil, nil, err
      }
      tc.exitScope()
    }
  }

  return root, tc.FuncTable.GetScope(), tc.TypeTable, nil
}

func (tc *TypeChecker) enterScope() {
  tc.VarTable.EnterScope()
  tc.FuncTable.EnterScope()
}

func (tc *TypeChecker) exitScope() {
  tc.VarTable.ExitScope()
  tc.FuncTable.ExitScope()
}

func (tc *TypeChecker) typeCheckFunction(funcNode *Node, sig *FunctionSig) error {
  val := funcNode.Value.(*parser.FunctionValue)
  body := funcNode.Children[0]
  name := IdentToStr(funcNode.Value.(*parser.FunctionValue).Name)

  for _, arg := range sig.ArgTypes {
    tc.VarTable.Define(arg.L, &VarSig{
      Type:    arg.R,
      Mutable: false,
    })
  }

  if body.Type == parser.NODE_TYPE_BLOCK {
    _, err := tc.typeCheckBlock(body, sig, false, true)
    return err
  }

  exprType, err := tc.typeCheckExpression(body, sig.RetType)
  if err != nil {
    return err
  }

  if !shared.CanCoerceTo(exprType, sig.RetType) {
    return shared.NewError(val.RetType.Loc, "function '%s' expects return type '%s' but returns '%s'",
      name, sig.RetType, exprType)
  }

  if exprType == shared.PRIMITIVE_UNTYPED_INT {
    body.ExprType = sig.RetType
  }
  return nil
}

func (tc *TypeChecker) typeCheckBlock(blockNode *Node, sig *FunctionSig, isLoop bool, isMainBody bool) (bool, error) {
  foundReturn := false

  for i, node := range blockNode.Children {
    switch node.Type {
    case parser.NODE_TYPE_DECLARATION:
      varName, varSig, err := tc.typeCheckDeclaration(node)
      if err != nil {
        return false, err
      }
      if ok := tc.VarTable.Define(varName, varSig); !ok {
        return false, shared.NewError(node.Left.Loc, "variable '%s' is already declared in this scope", varName)
      }

    case parser.NODE_TYPE_FUNCTION_DEF:
      fsig, err := tc.ExtractFunctionSig(node)
      if err != nil {
        return false, err
      }
      if ok := tc.FuncTable.Define(fsig.Name, fsig); !ok {
        return false, shared.NewError(node.Loc, "function '%s' is already defined", fsig.Name)
      }

      tc.enterScope()
      if err := tc.typeCheckFunction(node, fsig); err != nil {
        return false, err
      }
      tc.exitScope()

    case parser.NODE_TYPE_ASSIGNMENT:
      if err := tc.typeCheckAssignment(node); err != nil {
        return false, err
      }

    case parser.NODE_TYPE_FUNCTION_CALL:
      if _, err := tc.typeCheckFunctionCall(node); err != nil {
        return false, err
      }

    case parser.NODE_TYPE_CONTROL_KEYWORD:
      kw := node.Value.(string)
      switch kw {
      case "return":
        var retType Type = shared.PRIMITIVE_VOID
        if node.Right != nil {
          var err error
          retType, err = tc.typeCheckExpression(node.Right, sig.RetType)
          if err != nil {
            return false, err
          }

          if !shared.CanCoerceTo(retType, sig.RetType) {
            return false, shared.NewError(node.Loc, "wrong return type for function, expected '%s' got '%s'", sig.RetType, retType)
          }
          if retType == shared.PRIMITIVE_UNTYPED_INT {
            node.Right.ExprType = sig.RetType
          }
        } else if sig.RetType != shared.PRIMITIVE_VOID {
          return false, shared.NewError(node.Loc, "wrong return type for function, expected '%s' got void", sig.RetType)
        }
        foundReturn = true

        if i != len(blockNode.Children)-1 {
          return false, shared.NewError(blockNode.Children[i+1].Loc, "dead code following return statement")
        }

      case "break", "continue":
        if !isLoop {
          return false, shared.NewError(node.Loc, "'%s' statement outside a loop", kw)
        }
      }

    case parser.NODE_TYPE_BLOCK:
      tc.enterScope()
      returns, err := tc.typeCheckBlock(node, sig, isLoop, false)
      if err != nil {
        return false, err
      }
      tc.exitScope()

      if i == len(blockNode.Children)-1 && !foundReturn {
        foundReturn = returns
      }

    case parser.NODE_TYPE_FOR:
      if err := tc.typeCheckForLoop(node, sig); err != nil {
        return false, err
      }

    case parser.NODE_TYPE_IF:
      tc.enterScope()
      returns, err := tc.typeCheckIfStatement(node, sig, isLoop)
      if err != nil {
        return false, err
      }
      tc.exitScope()

      if i == len(blockNode.Children)-1 && !foundReturn {
        foundReturn = returns
      }

    default:
      panic("unexpected node in block")
    }
  }

  if isMainBody {
    if !foundReturn && sig.RetType != shared.PRIMITIVE_VOID {
      return false, shared.NewError(blockNode.Loc, "function with return type '%s' is missing a return statement", sig.RetType)
    }

    sig.ImplicitReturn = !foundReturn
  }

  return foundReturn, nil
}

func (tc *TypeChecker) typeCheckExpression(exprNode *Node, expectedType Type) (Type, error) {
  exprNode.ExprType = expectedType
  switch exprNode.Type {
  case parser.NODE_TYPE_IDENTIFIER:
    varName := exprNode.Value.(string)
    varSig, ok := tc.VarTable.Lookup(varName)
    if !ok {
      return shared.PRIMITIVE_VOID, shared.NewError(exprNode.Loc, "undefined variable '%s'", varName)
    }
    exprNode.ExprType = varSig.Type
    return varSig.Type, nil

  case parser.NODE_TYPE_NUMBER_LITERAL:
    exprType := shared.PRIMITIVE_UNTYPED_INT
    if expectedType != shared.PRIMITIVE_VOID && shared.CanCoerceTo(exprType, expectedType) {
      exprNode.ExprType = expectedType
      return expectedType, nil
    }
    exprNode.ExprType = exprType
    return exprType, nil

  case parser.NODE_TYPE_CHAR_LITERAL:
    exprNode.ExprType = shared.PRIMITIVE_CHAR
    return shared.PRIMITIVE_CHAR, nil

  case parser.NODE_TYPE_STRING_LITERAL:
    exprNode.ExprType = shared.PRIMITIVE_CSTRING
    return shared.PRIMITIVE_CSTRING, nil

  case parser.NODE_TYPE_BOOL_LITERAL:
    exprNode.ExprType = shared.PRIMITIVE_BOOL
    return shared.PRIMITIVE_BOOL, nil

  case parser.NODE_TYPE_FUNCTION_CALL:
    return tc.typeCheckFunctionCall(exprNode)

  case parser.NODE_TYPE_UNARY_OP:
    op := exprNode.Value.(string)
    operandType, err := tc.typeCheckExpression(exprNode.Right, shared.PRIMITIVE_VOID)
    if err != nil {
      return shared.PRIMITIVE_VOID, err
    }

    if expectedType != shared.PRIMITIVE_VOID && operandType == shared.PRIMITIVE_UNTYPED_INT && shared.IsNumericType(expectedType) {
      exprNode.Right.ExprType = expectedType
      operandType = expectedType
    }

    switch op {
    case "not":
      if operandType != shared.PRIMITIVE_BOOL {
        return shared.PRIMITIVE_VOID, shared.NewError(exprNode.Right.Loc, "unary operator 'not' expects a boolean operand, found '%s'", operandType)
      }
			exprNode.ExprType = shared.PRIMITIVE_BOOL
      return shared.PRIMITIVE_BOOL, nil
    case "-":
      if !shared.IsNumericType(operandType) {
        return shared.PRIMITIVE_VOID, shared.NewError(exprNode.Right.Loc, "unary operator '-' requires numeric operand, found '%s'", operandType)
      }
			exprNode.ExprType = operandType
      return operandType, nil
    case "*":
      if !operandType.IsPointer() {
        return shared.PRIMITIVE_VOID, shared.NewError(exprNode.Right.Loc, "unary operator '*' requires a pointer operand, found '%s'", operandType)
      }
			t := operandType.(shared.Pointer).To
			exprNode.ExprType = t
      return t, nil
    case "&":
			if exprNode.Right.Type != parser.NODE_TYPE_IDENTIFIER {
				return shared.PRIMITIVE_VOID, shared.NewError(exprNode.Right.Loc, "unary operator '&' requires an identifier operand, found '%s'", exprNode.Right.Type)
			}
			varName := IdentToStr(exprNode.Right)
			varSig, ok := tc.VarTable.Lookup(varName)
			if !ok {
				return shared.PRIMITIVE_VOID, shared.NewError(exprNode.Right.Loc, "undefined variable '%s'", varName)
			}
			t := shared.Pointer{
				To: operandType,
				Const: !varSig.Mutable,
			}
			exprNode.ExprType = t
      return t, nil
    default:
      return shared.PRIMITIVE_VOID, shared.NewError(exprNode.Loc, "unknown unary operator '%s'", op)
    }

  case parser.NODE_TYPE_BINARY_OP:
    op := exprNode.Value.(string)
    leftType, err := tc.typeCheckExpression(exprNode.Left, shared.PRIMITIVE_VOID)
    if err != nil {
      return shared.PRIMITIVE_VOID, err
    }
    rightType, err := tc.typeCheckExpression(exprNode.Right, shared.PRIMITIVE_VOID)
    if err != nil {
      return shared.PRIMITIVE_VOID, err
    }

    if expectedType != shared.PRIMITIVE_VOID {
      if leftType == shared.PRIMITIVE_UNTYPED_INT && shared.IsNumericType(expectedType) {
        exprNode.Left.ExprType = expectedType
        leftType = expectedType
      }
      if rightType == shared.PRIMITIVE_UNTYPED_INT && shared.IsNumericType(expectedType) {
        exprNode.Right.ExprType = expectedType
        rightType = expectedType
      }
    }

    switch op {
    case "and", "or":
      if leftType != shared.PRIMITIVE_BOOL || rightType != shared.PRIMITIVE_BOOL {
        return shared.PRIMITIVE_VOID, shared.NewError(exprNode.Loc, "operator '%s' expects boolean operands, found '%s' and '%s'", op, leftType, rightType)
      }
      exprNode.ExprType = shared.PRIMITIVE_BOOL
      return shared.PRIMITIVE_BOOL, nil

    case "==", "!=", "<", ">", "<=", ">=":
      if !shared.IsNumericType(leftType) || !shared.IsNumericType(rightType) {
        return shared.PRIMITIVE_VOID, shared.NewError(exprNode.Loc, "operator '%s' cannot be applied to operands of type '%s' and '%s'", op, leftType, rightType)
      }
      exprNode.ExprType = shared.PRIMITIVE_BOOL
      return shared.PRIMITIVE_BOOL, nil

    case "+", "-", "*", "/", "%":
      if !shared.IsNumericType(leftType) || !shared.IsNumericType(rightType) {
        return shared.PRIMITIVE_VOID, shared.NewError(exprNode.Loc, "operator '%s' requires numeric operands, found '%s' and '%s'", op, leftType, rightType)
      }
      commonType := shared.BiggerNumericType(leftType, rightType)
      exprNode.Left.ExprType = commonType
      exprNode.Right.ExprType = commonType
      exprNode.ExprType = commonType
      return commonType, nil

    default:
      return shared.PRIMITIVE_VOID, shared.NewError(exprNode.Loc, "unknown binary operator '%s'", op)
    }

  case parser.NODE_TYPE_CAST:
    tpeName := IdentToStr(exprNode.Left)
    tpe, ok := tc.TypeTable.Lookup(tpeName)
    if !ok {
      return shared.PRIMITIVE_VOID, shared.NewError(exprNode.Left.Loc, "no such type '%s'", tpeName)
    }
    exTpe, err := tc.typeCheckExpression(exprNode.Right, shared.PRIMITIVE_VOID)
    if err != nil {
      return shared.PRIMITIVE_VOID, err
    }

    if exTpe == shared.PRIMITIVE_UNTYPED_INT {
      exprNode.Right.ExprType = tpe
    }

    if shared.CanCastTo(exTpe, tpe) {
      exprNode.ExprType = tpe
      return tpe, nil
    } else {
      return shared.PRIMITIVE_VOID, shared.NewError(exprNode.Loc, "cannot cast type '%s' to '%s'", exTpe, tpe)
    }

  case parser.NODE_TYPE_IF_EXPR:
    ifExpr := exprNode.Value.(*parser.IfNodeValue)
    if ifExpr.IfBranch == nil {
      return shared.PRIMITIVE_VOID, shared.NewError(exprNode.Loc, "if expression missing if branch")
    }

    condType, err := tc.typeCheckExpression(ifExpr.IfBranch.Condition, shared.PRIMITIVE_BOOL)
    if err != nil {
      return shared.PRIMITIVE_VOID, err
    }
    if condType != shared.PRIMITIVE_BOOL {
      return shared.PRIMITIVE_VOID, shared.NewError(ifExpr.IfBranch.Condition.Loc, "if condition must be boolean, found '%s'", condType)
    }

    ifType, err := tc.typeCheckExpression(ifExpr.IfBranch.Node, expectedType)
    if err != nil {
      return shared.PRIMITIVE_VOID, err
    }
    commonType := ifType

    for _, elseIf := range ifExpr.ElseIfBranches {
      elseIfCondType, err := tc.typeCheckExpression(elseIf.Condition, shared.PRIMITIVE_BOOL)
      if err != nil {
        return shared.PRIMITIVE_VOID, err
      }
      if elseIfCondType != shared.PRIMITIVE_BOOL {
        return shared.PRIMITIVE_VOID, shared.NewError(elseIf.Condition.Loc, "else-if condition must be boolean, found '%s'", elseIfCondType)
      }

      elseIfType, err := tc.typeCheckExpression(elseIf.Node, expectedType)
      if err != nil {
        return shared.PRIMITIVE_VOID, err
      }
      if elseIfType != commonType {
        return shared.PRIMITIVE_VOID, shared.NewError(exprNode.Loc, "all branches must return same type, expected '%s' but found '%s'", commonType, elseIfType)
      }
    }

    if ifExpr.ElseBranch == nil {
      return shared.PRIMITIVE_VOID, shared.NewError(exprNode.Loc, "if expression requires else branch")
    }
    elseType, err := tc.typeCheckExpression(ifExpr.ElseBranch.Node, expectedType)
    if err != nil {
      return shared.PRIMITIVE_VOID, err
    }
    if elseType != commonType {
      return shared.PRIMITIVE_VOID, shared.NewError(ifExpr.ElseBranch.Node.Loc, "else branch must match type '%s', found '%s'", commonType, elseType)
    }

    exprNode.ExprType = commonType
    return commonType, nil

  default:
    return shared.PRIMITIVE_VOID, shared.NewError(exprNode.Loc, "unsupported expression type: '%s'", exprNode.Type)
  }
}

func (tc *TypeChecker) typeCheckFunctionCall(funccallNode *Node) (Type, error) {
  nameNode := funccallNode.Value.(*Node)
  fname := IdentToStr(nameNode)

  fsig, ok := tc.FuncTable.Lookup(fname)
  if !ok {
    return shared.PRIMITIVE_VOID, shared.NewError(nameNode.Loc, "undefined function '%s'", fname)
  }

  if (len(funccallNode.Children) > len(fsig.ArgTypes) && !fsig.HasVariadic) || len(funccallNode.Children) < len(fsig.ArgTypes) {
    return shared.PRIMITIVE_VOID, shared.NewError(nameNode.Loc, "argument count mismatch, expected at least %d, got %d", len(fsig.ArgTypes), len(funccallNode.Children))
  }

  for i, arg := range funccallNode.Children {
    var fsigArgType Type = shared.PRIMITIVE_VOID
    if i < len(fsig.ArgTypes) {
      fsigArgType = fsig.ArgTypes[i].R
    }
    argType, err := tc.typeCheckExpression(arg, fsigArgType)
    if err != nil {
      return shared.PRIMITIVE_VOID, err
    }
    if fsigArgType != argType && fsigArgType != shared.PRIMITIVE_VOID {
      return shared.PRIMITIVE_VOID, shared.NewError(arg.Loc,
        "argument %d of function '%s' has type '%s' but expected '%s'",
        i+1, fname, argType, fsigArgType,
      )
    }
  }

  funccallNode.ExprType = fsig.RetType
  return fsig.RetType, nil
}

func (tc *TypeChecker) typeCheckDeclaration(declNode *Node) (string, *VarSig, error) {
  varName := IdentToStr(declNode.Left)
  if strings.HasPrefix(varName, "___") {
    return "", nil, shared.NewError(declNode.Loc, "variables starting with '___' are reserved for the compiler")
  }
  info := declNode.Value.(*parser.DeclarationValue)
  mutable := info.Mutable
  var varType Type = shared.PRIMITIVE_VOID

  if info.Type != nil {
    varTypeStr := IdentToStr(info.Type)
    vt, ok := tc.TypeTable.Lookup(varTypeStr)
    if !ok {
      return "", nil, shared.NewError(info.Type.Loc, "variable '%s' has undefined type '%s'", varName, varTypeStr)
    }
    varType = vt
  }

  exprType, err := tc.typeCheckExpression(declNode.Right, varType)
  if err != nil {
    return "", nil, err
  }

  if varType == shared.PRIMITIVE_VOID {
    varType = exprType
  } else if !shared.CanCoerceTo(exprType, varType) {
    return "", nil, shared.NewError(declNode.Loc,
      "cannot assign value of type '%s' to variable '%s' of type '%s'",
      exprType, varName, varType,
    )
  }

  if varType == shared.PRIMITIVE_UNTYPED_INT {
		varType = shared.PRIMITIVE_I32
		exprType = shared.PRIMITIVE_I32
  }
  if varType == shared.PRIMITIVE_VOID {
    return "", nil, shared.NewError(declNode.Loc, "a variable cannot be of type void")
  }

  declNode.ExprType = varType
	declNode.Right.ExprType = exprType

  return varName, &VarSig{
    Type:    varType,
    Mutable: mutable,
  }, nil
}

func (tc *TypeChecker) typeCheckAssignment(asNode *Node) error {
  varName := IdentToStr(asNode.Left)
  varSig, ok := tc.VarTable.Lookup(varName)
  if !ok {
    return shared.NewError(asNode.Left.Loc, "undefined variable '%s'", varName)
  }

  if !varSig.Mutable {
    return shared.NewError(asNode.Loc, "cannot assign to immutable variable '%s'", varName)
  }

  exprType, err := tc.typeCheckExpression(asNode.Right, varSig.Type)
  if err != nil {
    return err
  }

  if exprType != varSig.Type {
    return shared.NewError(asNode.Loc,
      "cannot assign value of type '%s' to variable '%s' of type '%s'",
      exprType, varName, varSig.Type,
    )
  }

  asNode.ExprType = varSig.Type

  return nil
}

func (tc *TypeChecker) typeCheckForLoop(loopNode *Node, sig *FunctionSig) error {
  value := loopNode.Value.(*parser.ForLoopValue)

  tc.enterScope()

  if len(value.ExprsOrStmts) == 1 {
    expr := value.ExprsOrStmts[0]
    if !expr.Type.IsExpression() {
      return shared.NewError(expr.Loc, "for loop condition must be a boolean expression")
    }
    exprType, err := tc.typeCheckExpression(expr, shared.PRIMITIVE_BOOL)
    if err != nil {
      return err
    }

    if exprType != shared.PRIMITIVE_BOOL {
      return shared.NewError(expr.Loc, "for loop condition must be a boolean expression, found '%s'", exprType)
    }
  } else if len(value.ExprsOrStmts) == 3 {
    init := value.ExprsOrStmts[0]
    switch init.Type {
    case parser.NODE_TYPE_DECLARATION:
      varName, varSig, err := tc.typeCheckDeclaration(init)
      if err != nil {
        return err
      }

      tc.VarTable.Define(varName, varSig)
    case parser.NODE_TYPE_ASSIGNMENT:
      if err := tc.typeCheckAssignment(init); err != nil {
        return err
      }
    default:
      return shared.NewError(init.Loc, "for loop initialiser must be a declaration or assignment statement")
    }

    cond := value.ExprsOrStmts[1]
    if !cond.Type.IsExpression() {
      return shared.NewError(cond.Loc, "for loop condition must be a boolean expression")
    }
    exprType, err := tc.typeCheckExpression(cond, shared.PRIMITIVE_BOOL)
    if err != nil {
      return err
    }

    if exprType != shared.PRIMITIVE_BOOL {
      return shared.NewError(cond.Loc, "for loop condition must be a boolean expression, found '%s'", exprType)
    }

    reass := value.ExprsOrStmts[2]
    if reass.Type != parser.NODE_TYPE_ASSIGNMENT {
      return shared.NewError(reass.Loc, "for loop 'after' step must be an assignment statement")
    }

    if err := tc.typeCheckAssignment(reass); err != nil {
      return err
    }
  } else if len(value.ExprsOrStmts) != 0 {
    return shared.NewError(loopNode.Loc, "for loop must have either:\n  - no conditions\n  - a condition\n  - an initialiser, a condition and an 'after' assignment")
  }

  if _, err := tc.typeCheckBlock(value.Body, sig, true, false); err != nil {
    return err
  }

  tc.exitScope()
  return nil
}

func (tc *TypeChecker) typeCheckIfStatement(ifNode *Node, sig *FunctionSig, isLoop bool) (bool, error) {
  value := ifNode.Value.(*parser.IfNodeValue)

  ifCondType, err := tc.typeCheckExpression(value.IfBranch.Condition, shared.PRIMITIVE_BOOL)
  if err != nil {
    return false, err
  }
  if ifCondType != shared.PRIMITIVE_BOOL {
    return false, shared.NewError(value.ElseBranch.Condition.Loc, "if condition must be boolean, found '%s'", ifCondType)
  }

  ifReturns, err := tc.typeCheckBlock(value.IfBranch.Node, sig, isLoop, false)
  if err != nil {
    return false, err
  }
  alwaysReturns := ifReturns && value.ElseBranch != nil

  for _, elseIfBranch := range value.ElseIfBranches {
    elseIfCondType, err := tc.typeCheckExpression(elseIfBranch.Condition, shared.PRIMITIVE_BOOL)
    if err != nil {
      return false, err
    }
    if elseIfCondType != shared.PRIMITIVE_BOOL {
      return false, shared.NewError(elseIfBranch.Condition.Loc, "else-if condition must be boolean, found '%s'", elseIfCondType)
    }

    elseIfReturns, err := tc.typeCheckBlock(elseIfBranch.Node, sig, isLoop, false)
    if err != nil {
      return false, err
    }
    if !alwaysReturns {
      alwaysReturns = elseIfReturns
    }
  }

  if value.ElseBranch != nil {
    elseReturns, err := tc.typeCheckBlock(value.ElseBranch.Node, sig, isLoop, false)
    if err != nil {
      return false, err
    }
    if !alwaysReturns {
      alwaysReturns = elseReturns
    }
  }

  return alwaysReturns, nil
}

func (tc *TypeChecker) ExtractFunctionSig(functionNode *Node) (*FunctionSig, error) {
  functionVal := functionNode.Value.(*parser.FunctionValue)
  name := IdentToStr(functionVal.Name)
	var retType Type
	if (functionVal.RetType != nil) {
		retTypeStr := IdentToStr(functionVal.RetType)

		r, ok := tc.TypeTable.Lookup(retTypeStr)
		if !ok {
			return nil, shared.NewError(functionVal.RetType.Loc, "function '%s' has undefined return type '%s'", name, retTypeStr)
		}
		retType = r
	} else {
		if name == "main" {
			retType = shared.PRIMITIVE_I32
		} else {
			retType = shared.PRIMITIVE_VOID
		}
	}

  argTypes := make([]shared.Pair[string, Type], len(functionVal.Args))
  for i, arg := range functionVal.Args {
    tpe := IdentToStr(arg.R)
    resolved, ok := tc.TypeTable.Lookup(tpe)
    if !ok {
      return nil, shared.NewError(arg.R.Loc, "parameter '%s' has undefined type '%s'", IdentToStr(arg.L), tpe)
    }
    argTypes[i] = shared.Pair[string, Type]{L: IdentToStr(arg.L), R: resolved}
  }

  return &FunctionSig{
    ArgTypes:    argTypes,
    RetType:     retType,
    Name:        name,
    HasVariadic: functionVal.HasVariadic,
  }, nil
}

func IdentToStr(identNode *Node) string {
  return identNode.Value.(string)
}