/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.dubbo.rpc.protocol.dubbo;

import org.apache.dubbo.common.Version;
import org.apache.dubbo.common.io.Bytes;
import org.apache.dubbo.common.io.UnsafeByteArrayInputStream;
import org.apache.dubbo.common.logger.ErrorTypeAwareLogger;
import org.apache.dubbo.common.logger.LoggerFactory;
import org.apache.dubbo.common.serialize.ObjectInput;
import org.apache.dubbo.common.serialize.ObjectOutput;
import org.apache.dubbo.common.serialize.Serialization;
import org.apache.dubbo.common.threadpool.manager.ExecutorRepository;
import org.apache.dubbo.common.utils.StringUtils;
import org.apache.dubbo.remoting.Channel;
import org.apache.dubbo.remoting.exchange.HeartBeatRequest;
import org.apache.dubbo.remoting.exchange.HeartBeatResponse;
import org.apache.dubbo.remoting.exchange.Request;
import org.apache.dubbo.remoting.exchange.Response;
import org.apache.dubbo.remoting.exchange.codec.ExchangeCodec;
import org.apache.dubbo.remoting.transport.CodecSupport;
import org.apache.dubbo.rpc.AppResponse;
import org.apache.dubbo.rpc.Invocation;
import org.apache.dubbo.rpc.Result;
import org.apache.dubbo.rpc.RpcInvocation;
import org.apache.dubbo.rpc.model.FrameworkModel;

import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.util.Optional;
import java.util.concurrent.atomic.AtomicBoolean;

import static org.apache.dubbo.common.constants.CommonConstants.BYTE_ACCESSOR_KEY;
import static org.apache.dubbo.common.constants.CommonConstants.DUBBO_VERSION_KEY;
import static org.apache.dubbo.common.constants.CommonConstants.EXECUTOR_MANAGEMENT_MODE_ISOLATION;
import static org.apache.dubbo.common.constants.CommonConstants.INTERFACE_KEY;
import static org.apache.dubbo.common.constants.CommonConstants.PATH_KEY;
import static org.apache.dubbo.common.constants.CommonConstants.VERSION_KEY;
import static org.apache.dubbo.common.constants.LoggerCodeConstants.PROTOCOL_FAILED_DECODE;
import static org.apache.dubbo.rpc.protocol.dubbo.Constants.DEFAULT_DECODE_IN_IO_THREAD;

/**
 * Dubbo codec.
 */
public class DubboCodec extends ExchangeCodec {

    public static final String NAME = "dubbo";
    public static final String DUBBO_VERSION = Version.getProtocolVersion();
    public static final byte RESPONSE_WITH_EXCEPTION = 0;
    public static final byte RESPONSE_VALUE = 1;
    public static final byte RESPONSE_NULL_VALUE = 2;
    public static final byte RESPONSE_WITH_EXCEPTION_WITH_ATTACHMENTS = 3;
    public static final byte RESPONSE_VALUE_WITH_ATTACHMENTS = 4;
    public static final byte RESPONSE_NULL_VALUE_WITH_ATTACHMENTS = 5;
    public static final Object[] EMPTY_OBJECT_ARRAY = new Object[0];
    public static final Class<?>[] EMPTY_CLASS_ARRAY = new Class<?>[0];
    private static final ErrorTypeAwareLogger log = LoggerFactory.getErrorTypeAwareLogger(DubboCodec.class);

    private static final AtomicBoolean decodeInUserThreadLogged = new AtomicBoolean(false);
    private final CallbackServiceCodec callbackServiceCodec;
    private final FrameworkModel frameworkModel;
    private final ByteAccessor customByteAccessor;
    private static final String DECODE_IN_IO_THREAD_KEY = "decode.in.io.thread";

    public DubboCodec(FrameworkModel frameworkModel) {
        this.frameworkModel = frameworkModel;
        callbackServiceCodec = new CallbackServiceCodec(frameworkModel);
        customByteAccessor = Optional.ofNullable(System.getProperty(BYTE_ACCESSOR_KEY))
                .filter(StringUtils::isNotBlank)
                .map(key ->
                        frameworkModel.getExtensionLoader(ByteAccessor.class).getExtension(key))
                .orElse(null);
    }

