// Copyright (c) 2026 Huawei Technologies Co., Ltd.
// openFuyao is licensed under Mulan PSL v2.
// You can use this software according to the terms and conditions of the Mulan PSL v2.
// You may obtain a copy of Mulan PSL v2 at:
//          http://license.coscl.org.cn/MulanPSL2
// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
// EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
// MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
// See the Mulan PSL v2 for more details.

package driver

import (
	"fmt"
	"net"
	"os"
	"os/signal"
	"sync"
	"syscall"

	csi "github.com/container-storage-interface/spec/lib/go/csi"
	"google.golang.org/grpc"
	"k8s.io/klog/v2"
	registerapi "k8s.io/kubelet/pkg/apis/pluginregistration/v1"

	"gitcode.com/openFuyao/ub-ssu-csi/pkg/backend"
	"gitcode.com/openFuyao/ub-ssu-csi/pkg/kube"
	"gitcode.com/openFuyao/ub-ssu-csi/pkg/nvme"
	"gitcode.com/openFuyao/ub-ssu-csi/pkg/registry"
)

type Mode string

const (
	ControllerMode Mode = "controller"
	NodeMode       Mode = "node"
)

type Driver struct {
	name     string
	nodeID   string
	version  string
	endpoint string
	mode     Mode
	server   *grpc.Server
	stopChan chan struct{}
	stopOnce sync.Once
}

func NewDriver(name, nodeID, endpoint, version string, mode Mode) *Driver {
	return &Driver{
		name:     name,
		nodeID:   nodeID,
		version:  version,
		endpoint: endpoint,
		mode:     mode,
		stopChan: make(chan struct{}),
	}
}

func (d *Driver) Stop() {
	d.stopOnce.Do(func() {
		close(d.stopChan)
	})
}

func (d *Driver) Run(storage backend.StorageManager) error {
	d.server = grpc.NewServer()

	switch d.mode {
	case ControllerMode:
		pvcReader, err := kube.NewPVCReaderFromConfig()
		if err != nil {
			return fmt.Errorf("build pvc reader: %w", err)
		}
		csi.RegisterIdentityServer(d.server, NewIdentityServer(d.name, d.version))
		csi.RegisterControllerServer(d.server, NewControllerServer(storage, pvcReader))
		klog.InfoS("registered controller and identity service")
	case NodeMode:
		csi.RegisterNodeServer(d.server, NewNodeServer(
			d.nodeID, storage, nvme.NewConnector(),
		))
		registerapi.RegisterRegistrationServer(d.server, registry.New(d.name, d.endpoint, d.version))
		klog.InfoS("registered node and kubelet registration service")
	}

	listener, err := net.Listen("unix", d.endpoint)
	if err != nil {
		return fmt.Errorf("listen on %s: %w", d.endpoint, err)
	}

	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGTERM, syscall.SIGINT)

	go func() {
		klog.InfoS("starting CSI driver", "name", d.name, "mode", string(d.mode), "endpoint", d.endpoint)
		if err := d.server.Serve(listener); err != nil {
			klog.ErrorS(err, "grpc server serve failed")
			os.Exit(1)
		}
	}()

	select {
	case <-sigChan:
		klog.InfoS("received termination signal, shutting down...")
	case <-d.stopChan:
		klog.InfoS("stop called, shutting down...")
	}
	d.server.GracefulStop()

	return nil
}