package main

import (
	"slices"
	"strings"
)

type TypeType uint32

const (
	Type_Primitive TypeType = iota
	Type_Named
	Type_Array
	Type_Tuple
)

type Type struct {
	Type     TypeType
	Value    any
	Position uint64
}

type NamedType struct {
	TypeName string
}

type ArrayType struct {
	ElementType Type
}

type TupleType struct {
	Types []Type
}

type StatementType uint32

const (
	Statement_Expression StatementType = iota
	Statement_Block
	Statement_Return
	Statement_DeclareLocalVariable
	Statement_If
)

type Statement struct {
	Type     StatementType
	Value    any
	Position uint64
}

type ExpressionStatement struct {
	Expression Expression
}

type BlockStatement struct {
	Block *Block
}

type ReturnStatement struct {
	Value *Expression
}

type DeclareLocalVariableStatement struct {
	Variable     string
	VariableType Type
	Initializer  *Expression
}

type IfStatement struct {
	Condition        Expression
	ConditionalBlock *Block
	ElseBlock        *Block
}

type ExpressionType uint32

const (
	Expression_Assignment ExpressionType = iota
	Expression_Literal
	Expression_VariableReference
	Expression_Binary
	Expression_Tuple
	Expression_FunctionCall
	Expression_Negate
)

type Expression struct {
	Type      ExpressionType
	Value     any
	ValueType *Type
	Position  uint64
}

type AssignmentExpression struct {
	Variable string
	Value    Expression
}

type LiteralExpression struct {
	Literal Literal
}

type VariableReferenceExpression struct {
	Variable string
}

type Operation uint32

const (
	Operation_Add Operation = iota
	Operation_Sub
	Operation_Mul
	Operation_Div
	Operation_Mod
	Operation_Greater
	Operation_Less
	Operation_GreaterEquals
	Operation_LessEquals
	Operation_LogicalNot
	Operation_NotEquals
	Operation_Equals
)

type BinaryExpression struct {
	Operation  Operation
	Left       Expression
	Right      Expression
	ResultType *Type // Type to expand the operands to before performing the operation
}

type TupleExpression struct {
	Members []Expression
}

type FunctionCallExpression struct {
	Function   string
	Parameters *Expression
}

type NegateExpression struct {
	Value Expression
}

type Local struct {
	Name        string
	Type        Type
	IsParameter bool
	Index       int // unique, 0-based index of the local in the current function
}

type Block struct {
	Parent     *Block // TODO: implement
	Statements []Statement
	Locals     map[string]Local
}

type ParsedParameter struct {
	Name string
	Type Type
}

type ParsedFunction struct {
	Name       string
	Parameters []ParsedParameter
	ReturnType *Type
	Body       *Block
	Locals     []Local // All of the locals of the function, ordered by their index
}

type Import struct {
	Import string
}

type ParsedFile struct {
	Imports   []Import
	Functions []ParsedFunction
}

type Parser struct {
	Tokens   []LexToken
	Position uint64
}

func (p Parser) copy() Parser {
	return p
}

func (p *Parser) error(message string) CompilerError {
	return CompilerError{Position: p.Position, Message: message}
}

func (p *Parser) peekToken() *LexToken {
	if len(p.Tokens) == 0 {
		return nil
	}

	return &p.Tokens[0]
}

func (p *Parser) nextToken() *LexToken {
	if len(p.Tokens) == 0 {
		return nil
	}

	token := p.Tokens[0]
	p.Tokens = p.Tokens[1:]
	p.Position = token.Position
	return &token
}

func (p *Parser) trySeparator(separators ...Separator) (*Separator, error) {
	pCopy := p.copy()

	separator := pCopy.nextToken()
	if separator == nil || separator.Type != Type_Separator || !slices.Contains(separators, separator.Value.(Separator)) {
		return nil, nil
	}

	*p = pCopy
	sep := separator.Value.(Separator)
	return &sep, nil
}

func (p *Parser) expectSeparator(separators ...Separator) (Separator, error) {
	sep, err := p.trySeparator(separators...)
	if err != nil {
		return InvalidValue, err
	}

	if sep == nil {
		var separatorNames []string
		for _, sep := range separators {
			separatorNames = append(separatorNames, string(Separators[sep]))
		}

		return InvalidValue, p.error("expected one of " + strings.Join(separatorNames, " "))
	}

	return *sep, nil
}

