module main

import term

// ------------------------------------------- Symbol Table

type SymbolInfo = VarSymbolInfo | FuncSymbolInfo

struct VarSymbolInfo {
	type string
}

struct FuncSymbolInfo {
	type string
	block Block
}

struct SymbolTable {
mut:
	variable_scopes []map[string]VarSymbolInfo
	functions map[string]FuncSymbolInfo
}

fn (mut s SymbolTable) define_var(name string, typ string) {
  $if debug {
		dump(s.variable_scopes.len)
	}

	if s.variable_scopes.len == 0 {
    parse_error('No scope available')
  }

  if name in s.variable_scopes[s.variable_scopes.len-1] {
      parse_error('Variable ${name} already defined in this scope')
  }

  s.variable_scopes[s.variable_scopes.len-1][name] = VarSymbolInfo{
		type: typ
	}
}

fn (mut s SymbolTable) lookup_var(name string) ?VarSymbolInfo {
	$if debug {
		dump(s.variable_scopes.len)
	}
	if s.variable_scopes.len == 0 {return none}
	for variables in s.variable_scopes.reverse() {
		if name in variables {
			return variables[name]
		}
	}
	return none
}

fn (mut s SymbolTable) define_func(name string, typ string, block Block) {
	s.functions[name] = FuncSymbolInfo{type: typ, block: block}
}

fn (mut s SymbolTable) lookup_func(name string) ?FuncSymbolInfo {
	if name in s.functions {
		return s.functions[name]
	}
	return none
}

fn (mut s SymbolTable) is_in_global_scope() bool {
	println("scope count: ${s.variable_scopes.len}")
	return s.variable_scopes.len == 1
}

// ------------------------------------------- Expressions

type Expr = VoidExpr | BinaryExpr | IntegerLiteral | RealLiteral | BoolLiteral | Variable | TypeExpr | Function | TypeCast
type LiteralExpr = IntegerLiteral | RealLiteral | BoolLiteral

struct VoidExpr {}

struct BinaryExpr {
	left Expr
	op string
	right Expr
}

struct IntegerLiteral {
	val i32
}

struct RealLiteral {
	val f32
}

struct BoolLiteral {
	val bool
}

struct Variable {
	name string
}

struct TypeExpr {
	name string
}

struct Function {
	name string
}

struct TypeCast {
	expr Expr
	type string
}

// ------------------------------------------- Statements

type Stmt = VarDecl | ExprStmt | ReturnStmt | Block | FuncDecl

struct VarDecl {
	name string
	value Expr
	type string
}

struct FuncDecl {
	name string
	ret_type string
	block Block
}

struct ExprStmt {
	expr Expr
}

struct ReturnStmt {
	value Expr
}

struct Block {
	stmts []Stmt
}

// ------------------------------------------- Parser

struct Parser {
	tokens []Token
mut:
	symbols SymbolTable
	pos int
	statements []Stmt
}

fn (mut p Parser) peek() Token {
	return p.tokens[p.pos]
}

fn (mut p Parser) next() Token {
	token := p.tokens[p.pos]
	p.pos++
	return token
}

fn (mut p Parser) expect(type TokenType) {
	if p.peek().type != type {
		parse_error('Expected ${str_from_toktype(type)}, got ${str_from_toktype(p.peek().type)}')
	}
	p.next()
}

// ------------------------------------------- Debug

fn (mut p Parser) dump_token() {
	$if debug {
		dump(p.peek())
	}
}

fn (mut p Parser) dump_stmt() {
	$if debug {
		if p.statements.len > 0 {
			dump(p.statements[p.statements.len-1])
		}
	}
}

@[noreturn]
fn parse_error(str string) {
	eprintln(term.red("Parse Error: " + str))
	panic("")
}

// ------------------------------------------- Expressions

fn (mut p Parser) parse_primary() Expr {
	token := p.next()

	p.dump_token()

	return match token.type {
		.integer    {IntegerLiteral{token.text.int()}}
		.real       {RealLiteral{token.text.f32()}}
		.boolean    {BoolLiteral{token.text == 'true'}}
		.identifier {Variable{token.text}}
		.kw_fn      {Function{token.text}}
		.type       {p.parse_type(token.text)}
		else {parse_error("Unexpected Token")}
	}
}

fn (mut p Parser) parse_expr() Expr {
	mut left := p.parse_primary()

	match p.peek().type {
		.plus   {return p.parse_binary(left, '+')}
		.minus  {return p.parse_binary(left, '-')}
		.star   {return p.parse_binary(left, '*')}
		.slash  {return p.parse_binary(left, '/')}
		.equals {return p.parse_binary(left, '=')}
		else {return left}
	}
}

fn (mut p Parser) parse_binary(left Expr, op string) BinaryExpr {
	p.next()
	right := p.parse_expr()
	binary_expr := BinaryExpr{left, op, right}
	if !p.is_op_valid_for_type(p.get_expr_type(binary_expr), op) {
		parse_error("Illegal operation ${op} for type ${p.get_expr_type(binary_expr)}")
	}
	return binary_expr
}

fn (mut p Parser) parse_type(type string) Expr {

	if p.peek().type == .lparen {
		p.next()
		expr := p.parse_expr()
		p.expect(.rparen)

		return TypeCast {
			expr: expr
			type: type
		}
	}

	return TypeExpr {name: type}
}

