package criu

import (
	"errors"
	"fmt"
	"os"
	"os/exec"
	"sort"
	"strconv"
	"syscall"

	proto "github.com/checkpoint-restore/go-criu/v8/internal/proto"
	"github.com/checkpoint-restore/go-criu/v8/rpc"
)

// extraFilesStartFd is the first fd number assigned to cmd.ExtraFiles by os/exec.
// As documented in os/exec: "entry i becomes file descriptor 3+i"
const extraFilesStartFd = 3

// Criu struct
type Criu struct {
	swrkCmd    *exec.Cmd
	swrkSk     *os.File
	swrkPath   string
	inheritFds map[string]*os.File
}

// MakeCriu returns the Criu object required for most operations
func MakeCriu() *Criu {
	return &Criu{
		swrkPath: "criu",
	}
}

// SetCriuPath allows setting the path to the CRIU binary
// if it is in a non standard location
func (c *Criu) SetCriuPath(path string) {
	c.swrkPath = path
}

// AddInheritFd registers a file descriptor to be passed to CRIU.
// If opts.InheritFd is not set for an operation, it will be populated
// from these registrations using the same key order.
func (c *Criu) AddInheritFd(key string, file *os.File) {
	if c.inheritFds == nil {
		c.inheritFds = make(map[string]*os.File)
	}
	c.inheritFds[key] = file
}

func (c *Criu) inheritFdKeys() []string {
	if len(c.inheritFds) == 0 {
		return nil
	}
	keys := make([]string, 0, len(c.inheritFds))
	for k := range c.inheritFds {
		keys = append(keys, k)
	}
	sort.Strings(keys)
	return keys
}

func (c *Criu) ensureInheritFd(opts *rpc.CriuOpts) {
	if opts == nil || len(opts.GetInheritFd()) > 0 || len(c.inheritFds) == 0 {
		return
	}
	keys := c.inheritFdKeys()
	if len(keys) == 0 {
		return
	}
	opts.InheritFd = make([]*rpc.InheritFd, 0, len(keys))
	for i, key := range keys {
		fd := int32(extraFilesStartFd + i)
		opts.InheritFd = append(opts.InheritFd, &rpc.InheritFd{
			Key: proto.Ptr(key),
			Fd:  proto.Ptr(fd),
		})
	}
}

// Prepare sets up everything for the RPC communication to CRIU
func (c *Criu) Prepare() error {
	return c.doPrepare(nil)
}

func (c *Criu) doPrepare(opts *rpc.CriuOpts) error {
	fds, err := syscall.Socketpair(syscall.AF_LOCAL, syscall.SOCK_SEQPACKET, 0)
	if err != nil {
		return err
	}

	cln := os.NewFile(uintptr(fds[0]), "criu-xprt-cln")
	syscall.CloseOnExec(fds[0])
	srv := os.NewFile(uintptr(fds[1]), "criu-xprt-srv")
	defer func() { _ = srv.Close() }()

	args := []string{"swrk", strconv.Itoa(fds[1])}
	// #nosec G204
	cmd := exec.Command(c.swrkPath, args...)

	// Collect file descriptors to pass to child
	inheritKeys := c.inheritFdKeys()
	extraFiles := make([]*os.File, 0, len(inheritKeys))

	// Add fds from AddInheritFd (sorted for stable ordering)
	for _, k := range inheritKeys {
		extraFiles = append(extraFiles, c.inheritFds[k])
	}

	c.ensureInheritFd(opts)

	cmd.ExtraFiles = extraFiles

	err = cmd.Start()
	if err != nil {
		_ = cln.Close()
		return err
	}

	c.swrkCmd = cmd
	c.swrkSk = cln

	return nil
}

// Cleanup cleans up
func (c *Criu) Cleanup() error {
	var errs []error
	if c.swrkCmd != nil {
		if err := c.swrkSk.Close(); err != nil {
			errs = append(errs, err)
		}
		c.swrkSk = nil
		if err := c.swrkCmd.Wait(); err != nil {
			errs = append(errs, fmt.Errorf("criu swrk failed: %w", err))
		}
		c.swrkCmd = nil
	}
	return errors.Join(errs...)
}

func (c *Criu) sendAndRecv(reqB []byte) ([]byte, int, error) {
	cln := c.swrkSk
	_, err := cln.Write(reqB)
	if err != nil {
		return nil, 0, err
	}

	respB := make([]byte, 2*4096)
	n, err := cln.Read(respB)
	if err != nil {
		return nil, 0, err
	}

	return respB, n, nil
}

func (c *Criu) doSwrk(reqType rpc.CriuReqType, opts *rpc.CriuOpts, nfy Notify) error {
	resp, err := c.doSwrkWithResp(reqType, opts, nfy, nil)
	if err != nil {
		return err
	}
	respType := resp.GetType()
	if respType != reqType {
		return errors.New("unexpected CRIU RPC response")
	}

	return nil
}

