package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"sort"
	"strings"
)

// CloneDiskConfig is everything `clonetool clone-disk` needs. It embeds
// SyncConfig for the transport/deploy/ssh knobs shared with `sync`.
type CloneDiskConfig struct {
	SyncConfig
	Parts         []int
	Raw           []int
	FileLevel     []int
	FileAuto      bool
	AllowShrink   bool
	NoShrink      bool
	NewIDs        bool
	ReinstallBoot bool
	VSS           bool
	ImageSize     int64
}

func runCloneDisk(cfg CloneDiskConfig) error {
	srcSpec, err := parseSpec(cfg.Source)
	if err != nil {
		return fmt.Errorf("--source: %w", err)
	}
	dstSpec, err := parseSpec(cfg.Dest)
	if err != nil {
		return fmt.Errorf("--dest: %w", err)
	}
	if err := checkNotSame(srcSpec, dstSpec); err != nil {
		return err
	}

	srcRemoteBin, dstRemoteBin := cfg.RemoteBin, cfg.RemoteBin
	if !srcSpec.IsLocal() {
		if srcRemoteBin, err = resolveRemoteBin(&cfg.SyncConfig, srcSpec, "source"); err != nil {
			return err
		}
	}
	if !dstSpec.IsLocal() {
		if !srcSpec.IsLocal() && sameHost(srcSpec, dstSpec) {
			dstRemoteBin = srcRemoteBin
		} else if dstRemoteBin, err = resolveRemoteBin(&cfg.SyncConfig, dstSpec, "dest"); err != nil {
			return err
		}
	}

	// --- bring up both control agents -------------------------------------
	srcCtrl, srcLayout, srcSudo, err := bringUpProbe(srcSpec, "source", &cfg.SyncConfig, srcRemoteBin)
	if err != nil {
		return fmt.Errorf("source: %w", err)
	}
	defer srcCtrl.Close()

	dstCtrl, dstInfo, dstSudo, err := bringUpController(dstSpec, "dest", &cfg.SyncConfig, dstRemoteBin, dstSpec.Path)
	if err != nil {
		return fmt.Errorf("dest: %w", err)
	}
	defer dstCtrl.Close()

	isFile := !(dstInfo.Exists && dstInfo.IsDevice)
	destSize := dstInfo.Size
	fmt.Fprintf(os.Stderr, "source %s: %s, %s, %d partitions\n",
		srcSpec, srcLayout.Scheme, humanBytes(srcLayout.DiskSize), len(srcLayout.Partitions))
	reportMissingTools(srcLayout)

	// --- plan the target -------------------------------------------------
	plan, err := planTargetLayout(srcLayout, isFile, destSize, planOpts{
		Parts:       intSet(cfg.Parts),
		Raw:         intSet(cfg.Raw),
		FileLevel:   intSet(cfg.FileLevel),
		FileAuto:    cfg.FileAuto,
		AllowShrink: cfg.AllowShrink && !cfg.NoShrink,
		NewIDs:      cfg.NewIDs,
		ImageSize:   cfg.ImageSize,
	})
	if err != nil {
		return err
	}
	for _, n := range plan.Notes {
		fmt.Fprintf(os.Stderr, "  %s\n", n)
	}
	shrinks := false
	for _, pp := range plan.Parts {
		if pp.ShrinkToB > 0 {
			shrinks = true
		}
	}
	if shrinks && !cfg.Yes {
		return fmt.Errorf("shrinking resizes the SOURCE filesystem(s) in place before imaging; " +
			"re-run with --yes to confirm, or use a larger target")
	}
	if isFile && !cfg.Yes && dstInfo.Exists && dstInfo.Size > plan.ImageSize {
		if !confirmShrink(dstSpec, dstInfo.Size, plan.ImageSize) {
			return fmt.Errorf("aborted")
		}
	}

	// --- rebuild the partition table on the target ---------------------
	bo := buildOpts{Script: plan.Script, IsFile: isFile, ImageSize: plan.ImageSize}
	boJSON, _ := json.Marshal(bo)
	fmt.Fprintf(os.Stderr, "rebuilding partition table on %s ...\n", dstSpec)
	if err := dstCtrl.BuildLayout(dstSpec.Path, string(boJSON)); err != nil {
		return fmt.Errorf("build layout: %w", err)
	}

	bothLocal := srcSpec.IsLocal() && dstSpec.IsLocal()
	srcHost, srcUser := resolveConnectHost(srcSpec, &cfg.SyncConfig)
	dstHost, dstUser := resolveConnectHost(dstSpec, &cfg.SyncConfig)

	base := func(s Spec, remoteBin string) CtrlMsg {
		return CtrlMsg{
			BlockSize: cfg.BlockSize, SSHBin: cfg.SSHBin, SSHOpts: cfg.SSHOpts,
			ConnectTimeoutSec: cfg.ConnectTimeoutSec, PeerLocal: bothLocal,
		}
	}
	_ = base

	// --- copy the boot regions verbatim (offset windows) ---------------
	for _, r := range plan.BootRegions {
		if r.Length <= 0 {
			continue
		}
		fmt.Fprintf(os.Stderr, "boot region @%d (%s) %s ...\n", r.Offset, humanBytes(r.Length), r.Note)
		if err := runWindowTransfer(srcCtrl, dstCtrl, windowJob{
			srcPath: srcSpec.Path, dstPath: dstSpec.Path,
			srcBase: r.Offset, dstBase: r.Offset, size: r.Length,
			bothLocal: bothLocal, blockSize: cfg.BlockSize, cfg: &cfg.SyncConfig,
			srcHost: srcHost, srcUser: srcUser, dstHost: dstHost, dstUser: dstUser,
			srcRemoteBin: srcRemoteBin, dstRemoteBin: dstRemoteBin,
			srcSudo: srcSudo, dstSudo: dstSudo,
		}); err != nil {
			return fmt.Errorf("boot region @%d: %w", r.Offset, err)
		}
	}

	// --- per-partition data ------------------------------------------------
	for i, pp := range plan.Parts {
		tag := fmt.Sprintf("partition %d/%d (p%d %s, %s)", i+1, len(plan.Parts), pp.Num, orDash(pp.FSType), pp.Method)
		fmt.Fprintf(os.Stderr, "%s ...\n", tag)
		switch pp.Method {
		case "raw":
			sz := pp.SrcSizeB
			if pp.DstSizeB < sz {
				sz = pp.DstSizeB
			}
			if err := runWindowTransfer(srcCtrl, dstCtrl, windowJob{
				srcPath: srcSpec.Path, dstPath: dstSpec.Path,
				srcBase: pp.SrcStartB, dstBase: pp.DstStartB, size: sz,
				bothLocal: bothLocal, blockSize: cfg.BlockSize, cfg: &cfg.SyncConfig,
				srcHost: srcHost, srcUser: srcUser, dstHost: dstHost, dstUser: dstUser,
				srcRemoteBin: srcRemoteBin, dstRemoteBin: dstRemoteBin,
				srcSudo: srcSudo, dstSudo: dstSudo,
			}); err != nil {
				return fmt.Errorf("%s: %w", tag, err)
			}
		case "fs-image":
			if err := runFSImagePartition(srcCtrl, dstCtrl, pp, cloneCtx{
				srcPath: srcSpec.Path, dstPath: dstSpec.Path, bothLocal: bothLocal,
				cfg: &cfg.SyncConfig, srcHost: srcHost, srcUser: srcUser,
				dstHost: dstHost, dstUser: dstUser,
				srcRemoteBin: srcRemoteBin, dstRemoteBin: dstRemoteBin,
				srcSudo: srcSudo, dstSudo: dstSudo,
			}); err != nil {
				return fmt.Errorf("%s: %w", tag, err)
			}
		case "file-level":
			return fmt.Errorf("%s: file-level clone is not implemented yet", tag)
		}
	}

	// --- bootloader (opt-in) --------------------------------------------
	if cfg.ReinstallBoot {
		bootJSON, _ := json.Marshal(bootOpts{
			RootPart: plan.RootPart, ESPPart: plan.ESPPart, UEFI: plan.UEFI, DiskPath: dstSpec.Path,
		})
		fmt.Fprintf(os.Stderr, "reinstalling bootloader on %s ...\n", dstSpec)
		summary, err := dstCtrl.ReinstallBoot(dstSpec.Path, string(bootJSON))
		if err != nil {
			fmt.Fprintf(os.Stderr, "  bootloader: %v\n", err)
		} else {
			fmt.Fprintf(os.Stderr, "  bootloader: %s\n", summary)
		}
	}

	label := cfg.Job
	if label == "" {
		label = fmt.Sprintf("%s -> %s", srcSpec, dstSpec)
	}
	fmt.Fprintf(os.Stderr, "done: %s\n", label)
	return nil
}

