package grpc_conn

import (
	"common/log"
	"fmt"
	"google.golang.org/grpc"
	"math/rand"
)

type GrpcConnectionMgr struct {
	poolM map[int64]*GrpcConnection
	poolS []*GrpcConnection
}

func NewGrpcConnectionMgr() *GrpcConnectionMgr {
	return &GrpcConnectionMgr{
		poolM: make(map[int64]*GrpcConnection),
		poolS: make([]*GrpcConnection, 0),
	}
}

func (p *GrpcConnectionMgr) Store(sid int64, addr string) {
	pool, err := NewGrpcConnection(sid, addr)
	if err != nil {
		log.Errorf("create grpc err: %v, sid: %v, addr: %v", err, sid, addr)
		return
	}
	p.poolM[sid] = pool
	p.poolS = append(p.poolS, pool)
}

func (p *GrpcConnectionMgr) Delete(sid int64) int {
	delete(p.poolM, sid)
	for i, pool := range p.poolS {
		if pool.sid == sid {
			p.poolS = append(p.poolS[:i], p.poolS[i+1:]...)
			break
		}
	}
	return len(p.poolS)
}

func (p *GrpcConnectionMgr) Load(sid ...int64) (*grpc.ClientConn, error) {
	var pool *GrpcConnection
	if len(sid) > 0 && sid[0] > 0 {
		pool = p.poolM[sid[0]]
	} else {
		pool = p.poolS[rand.Intn(len(p.poolS))]
	}
	if pool == nil {
		return nil, fmt.Errorf("cannot find connection")
	}
	return pool.GetConnection(), nil
}

func (p *GrpcConnectionMgr) LoadAll() map[int64]*grpc.ClientConn {
	sidM := make(map[int64]*grpc.ClientConn)
	for sid, pool := range p.poolM {
		sidM[sid] = pool.GetConnection()
	}
	return sidM
}