package main

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"os/exec"
	"strconv"
	"strings"
)

// ----------------------------------------------------------------------------
// clone-disk agent state: loop-device / mount attachments held for the life
// of the control agent so several partition clones can reuse them.
// ----------------------------------------------------------------------------

var diskAttach = map[string]*diskAttachment{}

func attachedNode(diskPath string, num int, writable bool) (string, error) {
	a := diskAttach[diskPath]
	if a == nil {
		var err error
		a, err = attachDisk(diskPath, writable)
		if err != nil {
			return "", err
		}
		diskAttach[diskPath] = a
	}
	return partNode(a.blockPath, num), nil
}

func detachAllDisks() {
	for k, a := range diskAttach {
		if a.cleanup != nil {
			a.cleanup()
		}
		delete(diskAttach, k)
	}
}

// ----------------------------------------------------------------------------
// probe_disk
// ----------------------------------------------------------------------------

func handleProbeDisk(out *FrameWriter, m CtrlMsg) {
	d, err := probeDiskLayout(m.Path)
	if err != nil {
		_ = out.WriteJSON(CtrlMsg{Type: msgError, Message: err.Error(), NeedPriv: isPermErr(err)})
		return
	}
	_ = out.WriteJSON(CtrlMsg{Type: msgProbeDiskOK, LayoutJSON: d.JSON()})
}

// ----------------------------------------------------------------------------
// build_layout: create/size an image, write the partition table, re-read it.
// The boot-region bytes and the partition data are moved afterwards by the
// manager through ordinary push/pull transfers.
// ----------------------------------------------------------------------------

type buildOpts struct {
	Script    string `json:"script"`    // sfdisk restore script (Linux)
	IsFile    bool   `json:"isFile"`    // target is a regular file, not a device
	ImageSize int64  `json:"imageSize"` // truncate the file to this many bytes first
}

func handleBuildLayout(out *FrameWriter, m CtrlMsg) {
	var o buildOpts
	if err := json.Unmarshal([]byte(m.Options), &o); err != nil {
		_ = out.WriteJSON(CtrlMsg{Type: msgError, Message: "decode build opts: " + err.Error()})
		return
	}
	if o.IsFile {
		f, err := os.OpenFile(m.Path, os.O_CREATE|os.O_RDWR, 0644)
		if err != nil {
			_ = out.WriteJSON(CtrlMsg{Type: msgError, Message: err.Error(), NeedPriv: isPermErr(err)})
			return
		}
		if err := f.Truncate(o.ImageSize); err != nil {
			f.Close()
			_ = out.WriteJSON(CtrlMsg{Type: msgError, Message: fmt.Sprintf("truncate %s: %v", m.Path, err)})
			return
		}
		f.Close()
	}
	if strings.TrimSpace(o.Script) != "" {
		if err := applyPartitionTable(m.Path, o.Script); err != nil {
			_ = out.WriteJSON(CtrlMsg{Type: msgError, Message: err.Error(), NeedPriv: isPermErr(err)})
			return
		}
	} else {
		rereadPartTable(m.Path)
	}
	_ = out.WriteJSON(CtrlMsg{Type: msgBuildLayoutOK})
}

// ----------------------------------------------------------------------------
// clone_partition: fs-image and file-level methods. (Raw partitions and boot
// regions go through connect_push / connect_pull instead.)
//
// The message is sent to the SOURCE control agent first (Method set, peer
// fields as for connect_push). On a push failure the manager re-sends it to
// the DEST control agent with m.Reason == "pull".
// ----------------------------------------------------------------------------

func handleClonePartition(out *FrameWriter, m CtrlMsg) {
	o, err := decodeFSCloneOpts(m.Options)
	if err != nil {
		_ = out.WriteJSON(CtrlMsg{Type: msgError, Message: "decode clone opts: " + err.Error()})
		return
	}
	switch m.Method {
	case "fs-image":
		if m.Reason == "pull" {
			runFSPullDriver(m, o, out)
		} else {
			runFSPushDriver(m, o, out)
		}
	case "file-level":
		runFileLevelDriver(m, o, out)
	default:
		_ = out.WriteJSON(CtrlMsg{Type: msgError, Message: "unknown clone method " + m.Method})
	}
}

// localDevice resolves this side's partition node. m.Path is the whole disk
// (or image); m.PartIndex is the 1-based partition number.
func localDevice(diskPath string, num int, writable bool) (string, error) {
	if isDevicePath(diskPath) || looksLikeBlockDevice(diskPath) {
		// real disk: node is derived directly, no attach needed
		return partNode(diskPath, num), nil
	}
	return attachedNode(diskPath, num, writable)
}

func looksLikeBlockDevice(p string) bool {
	fi, err := os.Stat(p)
	return err == nil && fi.Mode()&os.ModeDevice != 0
}