// bringUpProbe starts a source control agent and runs probe_disk, retrying
// once under sudo on a permission error (mirrors bringUpController).
func bringUpProbe(spec Spec, tag string, cfg *SyncConfig, remoteBin string) (*Controller, *DiskLayout, bool, error) {
	sudo := cfg.Sudo == "always" && canElevate()
	c, err := startController(spec, tag, cfg, remoteBin, sudo)
	if err != nil {
		return nil, nil, sudo, err
	}
	lay, err := c.ProbeDisk(spec.Path)
	if err == nil {
		return c, lay, sudo, nil
	}
	if errors.Is(err, errNeedPriv) && cfg.Sudo == "auto" && !sudo && canElevate() {
		fmt.Fprintf(os.Stderr, "%s: permission denied on %s; retrying via sudo ...\n", tag, spec.Path)
		c.Close()
		sudo = true
		if c, err = startController(spec, tag, cfg, remoteBin, true); err != nil {
			return nil, nil, sudo, err
		}
		if lay, err = c.ProbeDisk(spec.Path); err == nil {
			return c, lay, sudo, nil
		}
	}
	c.Close()
	return nil, nil, sudo, err
}

func reportMissingTools(d *DiskLayout) {
	want := []string{"sfdisk", "ntfsclone", "partclone.extfs", "partclone.restore", "e2image"}
	var miss []string
	for _, t := range want {
		if !d.Tools[t] {
			miss = append(miss, t)
		}
	}
	if len(miss) > 0 {
		fmt.Fprintf(os.Stderr, "  note: not present on source: %s (affected partitions fall back to raw block copy)\n",
			strings.Join(miss, ", "))
	}
}

func orDash(s string) string {
	if s == "" {
		return "-"
	}
	return s
}

func intSet(xs []int) map[int]bool {
	m := make(map[int]bool, len(xs))
	for _, x := range xs {
		m[x] = true
	}
	return m
}

func sortedParts(ps []Partition) []Partition {
	out := append([]Partition(nil), ps...)
	sort.Slice(out, func(i, j int) bool { return out[i].Start < out[j].Start })
	return out
}