package saccharine

import (
	"errors"
	"fmt"

	"git.maximhutz.com/max/lambda/pkg/iterator"
	"git.maximhutz.com/max/lambda/pkg/saccharine/ast"
	"git.maximhutz.com/max/lambda/pkg/saccharine/token"
	"git.maximhutz.com/max/lambda/pkg/trace"
)

type TokenIterator = iterator.Iterator[token.Token]

func parseRawToken(i *TokenIterator, expected token.Type) (*token.Token, error) {
	return iterator.Do(i, func(i *TokenIterator) (*token.Token, error) {
		if tok, err := i.Next(); err != nil {
			return nil, err
		} else if tok.Type != expected {
			return nil, fmt.Errorf("expected token %v, got %v'", token.Name(expected), tok.Value)
		} else {
			return &tok, nil
		}
	})
}

func passSoftBreaks(i *TokenIterator) {
	for {
		if _, err := parseRawToken(i, token.SoftBreak); err != nil {
			return
		}
	}
}

func parseToken(i *TokenIterator, expected token.Type, ignoreSoftBreaks bool) (*token.Token, error) {
	return iterator.Do(i, func(i *TokenIterator) (*token.Token, error) {
		if ignoreSoftBreaks {
			passSoftBreaks(i)
		}

		return parseRawToken(i, expected)
	})
}

func parseString(i *TokenIterator) (string, error) {
	if tok, err := parseToken(i, token.Atom, true); err != nil {
		return "", trace.Wrap(err, "no variable (col %d)", i.Index())
	} else {
		return tok.Value, nil
	}
}

func parseBreak(i *TokenIterator) (*token.Token, error) {
	if tok, softErr := parseRawToken(i, token.SoftBreak); softErr == nil {
		return tok, nil
	} else if tok, hardErr := parseRawToken(i, token.HardBreak); hardErr == nil {
		return tok, nil
	} else {
		return nil, errors.Join(softErr, hardErr)
	}
}

func parseList[U any](i *TokenIterator, fn func(*TokenIterator) (U, error), minimum int) ([]U, error) {
	results := []U{}

	for {
		if u, err := fn(i); err != nil {
			if len(results) < minimum {
				return nil, trace.Wrap(err, "expected at least '%v' items, got only '%v'", minimum, len(results))
			}
			return results, nil
		} else {
			results = append(results, u)
		}
	}
}

func parseAbstraction(i *TokenIterator) (*ast.Abstraction, error) {
	return iterator.Do(i, func(i *TokenIterator) (*ast.Abstraction, error) {
		if _, err := parseToken(i, token.Slash, true); err != nil {
			return nil, trace.Wrap(err, "no function slash (col %d)", i.MustGet().Column)
		} else if parameters, err := parseList(i, parseString, 0); err != nil {
			return nil, err
		} else if _, err = parseToken(i, token.Dot, true); err != nil {
			return nil, trace.Wrap(err, "no function dot (col %d)", i.MustGet().Column)
		} else if body, err := parseExpression(i); err != nil {
			return nil, err
		} else {
			return ast.NewAbstraction(parameters, body), nil
		}
	})
}

func parseApplication(i *TokenIterator) (*ast.Application, error) {
	return iterator.Do(i, func(i *TokenIterator) (*ast.Application, error) {
		if _, err := parseToken(i, token.OpenParen, true); err != nil {
			return nil, trace.Wrap(err, "no openning brackets (col %d)", i.MustGet().Column)
		} else if expressions, err := parseList(i, parseExpression, 1); err != nil {
			return nil, err
		} else if _, err := parseToken(i, token.CloseParen, true); err != nil {
			return nil, trace.Wrap(err, "no closing brackets (col %d)", i.MustGet().Column)
		} else {
			return ast.NewApplication(expressions[0], expressions[1:]), nil
		}
	})
}

func parseAtom(i *TokenIterator) (*ast.Atom, error) {
	if tok, err := parseToken(i, token.Atom, true); err != nil {
		return nil, trace.Wrap(err, "no variable (col %d)", i.Index())
	} else {
		return ast.NewAtom(tok.Value), nil
	}
}

func parseStatements(i *TokenIterator) ([]ast.Statement, error) {
	statements := []ast.Statement{}

	//nolint:errcheck
	parseList(i, parseBreak, 0)

	for {
		if statement, err := parseStatement(i); err != nil {
			break
		} else if _, err := parseList(i, parseBreak, 1); err != nil && !i.Done() {
			break
		} else {
			statements = append(statements, statement)
		}
	}

	return statements, nil
}

func parseClause(i *TokenIterator, braces bool) (*ast.Clause, error) {
	if braces {
		if _, err := parseToken(i, token.OpenBrace, true); err != nil {
			return nil, err
		}
	}

	var stmts []ast.Statement
	var last *ast.DeclareStatement
	var err error
	var ok bool

	if stmts, err = parseStatements(i); err != nil {
		return nil, err
	} else if len(stmts) == 0 {
		return nil, fmt.Errorf("no statements in clause")
	} else if last, ok = stmts[len(stmts)-1].(*ast.DeclareStatement); !ok {
		return nil, fmt.Errorf("this clause contains no final return value (col %d)", i.MustGet().Column)
	}

	if braces {
		if _, err := parseToken(i, token.CloseBrace, true); err != nil {
			return nil, err
		}
	}

	return ast.NewClause(stmts[:len(stmts)-1], last.Value), nil
}

func parseExpression(i *TokenIterator) (ast.Expression, error) {
	return iterator.Do(i, func(i *TokenIterator) (ast.Expression, error) {
		passSoftBreaks(i)

		switch peek := i.MustGet(); peek.Type {
		case token.OpenParen:
			return parseApplication(i)
		case token.Slash:
			return parseAbstraction(i)
		case token.Atom:
			return parseAtom(i)
		case token.OpenBrace:
			return parseClause(i, true)
		default:
			return nil, fmt.Errorf("expected expression, got '%v' (col %d)", peek.Value, peek.Column)
		}
	})
}

func parseLet(i *TokenIterator) (*ast.LetStatement, error) {
	return iterator.Do(i, func(i *TokenIterator) (*ast.LetStatement, error) {
		if parameters, err := parseList(i, parseString, 1); err != nil {
			return nil, err
		} else if _, err := parseToken(i, token.Assign, true); err != nil {
			return nil, err
		} else if body, err := parseExpression(i); err != nil {
			return nil, err
		} else {
			return ast.NewLet(parameters[0], parameters[1:], body), nil
		}
	})
}

func parseDeclare(i *TokenIterator) (*ast.DeclareStatement, error) {
	if value, err := parseExpression(i); err != nil {
		return nil, err
	} else {
		return ast.NewDeclare(value), nil
	}
}

func parseStatement(i *TokenIterator) (ast.Statement, error) {
	if let, letErr := parseLet(i); letErr == nil {
		return let, nil
	} else if declare, declErr := parseDeclare(i); declErr == nil {
		return declare, nil
	} else {
		return nil, errors.Join(letErr, declErr)
	}
}

func Parse(tokens []token.Token) (ast.Expression, error) {
	i := iterator.Of(tokens)

	exp, err := parseClause(i, false)
	if err != nil {
		return nil, err
	}

	if !i.Done() {
		return nil, fmt.Errorf("expected EOF, found more code (col %d)", i.MustGet().Column)
	}

	return exp, nil
}