func (p *Parser) tryOperator(operators ...Operator) (*Operator, error) {
	pCopy := p.copy()

	operator := pCopy.nextToken()
	if operator == nil || operator.Type != Type_Operator || !slices.Contains(operators, operator.Value.(Operator)) {
		return nil, nil
	}

	*p = pCopy
	sep := operator.Value.(Operator)
	return &sep, nil
}

func (p *Parser) expectIdentifier() (string, error) {
	identifier := p.nextToken()
	if identifier == nil || identifier.Type != Type_Identifier {
		return "", p.error("expected identifier")
	}

	return identifier.Value.(string), nil
}

func (p *Parser) expectImport() (*Import, error) {
	var err error

	importToken := p.nextToken()
	if importToken == nil || importToken.Type != Type_Keyword || importToken.Value.(Keyword) != Keyword_Import {
		return nil, p.error("expected import")
	}

	identifier := p.nextToken()
	if identifier == nil || identifier.Type != Type_Identifier {
		return nil, p.error("expected identifier")
	}

	_, err = p.expectSeparator(Separator_Semicolon)
	if err != nil {
		return nil, err
	}

	return &Import{Import: identifier.Value.(string)}, nil
}

func (p *Parser) tryType() (*Type, error) {
	pCopy := p.copy()

	tok := pCopy.nextToken()
	if tok == nil {
		return nil, nil
	}

	if tok.Type == Type_Identifier {
		// TODO: array type

		index := slices.Index(PRIMITIVE_TYPE_NAMES, tok.Value.(string))
		if index != -1 {
			*p = pCopy
			return &Type{Type: Type_Primitive, Value: PrimitiveType(index), Position: tok.Position}, nil
		}

		*p = pCopy
		return &Type{Type: Type_Named, Value: tok.Value, Position: tok.Position}, nil
	}

	return nil, nil
}

func (p *Parser) expectType() (*Type, error) {
	t, err := p.tryType()
	if err != nil {
		return nil, err
	}

	if t != nil {
		return t, nil
	}

	return nil, p.error("expected type")
}

func (p *Parser) expectTypeOrTupleType() (*Type, error) {
	var err error

	tok := p.peekToken()
	if tok == nil {
		return nil, p.error("expected type or tuple type")
	}

	if tok.Type == Type_Separator && tok.Value.(Separator) == Separator_OpenParen {
		// Tuple type
		p.nextToken()
		var types []Type
		for {
			parsedType, err := p.expectType()
			if err != nil {
				return nil, err
			}

			types = append(types, *parsedType)

			var sep Separator
			sep, err = p.expectSeparator(Separator_Comma, Separator_CloseParen)
			if err != nil {
				return nil, err
			}

			if sep == Separator_CloseParen {
				break
			}
		}

		if len(types) == 0 {
			return nil, p.error("empty tuple")
		}

		return &Type{Type: Type_Tuple, Value: TupleType{Types: types}, Position: tok.Position}, nil
	}

	t, err := p.tryType()
	if err != nil {
		return nil, err
	}

	if t != nil {
		return t, nil
	}

	return nil, p.error("expected type or tuple type")
}

func (p *Parser) tryParanthesizedExpression() (*Expression, error) {
	pCopy := p.copy()

	token := pCopy.nextToken()
	if token == nil || token.Type != Type_Separator || token.Value.(Separator) != Separator_OpenParen {
		return nil, nil
	}

	expr, err := pCopy.tryExpression()
	if err != nil {
		*p = pCopy
		return nil, err
	}

	if expr == nil {
		return nil, nil
	}

	// Assuming we're in a paranthesized statement
	_, err = pCopy.expectSeparator(Separator_CloseParen)
	if err != nil {
		*p = pCopy
		return nil, err
	}

	*p = pCopy
	return expr, nil
}

