package com.tsi.comm.server.tcp.handler;

import com.tsi.comm.server.xnet.NettyUtils;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.ChannelPromise;
import io.netty.handler.logging.LogLevel;
import io.netty.handler.logging.LoggingHandler;
import org.slf4j.MDC;

import java.net.SocketAddress;

public class MdcLoggingHandler extends LoggingHandler {

    public MdcLoggingHandler(LogLevel level) {
        super(level);
    }

    private void setMdc(ChannelHandlerContext ctx) {
        final String remoteIpAddress = NettyUtils.getRemoteIpAddress(ctx.channel());
        MDC.put("id", remoteIpAddress);
    }

    private void clearMdc() {
        MDC.remove("id");
    }

    @Override
    public void channelRegistered(ChannelHandlerContext ctx) throws Exception {
        setMdc(ctx);
        try {
            super.channelRegistered(ctx);
        } finally {
            clearMdc();
        }
    }

    @Override
    public void channelUnregistered(ChannelHandlerContext ctx) throws Exception {
        setMdc(ctx);
        try {
            super.channelUnregistered(ctx);
        } finally {
            clearMdc();
        }
    }

    @Override
    public void channelActive(ChannelHandlerContext ctx) throws Exception {
        setMdc(ctx);
        try {
            super.channelActive(ctx);
        } finally {
            clearMdc();
        }
    }

    @Override
    public void channelInactive(ChannelHandlerContext ctx) throws Exception {
        setMdc(ctx);
        try {
            super.channelInactive(ctx);
        } finally {
            clearMdc();
        }
    }

    @Override
    public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception {
        setMdc(ctx);
        try {
            super.channelRead(ctx, msg);
        } finally {
            clearMdc();
        }
    }

    @Override
    public void channelReadComplete(ChannelHandlerContext ctx) throws Exception {
        setMdc(ctx);
        try {
            super.channelReadComplete(ctx);
        } finally {
            clearMdc();
        }
    }

    @Override
    public void userEventTriggered(ChannelHandlerContext ctx, Object evt) throws Exception {
        setMdc(ctx);
        try {
            super.userEventTriggered(ctx, evt);
        } finally {
            clearMdc();
        }
    }

    @Override
    public void channelWritabilityChanged(ChannelHandlerContext ctx) throws Exception {
        setMdc(ctx);
        try {
            super.channelWritabilityChanged(ctx);
        } finally {
            clearMdc();
        }
    }

    @Override
    public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws Exception {
        setMdc(ctx);
        try {
            super.exceptionCaught(ctx, cause);
        } finally {
            clearMdc();
        }
    }

    @Override
    public void bind(ChannelHandlerContext ctx, SocketAddress localAddress, ChannelPromise promise) throws Exception {
        setMdc(ctx);
        try {
            super.bind(ctx, localAddress, promise);
        } finally {
            clearMdc();
        }
    }

    @Override
    public void connect(ChannelHandlerContext ctx, SocketAddress remoteAddress, SocketAddress localAddress, ChannelPromise promise) throws Exception {
        setMdc(ctx);
        try {
            super.connect(ctx, remoteAddress, localAddress, promise);
        } finally {
            clearMdc();
        }
    }

    @Override
    public void disconnect(ChannelHandlerContext ctx, ChannelPromise promise) throws Exception {
        setMdc(ctx);
        try {
            super.disconnect(ctx, promise);
        } finally {
            clearMdc();
        }
    }

    @Override
    public void close(ChannelHandlerContext ctx, ChannelPromise promise) throws Exception {
        setMdc(ctx);
        try {
            super.close(ctx, promise);
        } finally {
            clearMdc();
        }
    }

    @Override
    public void deregister(ChannelHandlerContext ctx, ChannelPromise promise) throws Exception {
        setMdc(ctx);
        try {
            super.deregister(ctx, promise);
        } finally {
            clearMdc();
        }
    }

    @Override
    public void read(ChannelHandlerContext ctx) throws Exception {
        setMdc(ctx);
        try {
            super.read(ctx);
        } finally {
            clearMdc();
        }
    }

    @Override
    public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise promise) throws Exception {
        setMdc(ctx);
        try {
            super.write(ctx, msg, promise);
        } finally {
            clearMdc();
        }
    }

    @Override
    public void flush(ChannelHandlerContext ctx) throws Exception {
        setMdc(ctx);
        try {
            super.flush(ctx);
        } finally {
            clearMdc();
        }
    }
}