fn (mut p Parser) get_expr_type(expr Expr) string {
	return match expr {
		IntegerLiteral {'int'}
		RealLiteral    {'real'}
		BoolLiteral    {'bool'}
		VoidExpr       {'void'}
		BinaryExpr     {
			left_t := p.get_expr_type(expr.left)
			right_t := p.get_expr_type(expr.right)
			if left_t != right_t {
				parse_error ('Type mismatch in expression: ${left_t} and ${right_t}')
			}
			left_t
		}
		Variable {
			p.dump_stmt()
			info := p.symbols.lookup_var(expr.name) or {
				parse_error("Undefined variable ${expr.name}")
			}
			return info.type
		}
		TypeCast {expr.type}
		else {"Tried getting type of unexpected Expr"}
	}


}

fn (mut p Parser) is_op_valid_for_type(type string, op string) bool {
	legal_ops := match type {
		'int' {['+', '-', '*', '/', '<', '>', '=']}
		'real' {['+', '-', '*', '/', '<', '>', '=']}
		'bool' {['=']}
		else {[]}
	}
	return op in legal_ops
}

// ------------------------------------------- Statements

fn (mut p Parser) get_return_stmts_recursive(block Block) []ReturnStmt {
	mut returns := []ReturnStmt{}
	for stmt in block.stmts {
		if stmt is ReturnStmt {
			returns << stmt
		}
		if stmt is Block {
			returns << p.get_return_stmts_recursive(stmt)
		}
	}
	return returns
}

fn (mut p Parser) parse_statement() Stmt {
		match p.peek().type {
			.kw_let    {return p.parse_var_decl()}
			.kw_return {return p.parse_return_stmt()}
			.kw_fn     {return p.parse_func_decl()}
			.lbracket  {return p.parse_block(false)}
			else       {return p.parse_expr_stmt()}
		}
}



fn (mut p Parser) parse_var_decl() VarDecl {
	p.expect(.kw_let)

	name_tok := p.next()
	if name_tok.type != .identifier {
		parse_error("Expected variable name after let")
	}

	type_tok := p.next()
	if type_tok.type != .type {
		parse_error("Expected variable type after name when declaring ${name_tok.text}")
	}

	p.expect(.equals)
	val := p.parse_expr()

	if type_tok.text == 'void' {
		parse_error("Cannot declare a variable of type void")
	}
	if p.get_expr_type(val) != type_tok.text {
		parse_error("Mismatch between declared type (${type_tok.text}) and actual type (${p.get_expr_type(val)})")
	}
	p.expect(.semicolon)

	p.symbols.define_var(name_tok.text, type_tok.text)

	return VarDecl {
		name: name_tok.text
		value: val
		type: type_tok.text
	}
}

fn (mut p Parser) parse_func_decl() FuncDecl {

	if !p.symbols.is_in_global_scope() {parse_error("Tried to define a function in a non-global scope")}

	p.expect(.kw_fn)

	name_tok := p.next()
	if name_tok.type != .identifier {
		parse_error("Expected function name after let")
	}

	p.expect(.lparen)
	p.expect(.rparen)

	p.dump_token()

	type_tok := p.next()
	if type_tok.type != .type {
		parse_error("Expected function return type after name when declaring ${name_tok.text}")
	}

	p.symbols.variable_scopes << map[string]VarSymbolInfo{}
	block := p.parse_block(true)

	return_stmts := p.get_return_stmts_recursive(block)

	for return_stmt in return_stmts {
		if p.get_expr_type(return_stmt.value) != type_tok.text {
			parse_error("Mismatch between declared return type (${type_tok.text}) \
			and actual return type (${p.get_expr_type(return_stmt.value)})")
		}
	}

	p.symbols.variable_scopes.delete_last()

	p.symbols.define_func(name_tok.text, type_tok.text, block)

	return FuncDecl {
		name: name_tok.text
		ret_type: type_tok.text
		block: block
	}
}

fn (mut p Parser) parse_return_stmt() ReturnStmt {
	p.expect(.kw_return)

	token := p.peek()
	mut expr := Expr{}
	expr = match token.type {
		.integer    {IntegerLiteral{token.text.int()}}
		.real       {RealLiteral{token.text.f32()}}
		.boolean    {BoolLiteral{token.text == 'true'}}
		.identifier {Variable{token.text}}
		.semicolon  {VoidExpr{}}
		else {parse_error("Unexpected Token")}
	}
	p.next()
	if !(expr is VoidExpr) {
		p.expect(.semicolon)
	}
	return ReturnStmt {
		value: expr
	}
}

fn (mut p Parser) parse_expr_stmt() ExprStmt {
	expr := p.parse_expr()
	p.expect(.semicolon)

	return ExprStmt {
		expr: expr
	}
}


fn (mut p Parser) parse_block(no_scope bool) Block {
	p.expect(.lbracket)

	mut statements := []Stmt{}

	if !no_scope {
		p.symbols.variable_scopes << map[string]VarSymbolInfo{}
	}

	$if debug {
		println("entering scope")
	}

	for p.peek().type != .rbracket && p.peek().type != .eof {
		statements << p.parse_statement()
	}

	p.expect(.rbracket)


	return_stmts := (statements.filter(it is ReturnStmt).map(it as ReturnStmt))
	if return_stmts.len > 0 && (return_stmts.len > 1 || Stmt(return_stmts[0]) != statements[statements.len - 1]) {
		parse_error("Unexpected use of return. Unreachable code")
	}

	if !no_scope {
		p.symbols.variable_scopes.delete_last()
	}

	$if debug {
		println("exiting scope")
	}

	return Block {
		stmts: statements
	}

}

fn (mut p Parser) parse_program() []Stmt {
	p.symbols.variable_scopes << map[string]VarSymbolInfo{}
	for p.peek().type != .eof {
		p.statements << p.parse_statement()
	}
	p.symbols.variable_scopes.delete_last()
	$if debug {
		dump(p.symbols.functions)
	}
	return p.statements
}