package p.studio.compiler.pbs.parser;

import p.studio.compiler.pbs.ast.PbsAst;
import p.studio.compiler.pbs.lexer.PbsToken;
import p.studio.compiler.pbs.lexer.PbsTokenKind;

import java.util.ArrayList;

final class PbsTypeParser {
    private static final int MAX_NAMED_TUPLE_ARITY = 6;

    private final PbsParserContext context;
    private final PbsTokenCursor cursor;

    PbsTypeParser(final PbsParserContext context) {
        this.context = context;
        this.cursor = context.cursor();
    }

    ArrayList<PbsAst.Parameter> parseParametersUntilRightParen() {
        final var parameters = new ArrayList<PbsAst.Parameter>();
        if (cursor.check(PbsTokenKind.RIGHT_PAREN)) {
            return parameters;
        }

        do {
            final var parameterStart = cursor.peek();
            final var parameterName = consume(PbsTokenKind.IDENTIFIER, "Expected parameter name");
            consume(PbsTokenKind.COLON, "Expected ':' after parameter name");
            final var typeRef = parseTypeRef();
            parameters.add(new PbsAst.Parameter(
                    parameterName.lexeme(),
                    typeRef,
                    context.span(parameterStart.start(), typeRef.span().getEnd())));
        } while (cursor.match(PbsTokenKind.COMMA));

        return parameters;
    }

    ParsedReturnSpec parseCallableReturnSpec(final PbsToken anchorToken) {
        if (cursor.match(PbsTokenKind.ARROW)) {
            return parseReturnSpecFromMarker();
        }

        if (cursor.match(PbsTokenKind.COLON)) {
            context.report(cursor.previous(), ParseErrors.E_PARSE_INVALID_RETURN_ANNOTATION,
                    "Return annotations must use '->' syntax");
            return parseReturnSpecFromMarker();
        }

        final var inferredSpan = context.span(anchorToken.end(), anchorToken.end());
        return new ParsedReturnSpec(
                PbsAst.ReturnKind.INFERRED_UNIT,
                PbsAst.TypeRef.unit(inferredSpan),
                null);
    }

    PbsAst.TypeRef parseTypeRef() {
        if (cursor.match(PbsTokenKind.OPTIONAL)) {
            final var optionalToken = cursor.previous();
            final var inner = parseTypeRef();
            if (inner.kind() == PbsAst.TypeRefKind.UNIT) {
                context.report(optionalToken, ParseErrors.E_PARSE_INVALID_TYPE_SURFACE,
                        "'optional void' is not a valid type surface");
            }
            return PbsAst.TypeRef.optional(inner, context.span(optionalToken.start(), inner.span().getEnd()));
        }

        if (cursor.match(PbsTokenKind.VOID)) {
            final var token = cursor.previous();
            return PbsAst.TypeRef.unit(context.span(token.start(), token.end()));
        }

        if (cursor.match(PbsTokenKind.SELF)) {
            final var token = cursor.previous();
            return PbsAst.TypeRef.self(context.span(token.start(), token.end()));
        }

        if (cursor.match(PbsTokenKind.IDENTIFIER)) {
            final var token = cursor.previous();
            return PbsAst.TypeRef.simple(token.lexeme(), context.span(token.start(), token.end()));
        }

        if (cursor.match(PbsTokenKind.LEFT_PAREN)) {
            final var open = cursor.previous();
            if (cursor.match(PbsTokenKind.RIGHT_PAREN)) {
                return PbsAst.TypeRef.unit(context.span(open.start(), cursor.previous().end()));
            }

            if (cursor.check(PbsTokenKind.IDENTIFIER) && cursor.checkNext(PbsTokenKind.COLON)) {
                final var fields = parseNamedTupleTypeFields();
                final var close = consume(PbsTokenKind.RIGHT_PAREN, "Expected ')' after named tuple type");
                return PbsAst.TypeRef.namedTuple(
                        p.studio.utilities.structures.ReadOnlyList.wrap(fields),
                        context.span(open.start(), close.end()));
            }

            final var inner = parseTypeRef();
            final var close = consume(PbsTokenKind.RIGHT_PAREN, "Expected ')' after grouped type");
            return PbsAst.TypeRef.group(inner, context.span(open.start(), close.end()));
        }

        final var token = cursor.peek();
        context.report(token, ParseErrors.E_PARSE_EXPECTED_TOKEN, "Expected type surface");
        if (!cursor.isAtEnd()) {
            cursor.advance();
        }
        return PbsAst.TypeRef.error(context.span(token.start(), token.end()));
    }

    private ParsedReturnSpec parseReturnSpecFromMarker() {
        if (cursor.match(PbsTokenKind.RESULT)) {
            final var resultToken = cursor.previous();
            consume(PbsTokenKind.LESS, "Expected '<' after 'result'");
            final var errorType = parseTypeRef();
            consume(PbsTokenKind.GREATER, "Expected '>' after result error type");

            final PbsAst.TypeRef payloadType;
            if (isTypeStart(cursor.peek().kind())) {
                payloadType = parseTypeRef();
            } else {
                payloadType = PbsAst.TypeRef.unit(context.span(resultToken.end(), resultToken.end()));
            }

            return new ParsedReturnSpec(
                    PbsAst.ReturnKind.RESULT,
                    payloadType,
                    errorType);
        }

        final var typeRef = parseTypeRef();
        if (typeRef.kind() == PbsAst.TypeRefKind.UNIT) {
            return new ParsedReturnSpec(PbsAst.ReturnKind.EXPLICIT_UNIT, typeRef, null);
        }
        return new ParsedReturnSpec(PbsAst.ReturnKind.PLAIN, typeRef, null);
    }

    private ArrayList<PbsAst.NamedTypeField> parseNamedTupleTypeFields() {
        final var fields = new ArrayList<PbsAst.NamedTypeField>();
        do {
            final var label = consume(PbsTokenKind.IDENTIFIER, "Expected tuple field label");
            consume(PbsTokenKind.COLON, "Expected ':' after tuple field label");
            final var typeRef = parseTypeRef();
            fields.add(new PbsAst.NamedTypeField(
                    label.lexeme(),
                    typeRef,
                    context.span(label.start(), typeRef.span().getEnd())));
            if (fields.size() == MAX_NAMED_TUPLE_ARITY + 1) {
                context.report(label, ParseErrors.E_PARSE_INVALID_TYPE_SURFACE,
                        "Named tuple type arity must be between 1 and 6 fields");
            }
        } while (cursor.match(PbsTokenKind.COMMA) && !cursor.check(PbsTokenKind.RIGHT_PAREN));
        return fields;
    }

    private boolean isTypeStart(final PbsTokenKind kind) {
        return kind == PbsTokenKind.IDENTIFIER
                || kind == PbsTokenKind.SELF
                || kind == PbsTokenKind.OPTIONAL
                || kind == PbsTokenKind.VOID
                || kind == PbsTokenKind.LEFT_PAREN;
    }

    private PbsToken consume(final PbsTokenKind kind, final String message) {
        if (cursor.check(kind)) {
            return cursor.advance();
        }
        final var token = cursor.peek();
        context.report(token, ParseErrors.E_PARSE_EXPECTED_TOKEN, message + ", found " + token.kind());
        if (!cursor.isAtEnd()) {
            return cursor.advance();
        }
        return token;
    }

    record ParsedReturnSpec(
            PbsAst.ReturnKind kind,
            PbsAst.TypeRef returnType,
            PbsAst.TypeRef resultErrorType) {
    }
}