// runFSPushDriver runs in the SOURCE agent: spawn the peer fs-recv (ssh or
// local), run the save tool locally, pipe save-stdout -> peer-stdin.
func runFSPushDriver(m CtrlMsg, o fsCloneOpts, out *FrameWriter) {
	dev, err := localDevice(m.Path, m.PartIndex, false)
	if err != nil {
		_ = out.WriteJSON(CtrlMsg{Type: msgPushFailed, Reason: "source device: " + err.Error()})
		return
	}
	cl := pickCloner(dev, o.FSType, o.Tool)
	if cl == nil {
		_ = out.WriteJSON(CtrlMsg{Type: msgError, Message: "no fs-image cloner for " + o.FSType})
		return
	}
	if o.ShrinkToBytes > 0 {
		_ = out.WriteJSON(CtrlMsg{Type: msgProgress, Stage: "shrink source " + o.FSType + " p" + strconv.Itoa(m.PartIndex)})
		if err := preShrinkFS(dev, o.FSType, o.ShrinkToBytes); err != nil {
			_ = out.WriteJSON(CtrlMsg{Type: msgError, Message: "pre-shrink: " + err.Error()})
			return
		}
	}

	peerTail := []string{"agent", "--role", "fs-recv",
		"--path", strconv.Itoa(m.PartIndex), // partition number; peer resolves its own node
		"--fs", o.FSType, "--fstool", o.Tool,
		"--size", strconv.FormatInt(o.SizeBytes, 10),
		"--peerdisk", m.PeerPath,
	}
	peer := peerAgentCommand(m, peerTail)
	peerIn, err := peer.StdinPipe()
	if err != nil {
		_ = out.WriteJSON(CtrlMsg{Type: msgPushFailed, Reason: err.Error()})
		return
	}
	peerOut, _ := peer.StdoutPipe()
	peerErr := newLimitedBuffer(8192)
	peer.Stderr = peerErr
	if err := peer.Start(); err != nil {
		_ = out.WriteJSON(CtrlMsg{Type: msgPushFailed, Reason: "start peer: " + err.Error()})
		return
	}

	save := exec.Command(cl.save[0], cl.save[1:]...)
	save.Stdout = peerIn
	saveErr, _ := save.StderrPipe()
	if err := save.Start(); err != nil {
		_ = peer.Process.Kill()
		_ = out.WriteJSON(CtrlMsg{Type: msgError, Message: fmt.Sprintf("start %s: %v", cl.save[0], err)})
		return
	}

	var moved int64
	go scanCloneProgress(saveErr, o.SizeBytes, func(b int64) {
		moved = b
		_ = out.WriteJSON(CtrlMsg{Type: msgProgress, Stage: o.FSType + " p" + strconv.Itoa(m.PartIndex),
			BytesCopied: b, BytesTotal: o.SizeBytes, PartIndex: m.PartIndex})
	})

	saveWait := make(chan error, 1)
	go func() { saveWait <- save.Wait() }()
	sErr := <-saveWait
	_ = peerIn.Close()

	peerLine := readLine(peerOut)
	pErr := peer.Wait()

	if sErr != nil {
		_ = out.WriteJSON(CtrlMsg{Type: msgError, Message: fmt.Sprintf("%s: %v", cl.save[0], sErr)})
		return
	}
	if pErr != nil || !strings.HasPrefix(peerLine, "OK") {
		reason := strings.TrimSpace(peerLine + " " + peerErr.String())
		if moved == 0 {
			_ = out.WriteJSON(CtrlMsg{Type: msgPushFailed, Reason: reason})
			return
		}
		_ = out.WriteJSON(CtrlMsg{Type: msgError, Message: "peer restore failed: " + reason})
		return
	}
	_ = out.WriteJSON(CtrlMsg{Type: msgProgress, BytesCopied: o.SizeBytes, BytesTotal: o.SizeBytes, PartIndex: m.PartIndex})
	_ = out.WriteJSON(CtrlMsg{Type: msgPushOK})
}

