package instancectl

import (
	"context"
	"encoding/json"
	"errors"
	"io"
	"net"
	"os"
	"path/filepath"
	"strconv"
	"time"

	"github.com/carlosframework/platform/internal/provision"
	"github.com/carlosframework/platform/internal/safefile"
)

// ClaimExecPermit is root's one-launch reservation. SafeUntil is CLOCK_BOOTTIME
// nanoseconds on BootID; the trampoline compares it immediately before ForkExec.
type ClaimExecPermit struct {
	Reservation, Inputs, BootID string
	SafeUntil                   int64
	Containment                 uint8 `json:",omitempty"`
}

func (p ClaimExecPermit) Valid() bool {
	return p.Reservation != "" && p.Inputs != "" && p.BootID != "" && p.SafeUntil > 0
}

const ClaimExecSocket = "/run/carlos/claim-exec.sock"
const ClaimExecSnapshot = "/run/carlos/claim-exec.json"
const ClaimExecBootID = "/run/carlos/claim-boot-id"

func (u ClaimUnits) StartReserved(ctx context.Context, in ClaimUnit, l Launch, p ClaimExecPermit, check func() error) (ClaimUnitStarted, error) {
	if !p.Valid() || p.Containment != 1 || check == nil {
		return ClaimUnitStarted{}, ErrClaimUnit
	}
	return u.start(ctx, in, &l, check, &p)
}

type claimExecGate struct {
	listener *net.UnixListener
	path     string
	permit   ClaimExecPermit
}

// Keep the closed socket inode for ExecStopPost namespace setup. clearLaunch
// removes it only after the unit has stopped and its cgroup is empty.
func (g *claimExecGate) close() { g.listener.Close() }
func (u ClaimUnits) prepareExecGate(ctx context.Context, host string, p ClaimExecPermit) (*claimExecGate, error) {
	source, dropin, err := u.launchPaths(host)
	if err != nil {
		return nil, err
	}
	raw, err := json.Marshal(p)
	if err != nil {
		return nil, err
	}
	if err = safefile.WriteFile(filepath.Join(source, "permit.json"), raw, 0444); err != nil {
		return nil, err
	}
	socket := filepath.Join(source, "exec.sock")
	// An old socket is removed only after start proved the old unit stopped.
	if err = safefile.Remove(socket); err != nil && !errors.Is(err, os.ErrNotExist) {
		return nil, err
	}
	ln, err := claimListenExecSocket(source)
	if err != nil {
		return nil, err
	}
	ln.SetUnlinkOnClose(false)
	gate := &claimExecGate{ln, socket, p}
	fail := func(e error) (*claimExecGate, error) { gate.close(); return nil, e }
	// Host ancestry remains root-private; inside the namespace peer identity is
	// checked against the pinned unit MainPID, never against a supplied PID.
	body, err := safefile.ReadFile(filepath.Join(dropin, "60-claim-launch.conf"))
	if err != nil {
		return fail(err)
	}
	instance := "%I"
	if u.BoundUnit != "" {
		instance = provision.InstanceID(host)
	}
	body = append(body, []byte("BindReadOnlyPaths="+socket+":"+ClaimExecSocket+"\nBindReadOnlyPaths="+filepath.Join(source, "permit.json")+":"+ClaimExecSnapshot+"\nBindReadOnlyPaths=/proc/sys/kernel/random/boot_id:"+ClaimExecBootID+"\nExecStart=\nExecStart="+SystemBinary+" instance-exec -claim-reservation="+ClaimExecSnapshot+" run "+instance+"\n")...)
	if err = safefile.WriteFile(filepath.Join(dropin, "60-claim-launch.conf"), body, 0644); err != nil {
		return fail(err)
	}
	if _, err = u.Systemctl(ctx, "daemon-reload"); err != nil {
		return fail(err)
	}
	return gate, nil
}

func (g *claimExecGate) authorize(ctx context.Context, u ClaimUnits, s *ClaimUnitStarted, check func() error) error {
	deadline := time.Now().Add(u.budget())
	if d, ok := ctx.Deadline(); ok && d.Before(deadline) {
		deadline = d
	}
	if err := g.listener.SetDeadline(deadline); err != nil {
		return err
	}
	stop := context.AfterFunc(ctx, func() { g.listener.Close() })
	defer stop()
	conn, err := g.listener.AcceptUnix()
	if err != nil {
		return err
	}
	defer conn.Close()
	closeOnCancel := context.AfterFunc(ctx, func() { conn.Close() })
	defer closeOnCancel()
	if err = conn.SetDeadline(deadline); err != nil {
		return err
	}
	if err = u.pinExecPeer(ctx, conn, s); err != nil {
		return err
	}
	if err = u.Same(ctx, *s); err != nil {
		return err
	}
	var asked ClaimExecPermit
	dec := json.NewDecoder(io.LimitReader(conn, 4096))
	if err = dec.Decode(&asked); err != nil {
		return err
	}
	if asked != g.permit {
		return ErrClaimUnit
	}
	if err = check(); err != nil {
		return err
	}
	if err = u.Same(ctx, *s); err != nil {
		return err
	}
	if u.BoundAddr != "" {
		if err = u.restrictBoundNetwork(*s); err != nil {
			return err
		}
		if err = u.Same(ctx, *s); err != nil {
			return err
		}
	}
	// One connection, one permit. Closing/restarting root cannot replay it.
	return json.NewEncoder(conn).Encode(g.permit)
}

// PrivatePIDs may replace systemd's setup PID with the namespace init PID.
// Only this pre-exec handshake may pin that transition. The invocation and
// cgroup must stay fixed, and systemd must identify the socket peer as MainPID.
func (u ClaimUnits) pinExecPeer(ctx context.Context, conn net.Conn, s *ClaimUnitStarted) error {
	for {
		if err := ctx.Err(); err != nil {
			return err
		}
		if err := claimExecPeer(conn, s.MainPID); err == nil {
			return nil
		}
		state, err := u.show(ctx, s.Unit, "InvocationID", "MainPID", "ControlGroup", "ActiveState")
		if err != nil {
			return err
		}
		if state["InvocationID"] != s.InvocationID || state["ControlGroup"] != "/"+s.ControlGroup || state["ActiveState"] != "active" {
			return ErrClaimUnit
		}
		pid, err := strconv.Atoi(state["MainPID"])
		if err != nil || pid <= 0 {
			return ErrClaimUnit
		}
		if err := claimExecPeer(conn, pid); err == nil {
			s.MainPID = pid
			return nil
		}
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-time.After(u.poll()):
		}
	}
}