func (p *Parser) tryUnaryExpression() (*Expression, error) {
	pCopy := p.copy()

	token := pCopy.peekToken()
	if token == nil {
		return nil, nil
	}

	if token.Type == Type_Separator && token.Value.(Separator) == Separator_OpenParen {
		// TODO: cast

		paren, err := pCopy.tryParanthesizedExpression()
		if err != nil {
			*p = pCopy
			return nil, err
		}

		if paren != nil {
			*p = pCopy
			return paren, nil
		}
	}

	if token.Type == Type_Literal {
		pCopy.nextToken()
		*p = pCopy
		return &Expression{Type: Expression_Literal, Value: LiteralExpression{Literal: token.Value.(Literal)}, Position: token.Position}, nil
	}

	if token.Type == Type_Keyword {
		keyword := token.Value.(Keyword)
		if keyword == Keyword_True || keyword == KeyWord_False {
			pCopy.nextToken()
			*p = pCopy
			return &Expression{Type: Expression_Literal, Value: LiteralExpression{Literal: Literal{Type: Literal_Boolean, Primitive: Primitive_Bool, Value: keyword == Keyword_True}}, Position: token.Position}, nil
		}
	}

	if token.Type == Type_Operator {
		op := token.Value.(Operator)
		if op == Operator_Minus || op == Operator_Plus {
			pCopy.nextToken()
			expr, err := pCopy.tryUnaryExpression()
			if err != nil {
				return nil, err
			}

			if expr == nil {
				return nil, nil
			}

			if op == Operator_Minus {
				expr = &Expression{Type: Expression_Negate, Value: NegateExpression{Value: *expr}, Position: token.Position}
			}

			return expr, nil
		}
	}

	if token.Type == Type_Identifier {
		pCopy.nextToken()

		next, err := pCopy.trySeparator(Separator_OpenParen)
		if err != nil {
			return nil, err
		}

		if next != nil {
			// Function call
			params, err := pCopy.tryTupleExpression()
			if err != nil {
				return nil, err
			}

			_, err = pCopy.expectSeparator(Separator_CloseParen)
			if err != nil {
				return nil, err
			}

			*p = pCopy
			return &Expression{Type: Expression_FunctionCall, Value: FunctionCallExpression{Function: token.Value.(string), Parameters: params}, Position: token.Position}, nil
		}

		*p = pCopy
		return &Expression{Type: Expression_VariableReference, Value: VariableReferenceExpression{Variable: token.Value.(string)}, Position: token.Position}, nil
	}

	return nil, nil
}

func (p *Parser) tryBinaryExpression0(opFunc func() (*Expression, error), operators ...Operator) (*Expression, error) {
	left, err := opFunc()
	if err != nil {
		return nil, err
	}

	if left == nil {
		return nil, nil
	}

	for {
		op, err := p.tryOperator(operators...)
		if err != nil {
			return nil, err
		}

		if op == nil {
			return left, nil
		}

		right, err := opFunc()
		if err != nil {
			return nil, err
		}

		if right == nil {
			return nil, p.error("expected expression")
		}

		left = &Expression{Type: Expression_Binary, Value: BinaryExpression{Operation: getOperation(*op), Left: *left, Right: *right}, Position: left.Position}
	}
}

func (p *Parser) tryMultiplicativeExpression() (*Expression, error) {
	return p.tryBinaryExpression0(p.tryUnaryExpression, Operator_Multiply, Operator_Divide, Operator_Modulo)
}

func (p *Parser) tryAdditiveExpression() (*Expression, error) {
	return p.tryBinaryExpression0(p.tryMultiplicativeExpression, Operator_Plus, Operator_Minus)
}

func (p *Parser) tryRelationalExpression() (*Expression, error) {
	return p.tryBinaryExpression0(p.tryAdditiveExpression, Operator_Less, Operator_Greater, Operator_LessEquals, Operator_GreaterEquals)
}

func (p *Parser) tryEqualityExpression() (*Expression, error) {
	return p.tryBinaryExpression0(p.tryRelationalExpression, Operator_EqualsEquals, Operator_NotEquals)
}

func (p *Parser) tryBinaryExpression() (*Expression, error) {
	return p.tryEqualityExpression()
}

func (p *Parser) tryExpression() (*Expression, error) {
	return p.tryBinaryExpression()
}

func (p *Parser) expectExpression() (*Expression, error) {
	return p.expect(p.tryExpression, "expected expression")
}