// runFSPullDriver runs in the DEST agent: spawn the peer fs-send, run the
// restore tool locally, pipe peer-stdout -> restore-stdin.
func runFSPullDriver(m CtrlMsg, o fsCloneOpts, out *FrameWriter) {
	dev, err := localDevice(m.Path, m.PartIndex, true)
	if err != nil {
		_ = out.WriteJSON(CtrlMsg{Type: msgPullFailed, Reason: "dest device: " + err.Error()})
		return
	}
	cl := pickCloner(dev, o.FSType, o.Tool)
	if cl == nil {
		_ = out.WriteJSON(CtrlMsg{Type: msgError, Message: "no fs-image cloner for " + o.FSType})
		return
	}

	peerTail := []string{"agent", "--role", "fs-send",
		"--path", strconv.Itoa(m.PartIndex),
		"--fs", o.FSType, "--fstool", o.Tool,
		"--peerdisk", m.PeerPath,
		"--shrink", strconv.FormatInt(o.ShrinkToBytes, 10),
	}
	peer := peerAgentCommand(m, peerTail)
	peerOut, err := peer.StdoutPipe()
	if err != nil {
		_ = out.WriteJSON(CtrlMsg{Type: msgPullFailed, Reason: err.Error()})
		return
	}
	peerErr := newLimitedBuffer(8192)
	peer.Stderr = peerErr
	if err := peer.Start(); err != nil {
		_ = out.WriteJSON(CtrlMsg{Type: msgPullFailed, Reason: "start peer: " + err.Error()})
		return
	}

	restore := exec.Command(cl.restore[0], cl.restore[1:]...)
	restore.Stdin = peerOut
	restoreErr, _ := restore.StderrPipe()
	if err := restore.Start(); err != nil {
		_ = peer.Process.Kill()
		_ = out.WriteJSON(CtrlMsg{Type: msgError, Message: fmt.Sprintf("start %s: %v", cl.restore[0], err)})
		return
	}
	go scanCloneProgress(restoreErr, o.SizeBytes, func(b int64) {
		_ = out.WriteJSON(CtrlMsg{Type: msgProgress, Stage: o.FSType + " p" + strconv.Itoa(m.PartIndex),
			BytesCopied: b, BytesTotal: o.SizeBytes, PartIndex: m.PartIndex})
	})

	rErr := restore.Wait()
	pErr := peer.Wait()
	if pErr != nil {
		if rErr != nil {
			_ = out.WriteJSON(CtrlMsg{Type: msgPullFailed, Reason: strings.TrimSpace(peerErr.String())})
			return
		}
	}
	if rErr != nil {
		_ = out.WriteJSON(CtrlMsg{Type: msgError, Message: fmt.Sprintf("%s: %v", cl.restore[0], rErr)})
		return
	}
	_ = out.WriteJSON(CtrlMsg{Type: msgProgress, BytesCopied: o.SizeBytes, BytesTotal: o.SizeBytes, PartIndex: m.PartIndex})
	_ = out.WriteJSON(CtrlMsg{Type: msgPullOK})
}

// ----------------------------------------------------------------------------
// fs-send / fs-recv: one-shot stdio helpers spawned on the peer.
// ----------------------------------------------------------------------------

func runFSSendRole(partNum int, fsType, tool, peerDisk string, shrinkTo int64) error {
	dev, err := resolveHelperDevice(peerDisk, partNum, shrinkTo > 0)
	if err != nil {
		return err
	}
	cl := pickCloner(dev, fsType, tool)
	if cl == nil {
		return fmt.Errorf("no fs-image cloner for %s", fsType)
	}
	if shrinkTo > 0 {
		if err := preShrinkFS(dev, fsType, shrinkTo); err != nil {
			return fmt.Errorf("pre-shrink: %w", err)
		}
	}
	cmd := exec.Command(cl.save[0], cl.save[1:]...)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	return cmd.Run()
}

func runFSRecvRole(partNum int, fsType, tool, peerDisk string) error {
	dev, err := resolveHelperDevice(peerDisk, partNum, true)
	if err != nil {
		fmt.Fprintf(os.Stderr, "ERR: %v\n", err)
		return err
	}
	cl := pickCloner(dev, fsType, tool)
	if cl == nil {
		err := fmt.Errorf("no fs-image cloner for %s", fsType)
		fmt.Fprintf(os.Stdout, "ERR: %v\n", err)
		return err
	}
	cmd := exec.Command(cl.restore[0], cl.restore[1:]...)
	cmd.Stdin = os.Stdin
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		fmt.Fprintf(os.Stdout, "ERR: %v\n", err)
		return err
	}
	fmt.Fprintln(os.Stdout, "OK")
	return nil
}

// resolveHelperDevice is used by the one-shot fs-send/fs-recv helpers, which
// have no persistent attachment map of their own.
func resolveHelperDevice(disk string, num int, writable bool) (string, error) {
	if isDevicePath(disk) || looksLikeBlockDevice(disk) {
		return partNode(disk, num), nil
	}
	a, err := attachDisk(disk, writable)
	if err != nil {
		return "", err
	}
	// deliberately leak the loop device until process exit; a one-shot
	// helper is short-lived and the parent detaches on the control side.
	_ = a
	return partNode(a.blockPath, num), nil
}

// ----------------------------------------------------------------------------
// file-level: mkfs on the target partition + rsync the tree across.
// ----------------------------------------------------------------------------

func runFileLevelDriver(m CtrlMsg, o fsCloneOpts, out *FrameWriter) {
	_ = out.WriteJSON(CtrlMsg{Type: msgError,
		Message: "file-level clone (--file-level) is not implemented yet; use --raw or an fs-image type"})
}

// ----------------------------------------------------------------------------
// reinstall_boot: best-effort, never fatal to the clone.
// ----------------------------------------------------------------------------

func handleReinstallBoot(out *FrameWriter, m CtrlMsg) {
	msg := reinstallBootloader(m.Path, m.Options)
	_ = out.WriteJSON(CtrlMsg{Type: msgReinstallBootOK, Message: msg})
}

// readLine reads one \n-terminated line (used for the fs-recv status line).
func readLine(r io.Reader) string {
	if r == nil {
		return ""
	}
	br := bufio.NewReader(r)
	s, _ := br.ReadString('\n')
	return strings.TrimSpace(s)
}