package main

import (
	"slices"
	"strings"
)

type TypeType uint32

const (
	Type_Primitive TypeType = iota
	Type_Void
	Type_Named
	Type_Array
	Type_Tuple
)

type Type struct {
	Type  TypeType
	Value any
}

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
)

type Statement struct {
	Type  StatementType
	Value any
}

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 ExpressionType uint32

const (
	Expression_Assignment ExpressionType = iota
	Expression_Literal
	Expression_VariableReference
	Expression_Arithmetic
	Expression_Tuple
)

type Expression struct {
	Type      ExpressionType
	Value     any
	ValueType Type
}

type AssignmentExpression struct {
	Variable string
	Value    Expression
}

type LiteralExpression struct {
	Literal Literal
}

type VariableReferenceExpression struct {
	Variable string
}

type ArithmeticOperation uint32

const (
	Arithmetic_Add ArithmeticOperation = iota
	Arithmetic_Sub
	Arithmetic_Mul
	Arithmetic_Div
	Arithmetic_Mod
)

type ArithmeticExpression struct {
	Operation ArithmeticOperation
	Left      Expression
	Right     Expression
}

type TupleExpression struct {
	Members []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_Keyword && tok.Value.(Keyword) == Keyword_Void {
		*p = pCopy
		return &Type{Type: Type_Void, Value: 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)}, nil
		}

		*p = pCopy
		return &Type{Type: Type_Named, Value: tok.Value}, 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}}, 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)}}, 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}}}, nil
		}
	}

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

	return nil, nil
}

func (p *Parser) tryMultiplicativeExpression() (*Expression, error) {
	left, err := p.tryUnaryExpression()
	if err != nil {
		return nil, err
	}

	op, err := p.tryOperator(Operator_Multiply, Operator_Divide, Operator_Modulo)
	if err != nil {
		return nil, err
	}

	if op == nil {
		return left, nil
	}

	right, err := p.tryUnaryExpression()
	if err != nil {
		return nil, err
	}

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

	var operation ArithmeticOperation
	switch *op {
	case Operator_Multiply:
		operation = Arithmetic_Mul
	case Operator_Divide:
		operation = Arithmetic_Div
	case Operator_Modulo:
		fallthrough
	default:
		operation = Arithmetic_Mod
	}
	if *op == Operator_Plus {
		operation = Arithmetic_Add
	} else {
		operation = Arithmetic_Sub
	}

	return &Expression{Type: Expression_Arithmetic, Value: ArithmeticExpression{Operation: operation, Left: *left, Right: *right}}, nil
}

func (p *Parser) tryAdditiveExpression() (*Expression, error) {
	left, err := p.tryMultiplicativeExpression()
	if err != nil {
		return nil, err
	}

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

	if op == nil {
		return left, nil
	}

	right, err := p.tryMultiplicativeExpression()
	if err != nil {
		return nil, err
	}

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

	var operation ArithmeticOperation
	if *op == Operator_Plus {
		operation = Arithmetic_Add
	} else {
		operation = Arithmetic_Sub
	}

	return &Expression{Type: Expression_Arithmetic, Value: ArithmeticExpression{Operation: operation, Left: *left, Right: *right}}, nil
}

func (p *Parser) tryArithmeticExpression() (*Expression, error) {
	return p.tryAdditiveExpression()
}

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

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}}, 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}}, 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}}, nil
}

func (p *Parser) expectStatement() (*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}}, 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}}, nil
	}

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

		return &Statement{Type: Statement_Block, Value: BlockStatement{Block: *block}}, nil
	}

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

	if stmt != nil {
		return stmt, nil
	}

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

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

	var statements []Statement
	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()
		if err != nil {
			return nil, err
		}

		statements = append(statements, *stmt)
	}

	return &Block{Statements: statements}, nil
}

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

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

	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()
	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
}