    @Override
    protected Object decodeBody(Channel channel, InputStream is, byte[] header) throws IOException {
        byte flag = header[2], proto = (byte) (flag & SERIALIZATION_MASK);
        // get request id.
        long id = Bytes.bytes2long(header, 4);
        if ((flag & FLAG_REQUEST) == 0) {
            // decode response.
            Response res = new Response(id);
            if ((flag & FLAG_EVENT) != 0) {
                res.setEvent(true);
            }
            // get status.
            byte status = header[3];
            res.setStatus(status);
            try {
                if (status == Response.OK) {
                    Object data;
                    if (res.isEvent()) {
                        byte[] eventPayload = CodecSupport.getPayload(is);
                        if (CodecSupport.isHeartBeat(eventPayload, proto)) {
                            // heart beat response data is always null;
                            data = null;
                        } else {
                            ObjectInput in = CodecSupport.deserialize(
                                    channel.getUrl(), new ByteArrayInputStream(eventPayload), proto);
                            data = decodeEventData(channel, in, eventPayload);
                        }
                    } else {
                        DecodeableRpcResult result;
                        Invocation inv = (Invocation) getRequestData(channel, res, id);
                        if (channel.getUrl().getParameter(DECODE_IN_IO_THREAD_KEY, DEFAULT_DECODE_IN_IO_THREAD)) {
                            if (customByteAccessor != null) {
                                result = customByteAccessor.getRpcResult(
                                        channel, res, new UnsafeByteArrayInputStream(readMessageData(is)), inv, proto);
                            } else {
                                result = new DecodeableRpcResult(
                                        channel, res, new UnsafeByteArrayInputStream(readMessageData(is)), inv, proto);
                            }
                            result.decode();
                        } else {
                            if (customByteAccessor != null) {
                                result = customByteAccessor.getRpcResult(
                                        channel, res, new UnsafeByteArrayInputStream(readMessageData(is)), inv, proto);
                            } else {
                                result = new DecodeableRpcResult(
                                        channel, res, new UnsafeByteArrayInputStream(readMessageData(is)), inv, proto);
                            }
                        }
                        data = result;
                    }
                    res.setResult(data);
                } else {
                    ObjectInput in = CodecSupport.deserialize(channel.getUrl(), is, proto);
                    res.setErrorMessage(in.readUTF());
                }
            } catch (Throwable t) {
                if (log.isWarnEnabled()) {
                    log.warn(PROTOCOL_FAILED_DECODE, "", "", "Decode response failed: " + t.getMessage(), t);
                }
                res.setStatus(Response.CLIENT_ERROR);
                res.setErrorMessage(StringUtils.toString(t));
            }
            return res;
        } else {
            // decode request.
            Request req;
            try {
                Object data;
                if ((flag & FLAG_EVENT) != 0) {
                    byte[] eventPayload = CodecSupport.getPayload(is);
                    if (CodecSupport.isHeartBeat(eventPayload, proto)) {
                        // heart beat response data is always null;
                        req = new HeartBeatRequest(id);
                        req.setVersion(Version.getProtocolVersion());
                        req.setTwoWay((flag & FLAG_TWOWAY) != 0);
                        ((HeartBeatRequest) req).setProto(proto);
                        data = null;
                    } else {
                        req = new Request(id);
                        req.setVersion(Version.getProtocolVersion());
                        req.setTwoWay((flag & FLAG_TWOWAY) != 0);

                        ObjectInput in = CodecSupport.deserialize(
                                channel.getUrl(), new ByteArrayInputStream(eventPayload), proto);
                        data = decodeEventData(channel, in, eventPayload);
                    }
                    req.setEvent(true);
                } else {
                    req = new Request(id);
                    req.setVersion(Version.getProtocolVersion());
                    req.setTwoWay((flag & FLAG_TWOWAY) != 0);

                    // get data length.
                    int len = Bytes.bytes2int(header, 12);
                    req.setPayload(len);

                    DecodeableRpcInvocation inv;
                    if (isDecodeDataInIoThread(channel)) {
                        if (customByteAccessor != null) {
                            inv = customByteAccessor.getRpcInvocation(
                                    channel, req, new UnsafeByteArrayInputStream(readMessageData(is)), proto);
                        } else {
                            inv = new DecodeableRpcInvocation(
                                    frameworkModel,
                                    channel,
                                    req,
                                    new UnsafeByteArrayInputStream(readMessageData(is)),
                                    proto);
                        }
                        inv.decode();
                    } else {
                        if (customByteAccessor != null) {
                            inv = customByteAccessor.getRpcInvocation(
                                    channel, req, new UnsafeByteArrayInputStream(readMessageData(is)), proto);
                        } else {
                            inv = new DecodeableRpcInvocation(
                                    frameworkModel,
                                    channel,
                                    req,
                                    new UnsafeByteArrayInputStream(readMessageData(is)),
                                    proto);
                        }
                    }
                    data = inv;
                }
                req.setData(data);
            } catch (Throwable t) {
                if (log.isWarnEnabled()) {
                    log.warn(PROTOCOL_FAILED_DECODE, "", "", "Decode request failed: " + t.getMessage(), t);
                }
                // bad request
                req = new HeartBeatRequest(id);
                req.setBroken(true);
                req.setData(t);
            }

            return req;
        }
    }

