package p.studio.compiler.pbs.semantics;

import p.studio.compiler.pbs.ast.PbsAst;
import p.studio.compiler.pbs.semantics.PbsFlowSemanticSupport.Kind;
import p.studio.compiler.pbs.semantics.PbsFlowSemanticSupport.Model;
import p.studio.compiler.pbs.semantics.PbsFlowSemanticSupport.TupleField;
import p.studio.compiler.pbs.semantics.PbsFlowSemanticSupport.TypeView;

import java.util.ArrayList;
import java.util.List;

final class PbsFlowTypeOps {
    TypeView inferBinaryResult(final String operator, final TypeView left, final TypeView right) {
        if ("+".equals(operator) || "-".equals(operator) || "*".equals(operator) || "/".equals(operator) || "%".equals(operator)) {
            if (isFloat(left) || isFloat(right)) {
                return TypeView.floatType();
            }
            if (isInt(left) && isInt(right)) {
                return TypeView.intType();
            }
            return TypeView.unknown();
        }
        if ("==".equals(operator)
                || "!=".equals(operator)
                || "<".equals(operator)
                || "<=".equals(operator)
                || ">".equals(operator)
                || ">=".equals(operator)
                || "and".equals(operator)
                || "or".equals(operator)
                || "&&".equals(operator)
                || "||".equals(operator)) {
            return TypeView.bool();
        }
        return TypeView.unknown();
    }

    boolean compatible(final TypeView actual, final TypeView expected) {
        if (actual == null || expected == null) {
            return true;
        }
        if (actual.kind() == Kind.UNKNOWN || expected.kind() == Kind.UNKNOWN) {
            return true;
        }
        if (actual.kind() != expected.kind()) {
            return false;
        }
        return switch (actual.kind()) {
            case UNIT, INT, FLOAT, BOOL, STR -> true;
            case STRUCT, SERVICE, CONTRACT, CALLBACK, ENUM, ERROR, TYPE_REF, ASSET_NAMESPACE -> actual.name().equals(expected.name());
            case ADDRESSABLE -> true;
            case OPTIONAL -> compatible(actual.inner(), expected.inner());
            case RESULT -> compatible(actual.errorType(), expected.errorType()) && compatible(actual.inner(), expected.inner());
            case TUPLE -> tupleCompatible(actual, expected);
            default -> throw new IllegalStateException("Unexpected value: " + actual.kind());
        };
    }

    boolean inputCompatible(final List<TypeView> inputTypes, final TypeView argumentType) {
        if (inputTypes.isEmpty()) {
            return isUnit(argumentType) || argumentType.kind() == Kind.UNKNOWN;
        }
        if (inputTypes.size() == 1) {
            return compatible(argumentType, inputTypes.getFirst());
        }
        if (argumentType.kind() != Kind.TUPLE) {
            return false;
        }
        if (argumentType.tupleFields().size() != inputTypes.size()) {
            return false;
        }
        for (int i = 0; i < inputTypes.size(); i++) {
            if (!compatible(argumentType.tupleFields().get(i).type(), inputTypes.get(i))) {
                return false;
            }
        }
        return true;
    }

    TypeView callableReturnType(
            final PbsAst.ReturnKind returnKind,
            final PbsAst.TypeRef returnType,
            final PbsAst.TypeRef resultErrorType,
            final Model model) {
        return switch (returnKind) {
            case INFERRED_UNIT, EXPLICIT_UNIT -> TypeView.unit();
            case PLAIN -> collapseReturnPayload(typeFromTypeRef(returnType, model, null));
            case RESULT -> TypeView.result(
                    typeFromTypeRef(resultErrorType, model, null),
                    collapseReturnPayload(typeFromTypeRef(returnType, model, null)));
        };
    }

    TypeView typeFromTypeRef(
            final PbsAst.TypeRef typeRef,
            final Model model,
            final TypeView receiverType) {
        if (typeRef == null) {
            return TypeView.unit();
        }
        return switch (typeRef.kind()) {
            case UNIT -> TypeView.unit();
            case SELF -> receiverType == null ? TypeView.unknown() : receiverType;
            case SIMPLE -> simpleType(typeRef.name(), model);
            case OPTIONAL -> TypeView.optional(typeFromTypeRef(typeRef.inner(), model, receiverType));
            case GROUP -> typeFromTypeRef(typeRef.inner(), model, receiverType);
            case NAMED_TUPLE -> {
                final var fields = new ArrayList<TupleField>(typeRef.fields().size());
                for (final var field : typeRef.fields()) {
                    fields.add(new TupleField(field.label(), typeFromTypeRef(field.typeRef(), model, receiverType)));
                }
                yield TypeView.tuple(fields);
            }
            case ERROR -> TypeView.unknown();
        };
    }

    boolean isScalarComparable(final TypeView type) {
        return isBool(type) || isInt(type) || isFloat(type) || isStr(type);
    }

    boolean isUnit(final TypeView type) {
        return type.kind() == Kind.UNIT || type.kind() == Kind.UNKNOWN;
    }

    boolean isBool(final TypeView type) {
        return type.kind() == Kind.BOOL || type.kind() == Kind.UNKNOWN;
    }

    boolean isInt(final TypeView type) {
        return type.kind() == Kind.INT || type.kind() == Kind.UNKNOWN;
    }

    boolean isFloat(final TypeView type) {
        return type.kind() == Kind.FLOAT || type.kind() == Kind.UNKNOWN;
    }

    boolean isStr(final TypeView type) {
        return type.kind() == Kind.STR || type.kind() == Kind.UNKNOWN;
    }

    boolean isMaterializedValue(final TypeView type) {
        if (type == null) {
            return false;
        }
        return switch (type.kind()) {
            case UNKNOWN, UNIT, ASSET_NAMESPACE -> false;
            default -> true;
        };
    }

    private boolean tupleCompatible(final TypeView actual, final TypeView expected) {
        if (actual.tupleFields().size() != expected.tupleFields().size()) {
            return false;
        }
        for (int i = 0; i < actual.tupleFields().size(); i++) {
            if (!compatible(actual.tupleFields().get(i).type(), expected.tupleFields().get(i).type())) {
                return false;
            }
        }
        return true;
    }

    private TypeView collapseReturnPayload(final TypeView type) {
        if (type == null) {
            return TypeView.unit();
        }
        if (type.kind() == Kind.TUPLE && type.tupleFields().size() == 1) {
            return type.tupleFields().getFirst().type();
        }
        return type;
    }

    private TypeView simpleType(final String name, final Model model) {
        if ("int".equals(name)) {
            return TypeView.intType();
        }
        if ("float".equals(name)) {
            return TypeView.floatType();
        }
        if ("bool".equals(name)) {
            return TypeView.bool();
        }
        if ("str".equals(name)) {
            return TypeView.str();
        }
        if ("Addressable".equals(name)) {
            return TypeView.addressable(name);
        }
        if (model.structs.containsKey(name)) {
            return TypeView.struct(name);
        }
        if (model.services.containsKey(name)) {
            return TypeView.service(name);
        }
        if (model.contracts.containsKey(name)) {
            return TypeView.contract(name);
        }
        if (model.enums.containsKey(name)) {
            return TypeView.enumType(name);
        }
        if (model.errors.containsKey(name)) {
            return TypeView.error(name);
        }
        final var callbackSignature = model.callbacks.get(name);
        if (callbackSignature != null) {
            return TypeView.callback(name, callbackSignature.inputTypes(), callbackSignature.outputType());
        }
        return TypeView.unknown();
    }
}
