package grpctool

import (
	"context"
	"errors"
	"fmt"
	"log/slog"
	"net/url"
	"sync"
	"time"

	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/errz"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/fieldz"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/logz"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/credentials/insecure"
	"k8s.io/utils/clock"
)

const (
	evictIdleConnAfter = 1 * time.Hour
)

type PoolConn interface {
	grpc.ClientConnInterface
	Done()
}

type Target interface {
	comparable
	String() string
}

type PoolInterface[T Target] interface {
	Dial(target T) (PoolConn, error)
	Shutdown(deadline time.Duration)
}

type Pool[T Target] struct {
	mu            sync.Mutex
	log           *slog.Logger
	errRep        errz.ErrReporter
	newConnection func(T) (string, []grpc.DialOption, error)
	conns         map[T]*connHolder // target -> conn
	clk           clock.PassiveClock
	closed        bool
	closePoolChan chan struct{}
}

func NewPool[T Target](log *slog.Logger, errRep errz.ErrReporter, newConnection func(T) (string, []grpc.DialOption, error)) *Pool[T] {
	return &Pool[T]{
		log:           log,
		errRep:        errRep,
		newConnection: newConnection,
		conns:         map[T]*connHolder{},
		clk:           clock.RealClock{},
		closePoolChan: make(chan struct{}),
	}
}

func (p *Pool[T]) Dial(target T) (PoolConn, error) {
	p.mu.Lock()
	defer p.mu.Unlock()
	if p.closed {
		return nil, errors.New("pool has been closed, cannot dial up new connections")
	}
	conn := p.conns[target]
	if conn == nil {
		clientTarget, opts, err := p.newConnection(target)
		if err != nil {
			return nil, fmt.Errorf("pool new conn: %w", err)
		}
		grpcConn, err := grpc.NewClient(clientTarget, opts...)
		if err != nil {
			return nil, fmt.Errorf("pool gRPC dial: %w", err)
		}
		conn = &connHolder{
			ClientConn: grpcConn,
		}
		p.conns[target] = conn
	}
	conn.numUsers++
	return &poolConn{
		connHolder: conn,
		done:       p.connDone,
	}, nil
}

func (p *Pool[T]) usersExistLocked() bool {
	for _, conn := range p.conns {
		if conn.numUsers != 0 {
			return true
		}
	}
	return false
}

func (p *Pool[T]) Shutdown(deadline time.Duration) {
	p.mu.Lock()
	defer p.mu.Unlock()
	p.closed = true

	if p.usersExistLocked() {
		p.mu.Unlock()
		select {
		case <-time.After(deadline):
		case <-p.closePoolChan:
		}
		p.mu.Lock()
	}

	for targetURL, conn := range p.conns {
		log := p.log.With(logz.URL(targetURL.String()))
		if conn.numUsers > 0 {
			p.errRep.HandleProcessingError(context.Background(), log, "Closing pooled connection",
				errors.New("connection is in use"), fieldz.New("num_users", conn.numUsers))
		}
		err := conn.Close()
		if err != nil {
			p.errRep.HandleProcessingError(context.Background(), log, "Error closing pooled connection", err)
		} else {
			log.Debug("Closed pooled connection")
		}
		conn.numUsers = 0
	}
	clear(p.conns)
	select {
	case <-p.closePoolChan:
	default:
		close(p.closePoolChan)
	}
}

func (p *Pool[T]) connDone(conn *connHolder) {
	p.mu.Lock()
	defer p.mu.Unlock()
	// A call to Shutdown has occurred, no need to handle closing connection
	if conn.numUsers == 0 {
		return
	}
	conn.numUsers--
	conn.lastUsed = p.clk.Now()
	p.runGCLocked()
}

func (p *Pool[T]) runGCLocked() {
	expireAt := p.clk.Now().Add(-evictIdleConnAfter)
	for targetURL, conn := range p.conns {
		if conn.numUsers == 0 && conn.lastUsed.Before(expireAt) {
			delete(p.conns, targetURL)
			err := conn.Close()
			if err != nil {
				p.errRep.HandleProcessingError(
					context.Background(),
					p.log.With(logz.URL(targetURL.String())),
					"Error closing idle pooled connection", err,
				)
			} else {
				p.log.Debug("Closed idle pooled connection", logz.URL(targetURL.String()))
			}
		}
	}
	if p.closed && !p.usersExistLocked() {
		close(p.closePoolChan)
	}
}

type connHolder struct {
	*grpc.ClientConn
	lastUsed time.Time
	numUsers int32 // protected by mutex
}

type poolConn struct {
	*connHolder
	done func(conn *connHolder)
}

func (c *poolConn) Done() {
	if c.done == nil {
		panic("pooled connection Done() called more than once")
	}
	done := c.done
	c.done = nil
	done(c.connHolder)
}

type URLTarget string

func (t URLTarget) String() string {
	return string(t)
}

func DefaultNewConnection(tlsCreds credentials.TransportCredentials, dialOpts ...grpc.DialOption) func(URLTarget) (string, []grpc.DialOption, error) {
	return func(targetURL URLTarget) (string, []grpc.DialOption, error) {
		u, err := url.Parse(string(targetURL))
		if err != nil {
			return "", nil, err
		}
		var creds credentials.TransportCredentials
		var target string
		switch u.Scheme {
		case "grpc":
			// See https://github.com/grpc/grpc/blob/master/doc/naming.md.
			target = "dns:" + HostWithPort(u)
			creds = insecure.NewCredentials()
		case "grpcs":
			// See https://github.com/grpc/grpc/blob/master/doc/naming.md.
			target = "dns:" + HostWithPort(u)
			creds = tlsCreds
		case "unix":
			// See https://github.com/grpc/grpc/blob/master/doc/naming.md.
			target = string(targetURL)
			creds = insecure.NewCredentials()
		default:
			return "", nil, fmt.Errorf("unsupported URL scheme in %s", targetURL)
		}
		l := len(dialOpts)
		opts := append(
			dialOpts[:l:l], // set cap=len to ensure append() will copy the slice
			grpc.WithTransportCredentials(creds),
		)
		return target, opts, nil
	}
}