func (p *Parser) tryTupleExpression() (*Expression, error) {
	pCopy := p.copy()

	var members []Expression
	for {
		expr, err := pCopy.tryExpression()
		if err != nil {
			return nil, err
		}

		if expr == nil {
			return nil, nil
		}

		members = append(members, *expr)

		token := pCopy.peekToken()
		if token == nil || token.Type != Type_Separator || token.Value.(Separator) != Separator_Comma {
			if len(members) == 1 {
				*p = pCopy
				return expr, nil
			}

			break
		}

		pCopy.nextToken()
	}

	*p = pCopy
	return &Expression{Type: Expression_Tuple, Value: TupleExpression{Members: members}, Position: members[0].Position}, nil
}

func (p *Parser) expectTupleExpression() (*Expression, error) {
	return p.expect(p.tryTupleExpression, "expected tuple expression")
}

func (p *Parser) expect(try func() (*Expression, error), message string) (*Expression, error) {
	expr, err := try()
	if err != nil {
		return nil, err
	}

	if expr == nil {
		return nil, p.error(message)
	}

	return expr, nil
}

func (p *Parser) tryDeclareLocalVariableStatement() (*Statement, error) {
	pCopy := p.copy()

	variableType, err := pCopy.tryType()
	if err != nil {
		*p = pCopy
		return nil, err
	}

	if variableType == nil {
		return nil, nil
	}

	name := pCopy.nextToken()
	if name == nil || name.Type != Type_Identifier {
		return nil, nil
	}

	variableName := name.Value.(string)

	token := pCopy.nextToken()
	if token.Type == Type_Separator && token.Value.(Separator) == Separator_Semicolon {
		*p = pCopy
		return &Statement{Type: Statement_DeclareLocalVariable, Value: DeclareLocalVariableStatement{Variable: variableName, VariableType: *variableType, Initializer: nil}, Position: variableType.Position}, nil
	}

	if token.Type != Type_Operator || token.Value.(Operator) != Operator_Equals {
		return nil, nil
	}

	initializer, err := pCopy.expect(pCopy.tryExpression, "expected initializer expression")
	if err != nil {
		*p = pCopy
		return nil, err
	}

	_, err = pCopy.expectSeparator(Separator_Semicolon)
	if err != nil {
		*p = pCopy
		return nil, err
	}

	*p = pCopy
	return &Statement{Type: Statement_DeclareLocalVariable, Value: DeclareLocalVariableStatement{Variable: variableName, VariableType: *variableType, Initializer: initializer}, Position: variableType.Position}, nil
}

func (p *Parser) expectStatement(block *Block) (*Statement, error) {
	token := p.peekToken()
	if token == nil {
		return nil, p.error("expected statement")
	}

	if token.Type == Type_Keyword && token.Value.(Keyword) == Keyword_Return {
		p.nextToken()

		token = p.peekToken()
		if token == nil {
			return nil, p.error("expected expression or ;")
		}

		if token.Type == Type_Separator && token.Value.(Separator) == Separator_Semicolon {
			p.nextToken()
			return &Statement{Type: Statement_Return, Value: ReturnStatement{Value: nil}, Position: token.Position}, nil
		}

		expr, err := p.expectTupleExpression()
		if err != nil {
			return nil, err
		}

		_, err = p.expectSeparator(Separator_Semicolon)
		if err != nil {
			return nil, err
		}

		return &Statement{Type: Statement_Return, Value: ReturnStatement{Value: expr}, Position: token.Position}, nil
	}

	if token.Type == Type_Keyword && token.Value.(Keyword) == Keyword_If {
		p.nextToken()

		_, err := p.expectSeparator(Separator_OpenParen)
		if err != nil {
			return nil, err
		}

		cond, err := p.expectExpression()
		if err != nil {
			return nil, err
		}

		_, err = p.expectSeparator(Separator_CloseParen)
		if err != nil {
			return nil, err
		}

		conditionalBlock, err := p.expectBlock(block)
		if err != nil {
			return nil, err
		}

		tok := p.peekToken()
		if tok == nil || tok.Type != Type_Keyword || tok.Value.(Keyword) != Keyword_Else {
			return &Statement{Type: Statement_If, Value: IfStatement{Condition: *cond, ConditionalBlock: conditionalBlock, ElseBlock: nil}, Position: token.Position}, nil
		}

		p.nextToken()

		var elseBlock *Block

		tok = p.peekToken()
		if tok.Type == Type_Keyword && tok.Value.(Keyword) == Keyword_If {
			stmt, err := p.expectStatement(block)
			if err != nil {
				return nil, err
			}

			elseBlock = &Block{Parent: block, Statements: []Statement{*stmt}}
		} else {
			elseBlock, err = p.expectBlock(block)
			if err != nil {
				return nil, err
			}
		}

		return &Statement{Type: Statement_If, Value: IfStatement{Condition: *cond, ConditionalBlock: conditionalBlock, ElseBlock: elseBlock}, Position: token.Position}, nil
	}

	if token.Type == Type_Separator && token.Value.(Separator) == Separator_OpenCurly {
		block, err := p.expectBlock(block)
		if err != nil {
			return nil, err
		}

		return &Statement{Type: Statement_Block, Value: BlockStatement{Block: block}, Position: token.Position}, nil
	}

	stmt, err := p.tryDeclareLocalVariableStatement()
	if err != nil {
		return nil, err
	}

	if stmt != nil {
		return stmt, nil
	}

	expr, err := p.tryExpression()
	if err != nil {
		return nil, err
	}

	if expr != nil {
		_, err := p.expectSeparator(Separator_Semicolon)
		if err != nil {
			return nil, err
		}

		return &Statement{Type: Statement_Expression, Value: ExpressionStatement{Expression: *expr}, Position: expr.Position}, nil
	}

	return nil, p.error("expected statement")
}

