//go:generate ${GOBIN}/mockgen -source registry.go -destination registry_mock.go -package agentsidecar

package agentsidecar

import (
	"errors"
	"fmt"
	"time"

	"google.golang.org/grpc"

	apiproxyapi "github.com/tidbcloud/aisupport/packages/sdk-go/contracts/api/agent-sidecar/apiproxy/v1"
	projectapi "github.com/tidbcloud/aisupport/packages/sdk-go/contracts/api/agent-sidecar/project/v1"
	secretstoreapi "github.com/tidbcloud/aisupport/packages/sdk-go/contracts/api/agent-sidecar/secretstore/v1"
	grpcutil "github.com/tidbcloud/aisupport/packages/sdk-go/infra/utils/grpc"
	"github.com/tidbcloud/aisupport/services/agent-gateway/config"
)

type ProjectClient interface {
	projectapi.ProjectServiceClient
}

type APIProxyClient interface {
	apiproxyapi.APIProxyServiceClient
}

type SecretStoreClient interface {
	secretstoreapi.SecretStoreServiceClient
}

type Client struct {
	APIProxy    APIProxyClient
	SecretStore SecretStoreClient
	Project     ProjectClient
	Timeout     time.Duration
	conn        *grpc.ClientConn
}

type Registry struct {
	clients map[string]*Client
}

var newClientConnection = grpc.NewClient

func NewRegistry(conf *config.RootConfig) (*Registry, error) {
	registry := &Registry{clients: make(map[string]*Client)}
	for agentID, agent := range conf.AgentRuntime.Agents {
		if agent.Sidecar == nil {
			continue
		}
		conn, err := newClientConnection(agent.Sidecar.Endpoint, grpcutil.PresetDailOption().Default()...)
		if err != nil {
			_ = registry.Close()
			return nil, fmt.Errorf("create sidecar client for agent %q failed: %w", agentID, err)
		}
		registry.clients[agentID] = &Client{
			Project:     projectapi.NewProjectServiceClient(conn),
			APIProxy:    apiproxyapi.NewAPIProxyServiceClient(conn),
			SecretStore: secretstoreapi.NewSecretStoreServiceClient(conn),
			Timeout:     agent.Sidecar.Timeout,
			conn:        conn,
		}
	}
	return registry, nil
}

func (r *Registry) Get(agentID string) (*Client, bool) {
	client, ok := r.clients[agentID]
	return client, ok
}

func (r *Registry) Close() error {
	var errs []error
	for _, client := range r.clients {
		errs = append(errs, client.conn.Close())
	}
	return errors.Join(errs...)
}