func (c *Criu) doSwrkWithResp(reqType rpc.CriuReqType, opts *rpc.CriuOpts, nfy Notify, features *rpc.CriuFeatures) (resp *rpc.CriuResp, retErr error) {
	c.ensureInheritFd(opts)

	req := rpc.CriuReq{
		Type: &reqType,
		Opts: opts,
	}

	if nfy != nil {
		opts.NotifyScripts = proto.Ptr(true)
	}

	if features != nil {
		req.Features = features
	}

	if c.swrkCmd == nil {
		err := c.doPrepare(opts)
		if err != nil {
			return nil, err
		}

		defer func() {
			// append any cleanup errors to the returned error
			err := c.Cleanup()
			if err != nil {
				retErr = errors.Join(retErr, err)
			}
		}()
	}

	for {
		reqB, err := req.MarshalVT()
		if err != nil {
			return nil, err
		}

		respB, respS, err := c.sendAndRecv(reqB)
		if err != nil {
			return nil, err
		}

		resp = &rpc.CriuResp{}
		err = resp.UnmarshalVT(respB[:respS])
		if err != nil {
			return nil, err
		}

		if !resp.GetSuccess() {
			return resp, fmt.Errorf("operation failed (msg:%s err:%d)",
				resp.GetCrErrmsg(), resp.GetCrErrno())
		}

		respType := resp.GetType()
		if respType != rpc.CriuReqType_NOTIFY {
			break
		}
		if nfy == nil {
			return resp, errors.New("unexpected notify")
		}

		notify := resp.GetNotify()
		switch notify.GetScript() {
		case "pre-dump":
			err = nfy.PreDump()
		case "post-dump":
			err = nfy.PostDump()
		case "pre-restore":
			err = nfy.PreRestore()
		case "post-restore":
			err = nfy.PostRestore(notify.GetPid())
		case "network-lock":
			err = nfy.NetworkLock()
		case "network-unlock":
			err = nfy.NetworkUnlock()
		case "setup-namespaces":
			err = nfy.SetupNamespaces(notify.GetPid())
		case "post-setup-namespaces":
			err = nfy.PostSetupNamespaces()
		case "post-resume":
			err = nfy.PostResume()
		default:
			err = nil
		}

		if err != nil {
			return resp, err
		}

		req = rpc.CriuReq{
			Type:          &respType,
			NotifySuccess: proto.Ptr(true),
		}
	}

	return resp, nil
}

// Dump dumps a process
func (c *Criu) Dump(opts *rpc.CriuOpts, nfy Notify) error {
	return c.doSwrk(rpc.CriuReqType_DUMP, opts, nfy)
}

// Restore restores a process
func (c *Criu) Restore(opts *rpc.CriuOpts, nfy Notify) error {
	return c.doSwrk(rpc.CriuReqType_RESTORE, opts, nfy)
}

// PreDump does a pre-dump
func (c *Criu) PreDump(opts *rpc.CriuOpts, nfy Notify) error {
	return c.doSwrk(rpc.CriuReqType_PRE_DUMP, opts, nfy)
}

// StartPageServer starts the page server
func (c *Criu) StartPageServer(opts *rpc.CriuOpts) error {
	return c.doSwrk(rpc.CriuReqType_PAGE_SERVER, opts, nil)
}

// StartPageServerChld starts the page server and returns PID and port
func (c *Criu) StartPageServerChld(opts *rpc.CriuOpts) (int, int, error) {
	resp, err := c.doSwrkWithResp(rpc.CriuReqType_PAGE_SERVER_CHLD, opts, nil, nil)
	if err != nil {
		return 0, 0, err
	}

	return int(resp.GetPs().GetPid()), int(resp.GetPs().GetPort()), nil
}

// GetCriuVersion executes the VERSION RPC call and returns the version
// as an integer. Major * 10000 + Minor * 100 + SubLevel
func (c *Criu) GetCriuVersion() (int, error) {
	resp, err := c.doSwrkWithResp(rpc.CriuReqType_VERSION, nil, nil, nil)
	if err != nil {
		return 0, err
	}

	if resp.GetType() != rpc.CriuReqType_VERSION {
		return 0, errors.New("unexpected CRIU RPC response")
	}

	version := resp.GetVersion().GetMajorNumber() * 10000
	version += resp.GetVersion().GetMinorNumber() * 100
	if resp.GetVersion().GetSublevel() != 0 {
		version += resp.GetVersion().GetSublevel()
	}

	if resp.GetVersion().GetGitid() != "" {
		// taken from runc: if it is a git release -> increase minor by 1
		version -= (version % 100)
		version += 100
	}

	return int(version), nil
}

// IsCriuAtLeast checks if the version is at least the same
// as the parameter version
func (c *Criu) IsCriuAtLeast(version int) (bool, error) {
	criuVersion, err := c.GetCriuVersion()
	if err != nil {
		return false, err
	}

	if criuVersion >= version {
		return true, nil
	}

	return false, nil
}