    private boolean isDecodeDataInIoThread(Channel channel) {
        Object obj = channel.getAttribute(DECODE_IN_IO_THREAD_KEY);
        if (obj instanceof Boolean) {
            return (Boolean) obj;
        }

        String mode = ExecutorRepository.getMode(channel.getUrl().getOrDefaultApplicationModel());
        boolean isIsolated = EXECUTOR_MANAGEMENT_MODE_ISOLATION.equals(mode);

        if (isIsolated && !decodeInUserThreadLogged.compareAndSet(false, true)) {
            channel.setAttribute(DECODE_IN_IO_THREAD_KEY, true);
            return true;
        }

        boolean decodeDataInIoThread =
                channel.getUrl().getParameter(DECODE_IN_IO_THREAD_KEY, DEFAULT_DECODE_IN_IO_THREAD);
        if (isIsolated && !decodeDataInIoThread) {
            log.info("Because thread pool isolation is enabled on the dubbo protocol, the body can only be decoded "
                    + "on the io thread, and the parameter[" + DECODE_IN_IO_THREAD_KEY + "] will be ignored");
            // Why? because obtaining the isolated thread pool requires the serviceKey of the service,
            // and this part must be decoded before it can be obtained (more see DubboExecutorSupport)
            channel.setAttribute(DECODE_IN_IO_THREAD_KEY, true);
            return true;
        }
        channel.setAttribute(DECODE_IN_IO_THREAD_KEY, decodeDataInIoThread);
        return decodeDataInIoThread;
    }

    private byte[] readMessageData(InputStream is) throws IOException {
        if (is.available() > 0) {
            byte[] result = new byte[is.available()];
            is.read(result);
            return result;
        }
        return new byte[] {};
    }

    @Override
    protected void encodeRequestData(Channel channel, ObjectOutput out, Object data) throws IOException {
        encodeRequestData(channel, out, data, DUBBO_VERSION);
    }

    @Override
    protected void encodeResponseData(Channel channel, ObjectOutput out, Object data) throws IOException {
        encodeResponseData(channel, out, data, DUBBO_VERSION);
    }

    @Override
    protected void encodeRequestData(Channel channel, ObjectOutput out, Object data, String version)
            throws IOException {
        RpcInvocation inv = (RpcInvocation) data;

        out.writeUTF(version);
        // https://github.com/apache/dubbo/issues/6138
        String serviceName = inv.getAttachment(INTERFACE_KEY);
        if (serviceName == null) {
            serviceName = inv.getAttachment(PATH_KEY);
        }
        out.writeUTF(serviceName);
        out.writeUTF(inv.getAttachment(VERSION_KEY));

        out.writeUTF(inv.getMethodName());
        out.writeUTF(inv.getParameterTypesDesc());
        Object[] args = inv.getArguments();
        if (args != null) {
            for (int i = 0; i < args.length; i++) {
                out.writeObject(callbackServiceCodec.encodeInvocationArgument(channel, inv, i));
            }
        }
        out.writeAttachments(inv.getObjectAttachments());
    }

    @Override
    protected void encodeResponseData(Channel channel, ObjectOutput out, Object data, String version)
            throws IOException {
        Result result = (Result) data;
        // currently, the version value in Response records the version of Request
        boolean attach = Version.isSupportResponseAttachment(version);
        Throwable th = result.getException();
        if (th == null) {
            Object ret = result.getValue();
            if (ret == null) {
                out.writeByte(attach ? RESPONSE_NULL_VALUE_WITH_ATTACHMENTS : RESPONSE_NULL_VALUE);
            } else {
                out.writeByte(attach ? RESPONSE_VALUE_WITH_ATTACHMENTS : RESPONSE_VALUE);
                out.writeObject(ret);
            }
        } else {
            out.writeByte(attach ? RESPONSE_WITH_EXCEPTION_WITH_ATTACHMENTS : RESPONSE_WITH_EXCEPTION);
            out.writeThrowable(th);
        }

        if (attach) {
            // returns current version of Response to consumer side.
            result.getObjectAttachments().put(DUBBO_VERSION_KEY, Version.getProtocolVersion());
            out.writeAttachments(result.getObjectAttachments());
        }
    }

    @Override
    protected Serialization getSerialization(Channel channel, Request req) {
        if (!(req.getData() instanceof Invocation)) {
            return super.getSerialization(channel, req);
        }
        return DubboCodecSupport.getRequestSerialization(channel.getUrl(), (Invocation) req.getData());
    }

    @Override
    protected Serialization getSerialization(Channel channel, Response res) {
        if (res instanceof HeartBeatResponse) {
            return CodecSupport.getSerializationById(((HeartBeatResponse) res).getProto());
        }
        if (!(res.getResult() instanceof AppResponse)) {
            return super.getSerialization(channel, res);
        }
        return DubboCodecSupport.getResponseSerialization(channel.getUrl(), (AppResponse) res.getResult());
    }
}
