/*
   Copyright The containerd 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 ttrpc

import (
	"context"
	"fmt"
	"net"
	"os"

	"github.com/containerd/containerd/v2/defaults"
	"github.com/containerd/containerd/v2/pkg/sys"
	"github.com/containerd/containerd/v2/plugins"
	"github.com/containerd/containerd/v2/plugins/server/internal"
	"github.com/containerd/plugin"
	"github.com/containerd/plugin/registry"
	"github.com/containerd/ttrpc"
)

type config struct {
	Address string `toml:"address"`
	UID     int    `toml:"uid"`
	GID     int    `toml:"gid"`
}

func (c *config) GetAddress() string {
	return c.Address
}

func (c *config) SetAddress(s string) {
	c.Address = s
}

func init() {
	// Keep the default the same for compatibility
	defaultAddress := defaults.DefaultAddress + ".ttrpc"
	registry.Register(&plugin.Registration{
		Type: plugins.ServerPlugin,
		ID:   "ttrpc",
		Requires: []plugin.Type{
			plugins.TTRPCPlugin,
			plugins.GRPCPlugin,
		},
		Config: &config{
			Address: defaultAddress,
			UID:     os.Geteuid(),
			GID:     os.Getegid(),
		},
		InitFn: func(ic *plugin.InitContext) (any, error) {
			c := ic.Config.(*config)
			if c.Address == "" {
				return nil, plugin.ErrSkipPlugin
			}

			s, err := newTTRPCServer()
			if err != nil {
				return nil, err
			}
			// ttrpcService allows TTRPC services to be registered with the underlying server
			type ttrpcService interface {
				RegisterTTRPC(*ttrpc.Server) error
			}
			// Iterate plugins directly rather than using GetByType, which
			// short-circuits on the first plugin that failed to initialize.
			// Plugins that failed (e.g. CRI under rootless) have already been
			// logged and should not prevent the TTRPC server from starting.
			var hasService bool
			for _, p := range ic.Plugins().GetAll() {
				if p.Registration.Type != plugins.TTRPCPlugin && p.Registration.Type != plugins.GRPCPlugin {
					continue
				}
				instance, err := p.Instance()
				if err != nil {
					continue
				}
				if gs, ok := instance.(ttrpcService); ok {
					if err := gs.RegisterTTRPC(s); err != nil {
						return nil, fmt.Errorf("failed to register ttrpc service: %w", err)
					}
					hasService = true
				}
			}
			if !hasService {
				return nil, fmt.Errorf("no ttrpc services configured: %w", plugin.ErrSkipPlugin)
			}
			return server{
				Server: s,
				config: *c,
			}, nil
		},
	})
}

type server struct {
	*ttrpc.Server
	config config
}

func (s server) Start(ctx context.Context) error {
	// setup the ttrpc endpoint
	tl, err := sys.GetLocalListener(s.config.Address, s.config.UID, s.config.GID)
	if err != nil {
		return fmt.Errorf("failed to get listener for main ttrpc endpoint: %w", err)
	}

	internal.Serve(ctx, tl, func(l net.Listener) error {
		return s.Serve(context.WithoutCancel(ctx), l)
	})

	return nil
}