func (p *Parser) expectBlock(parent *Block) (*Block, error) {
	_, err := p.expectSeparator(Separator_OpenCurly)
	if err != nil {
		return nil, err
	}

	block := &Block{Parent: parent}

	for {
		token := p.peekToken()
		if token == nil {
			return nil, p.error("expected statement or }")
		}

		if token.Type == Type_Separator && token.Value.(Separator) == Separator_CloseCurly {
			p.nextToken()
			break
		}

		stmt, err := p.expectStatement(block)
		if err != nil {
			return nil, err
		}

		block.Statements = append(block.Statements, *stmt)
	}

	return block, nil
}

func (p *Parser) expectFunction() (*ParsedFunction, error) {
	var err error

	var name string
	var parameters []ParsedParameter
	var returnType *Type
	var body *Block

	tok := p.peekToken()
	if tok.Type == Type_Keyword && tok.Value.(Keyword) == Keyword_Void {
		p.nextToken()
		returnType = nil
	} else {
		returnType, err = p.expectTypeOrTupleType()
		if err != nil {
			return nil, err
		}
	}

	name, err = p.expectIdentifier()
	if err != nil {
		return nil, err
	}

	_, err = p.expectSeparator(Separator_OpenParen)
	if err != nil {
		return nil, err
	}

	for {
		token := p.peekToken()
		if token == nil {
			return nil, p.error("incomplete function declaration")
		}

		if token.Type == Type_Separator && token.Value.(Separator) == Separator_CloseParen {
			p.nextToken()
			break
		}

		if len(parameters) != 0 {
			_, err := p.expectSeparator(Separator_Comma)
			if err != nil {
				return nil, err
			}
		}

		var paramType *Type
		paramType, err = p.expectType()
		if err != nil {
			return nil, err
		}

		var paramName string
		paramName, err = p.expectIdentifier()
		if err != nil {
			return nil, err
		}

		parameters = append(parameters, ParsedParameter{Name: paramName, Type: *paramType})
	}

	body, err = p.expectBlock(nil)
	if err != nil {
		return nil, err
	}

	return &ParsedFunction{Name: name, Parameters: parameters, ReturnType: returnType, Body: body}, nil
}

func (p *Parser) parseFile() (*ParsedFile, error) {
	var err error

	var functions []ParsedFunction
	var imports []Import

	for {
		token := p.peekToken()
		if token == nil {
			break
		}

		if token.Type == Type_Keyword && token.Value.(Keyword) == Keyword_Import {
			var parsedImport *Import
			parsedImport, err = p.expectImport()
			if err != nil {
				return nil, err
			}

			imports = append(imports, *parsedImport)
			continue
		}

		var parsedFunction *ParsedFunction
		parsedFunction, err = p.expectFunction()
		if err != nil {
			return nil, err
		}

		functions = append(functions, *parsedFunction)
	}

	return &ParsedFile{Imports: imports, Functions: functions}, nil
}