package main

import (
	"errors"
	"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 BlockStatement struct {
	Block Block
}

type ReturnStatement struct {
	Value *Expression
}

type DeclareLocalVariableStatement struct {
	Variable    string
	Initializer Expression
}

type ExpressionType uint32

const (
	Expression_Assignment ExpressionType = iota
	Expression_Literal
	Expression_VariableReference
	Expression_Arithmetic
)

type Expression struct {
	Type  ExpressionType
	Value any
}

type AssignmentExpression struct {
	Variable string
	Value    Expression
}

type LiteralExpression struct {
	Type  PrimitiveType
	Value any
}

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 Block struct {
	Statements []Statement
}

type ParsedFunction struct {
	Parameters []Type
	ReturnType Type
	Body       Block
}

type Import struct {
	Import string
}

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

func expectSeparator(tokens []LexToken, separators ...Separator) (Separator, []LexToken, error) {
	var separatorNames []string
	for _, sep := range separators {
		separatorNames = append(separatorNames, string(Separators[sep]))
	}

	if len(tokens) == 0 {
		return InvalidValue, nil, errors.New("expected one of " + strings.Join(separatorNames, " "))
	}

	separator := tokens[0]
	if separator.Type != Type_Separator || !slices.Contains(separators, separator.Value.(Separator)) {
		return InvalidValue, nil, errors.New("expected one of " + strings.Join(separatorNames, ", "))
	}

	return separator.Value.(Separator), tokens[1:], nil
}

func parseImport(tokens []LexToken) (*Import, []LexToken, error) {
	var err error

	if len(tokens) < 3 {
		return nil, nil, errors.New("incomplete import")
	}

	// tokens[0] == import keyword

	identifier := tokens[1]
	if identifier.Type != Type_Identifier {
		return nil, nil, errors.New("expected identifier")
	}

	_, tokens, err = expectSeparator(tokens[2:], Separator_Semicolon)
	if err != nil {
		return nil, nil, err
	}

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

func parseType(tokens []LexToken) (*Type, []LexToken, error) {
	var err error

	if len(tokens) == 0 {
		return nil, nil, errors.New("expected type")
	}

	tok := tokens[0]
	if tok.Type == Type_Keyword && tok.Value.(Keyword) == Keyword_Void {
		return &Type{Type: Type_Void, Value: nil}, tokens[1:], nil
	}

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

			types = append(types, *parsedType)

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

			if sep == Separator_CloseParen {
				break
			}
		}

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

		return &Type{Type: Type_Tuple, Value: TupleType{Types: types}}, tokens, nil
	}

	if tok.Type == Type_Identifier {
		return &Type{Type: Type_Named, Value: tok.Value}, tokens[1:], nil
	}

	return nil, nil, errors.New("expected type")
}

func parseFunction(tokens []LexToken) (*ParsedFunction, []LexToken, error) {
	var err error

	var parameters []Type
	var returnType *Type
	var body Block

	returnType, tokens, err = parseType(tokens)
	if err != nil {
		return nil, nil, err
	}

	// TODO: function name, parameters, body

	return &ParsedFunction{Parameters: parameters, ReturnType: *returnType, Body: body}, tokens, nil
}

func parser(tokens []LexToken) (*ParsedFile, error) {
	var err error

	var functions []ParsedFunction
	var imports []Import

	for len(tokens) > 0 {
		if tokens[0].Type == Type_Keyword && tokens[0].Value.(Keyword) == Keyword_Import {
			var parsedImport *Import
			parsedImport, tokens, err = parseImport(tokens)
			if err != nil {
				return nil, err
			}

			imports = append(imports, *parsedImport)
			continue
		}

		var parsedFunction *ParsedFunction
		parsedFunction, tokens, err = parseFunction(tokens)
		if err != nil {
			return nil, err
		}

		functions = append(functions, *parsedFunction)
	}

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