/*
 * Copyright 2024 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 ttstream

import (
	"context"
	"errors"
	"runtime"
	"sync"
	"sync/atomic"
	"time"

	"github.com/cloudwego/netpoll"

	"github.com/cloudwego/kitex/pkg/gofunc"
)

var DefaultMuxConnConfig = MuxConnConfig{
	PoolSize:       runtime.GOMAXPROCS(0),
	MaxIdleTimeout: time.Minute,
}

type MuxConnConfig struct {
	PoolSize       int
	MaxIdleTimeout time.Duration
}

var _ transPool = (*muxConnTransPool)(nil)

type muxConnTransList struct {
	L          sync.RWMutex
	size       int
	cursor     uint32
	transports []*clientTransport
	pool       *muxConnTransPool
}

func newMuxConnTransList(size int, pool *muxConnTransPool) *muxConnTransList {
	tl := new(muxConnTransList)
	if size <= 0 {
		size = runtime.GOMAXPROCS(0)
	}
	tl.size = size
	tl.transports = make([]*clientTransport, size)
	tl.pool = pool
	return tl
}

func (tl *muxConnTransList) Close() {
	tl.L.Lock()
	for i, t := range tl.transports {
		if t == nil {
			continue
		}
		_ = t.Close(nil)
		tl.transports[i] = nil
	}
	tl.L.Unlock()
}

func (tl *muxConnTransList) Get(network, addr string) (*clientTransport, error) {
	// fast path
	idx := atomic.AddUint32(&tl.cursor, 1) % uint32(tl.size)
	tl.L.RLock()
	trans := tl.transports[idx]
	tl.L.RUnlock()
	if trans != nil {
		if trans.IsActive() {
			return trans, nil
		}
		_ = trans.Close(nil)
	}

	// slow path
	tl.L.Lock()
	trans = tl.transports[idx]
	if trans != nil && trans.IsActive() {
		// another goroutine already create the new transport
		tl.L.Unlock()
		return trans, nil
	}
	// it may create more than tl.size transport if multi client try to get transport concurrently
	conn, err := dialer.DialConnection(network, addr, time.Second)
	if err != nil {
		return nil, err
	}
	trans = newClientTransport(conn, tl.pool)
	_ = conn.AddCloseCallback(func(connection netpoll.Connection) error {
		// peer close
		_ = trans.Close(errTransport.newBuilder().withCause(errors.New("connection closed by peer")))
		return nil
	})
	tl.transports[idx] = trans
	tl.L.Unlock()

	return trans, nil
}

func newMuxConnTransPool(config MuxConnConfig) transPool {
	t := new(muxConnTransPool)
	t.config = config
	return t
}

type muxConnTransPool struct {
	config      MuxConnConfig
	pool        sync.Map // addr:*muxConnTransList
	activity    sync.Map // addr:lastActive
	cleanerOnce int32
}

func (p *muxConnTransPool) Get(network, addr string) (trans *clientTransport, err error) {
	v, ok := p.pool.Load(addr)
	if !ok {
		// multi concurrent Get should get the same TransList object
		v, _ = p.pool.LoadOrStore(addr, newMuxConnTransList(p.config.PoolSize, p))
	}
	return v.(*muxConnTransList).Get(network, addr)
}

func (p *muxConnTransPool) Put(trans *clientTransport) {
	p.activity.Store(trans.conn.RemoteAddr().String(), time.Now())

	idleTimeout := p.config.MaxIdleTimeout
	if idleTimeout == 0 {
		return
	}
	if !atomic.CompareAndSwapInt32(&p.cleanerOnce, 0, 1) {
		return
	}
	// start cleaning background goroutine
	gofunc.RecoverGoFuncWithInfo(context.Background(), func() {
		for {
			now := time.Now()
			count := 0
			p.activity.Range(func(key, value interface{}) bool {
				addr := key.(string)
				count++
				lastActive := value.(time.Time)
				idleTime := now.Sub(lastActive)
				if lastActive.IsZero() || idleTime < 0 || now.Sub(lastActive) < idleTimeout {
					return true
				}
				// clean transport
				v, _ := p.pool.Load(addr)
				if v == nil {
					return true
				}
				p.pool.Delete(addr)
				p.activity.Delete(addr)
				return true
			})
			time.Sleep(idleTimeout)
		}
	}, gofunc.NewBasicInfo("", trans.Addr().String()))
}
