/*
 * Copyright 2021 CloudWeGo Authors
 *
 * Licensed 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 remotecli

import (
	"context"
	"net"
	"sync"
	"time"

	"github.com/cloudwego/kitex/pkg/kerrors"
	"github.com/cloudwego/kitex/pkg/remote"
	"github.com/cloudwego/kitex/pkg/rpcinfo"
	"github.com/cloudwego/kitex/pkg/stats"
)

var connWrapperPool sync.Pool

func init() {
	connWrapperPool.New = newConnWrapper
}

var _ ConnReleaser = &ConnWrapper{}

// ConnReleaser helps to release the raw connection.
type ConnReleaser interface {
	ReleaseConn(err error, ri rpcinfo.RPCInfo)
}

// ConnWrapper wraps a connection.
type ConnWrapper struct {
	connPool remote.ConnPool
	conn     net.Conn
}

// NewConnWrapper returns a new ConnWrapper using the given connPool and logger.
func NewConnWrapper(connPool remote.ConnPool) *ConnWrapper {
	cm := connWrapperPool.Get().(*ConnWrapper)
	cm.connPool = connPool
	return cm
}

// GetConn returns a connection using the given Dialer and RPCInfo.
func (cm *ConnWrapper) GetConn(ctx context.Context, d remote.Dialer, ri rpcinfo.RPCInfo) (net.Conn, error) {
	configs := ri.Config()
	timeout := configs.ConnectTimeout()
	if cm.connPool != nil {
		conn, err := cm.getConnWithPool(ctx, cm.connPool, d, timeout, ri)
		if err != nil {
			return nil, err
		}
		cm.conn = conn
		if raw, ok := conn.(remote.RawConn); ok {
			rawConn := raw.RawConn()
			return rawConn, nil
		}
		return cm.conn, nil
	}

	var err error
	cm.conn, err = cm.getConnWithDialer(ctx, d, timeout, ri)
	return cm.conn, err
}

// ReleaseConn should notice that ri may nil when oneway
// TODO duplicate release may cause problem?
func (cm *ConnWrapper) ReleaseConn(err error, ri rpcinfo.RPCInfo) {
	if cm.conn == nil {
		return
	}
	if cm.connPool != nil {
		if err == nil {
			_, ok := ri.To().Tag(rpcinfo.ConnResetTag)
			if ok || ri.Config().InteractionMode() == rpcinfo.Oneway {
				cm.connPool.Discard(cm.conn)
			} else {
				cm.connPool.Put(cm.conn)
			}
		} else {
			cm.connPool.Discard(cm.conn)
		}
	} else {
		cm.conn.Close()
	}

	cm.zero()
	connWrapperPool.Put(cm)
}

func newConnWrapper() interface{} {
	return &ConnWrapper{}
}

func (cm *ConnWrapper) zero() {
	cm.connPool = nil
	cm.conn = nil
}

func (cm *ConnWrapper) getConnWithPool(ctx context.Context, cp remote.ConnPool, d remote.Dialer,
	timeout time.Duration, ri rpcinfo.RPCInfo,
) (net.Conn, error) {
	addr := ri.To().Address()
	if addr == nil {
		return nil, kerrors.ErrNoDestAddress
	}
	opt := remote.ConnOption{Dialer: d, ConnectTimeout: timeout}
	ri.Stats().Record(ctx, stats.ClientConnStart, stats.StatusInfo, "")
	conn, err := cp.Get(ctx, addr.Network(), addr.String(), opt)
	if err != nil {
		ri.Stats().Record(ctx, stats.ClientConnFinish, stats.StatusError, err.Error())
		return nil, kerrors.ErrGetConnection.WithCause(err)
	}
	ri.Stats().Record(ctx, stats.ClientConnFinish, stats.StatusInfo, "")
	return conn, nil
}

func (cm *ConnWrapper) getConnWithDialer(ctx context.Context, d remote.Dialer,
	timeout time.Duration, ri rpcinfo.RPCInfo,
) (net.Conn, error) {
	addr := ri.To().Address()
	if addr == nil {
		return nil, kerrors.ErrNoDestAddress
	}

	ri.Stats().Record(ctx, stats.ClientConnStart, stats.StatusInfo, "")
	conn, err := d.DialTimeout(addr.Network(), addr.String(), timeout)
	if err != nil {
		ri.Stats().Record(ctx, stats.ClientConnFinish, stats.StatusError, err.Error())
		return nil, kerrors.ErrGetConnection.WithCause(err)
	}
	ri.Stats().Record(ctx, stats.ClientConnFinish, stats.StatusInfo, "")
	return conn, nil
}
