package main

import (
	"errors"
	"log"
	"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 ParsedParameter struct {
	Name string
	Type Type
}

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

type Import struct {
	Import string
}

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

type Parser struct {
	tokens []LexToken
}

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:]
	return &token
}

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

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

	return separator.Value.(Separator), nil
}

func (p *Parser) expectIdentifier() (string, error) {
	identifier := p.nextToken()
	if identifier == nil || identifier.Type != Type_Separator && identifier.Type != Type_Identifier {
		return "", errors.New("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, errors.New("expected import")
	}

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

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

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

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

	tok := p.nextToken()
	if tok == nil {
		return nil, errors.New("expected type")
	}

	if tok.Type == Type_Keyword && tok.Value.(Keyword) == Keyword_Void {
		return &Type{Type: Type_Void, Value: nil}, nil
	}

	if tok.Type == Type_Separator && tok.Value.(Separator) == Separator_OpenParen {
		// Tuple type
		var types []Type
		for {
			var parsedType *Type
			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, errors.New("empty tuple")
		}

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

	if tok.Type == Type_Identifier {
		return &Type{Type: Type_Named, Value: tok.Value}, nil
	}

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

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

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

	returnType, err = p.expectType()
	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, errors.New("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})
	}

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

	// TODO: body, closing curly

	_, err = p.expectSeparator(Separator_CloseCurly)
	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
		}

		log.Printf("%+#v\n", token)

		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
}