//go:build windows

package main

import (
	"encoding/binary"
	"fmt"
	"os"
	"os/exec"
	"strings"
	"syscall"
)

const (
	ioctlDiskGetDriveLayoutEx = 0x00070050 // IOCTL_DISK_GET_DRIVE_LAYOUT_EX
	ioctlDiskUpdateProperties = 0x00070140 // IOCTL_DISK_UPDATE_PROPERTIES
)

func logicalSectorSize(path string) (int64, error) {
	if !isDevicePath(path) {
		return 512, nil
	}
	return alignmentFor(path), nil
}

// guidString renders a little-endian Win32 GUID (4-2-2 + 8 bytes) as the
// canonical upper-case string.
func guidString(b []byte) string {
	if len(b) < 16 {
		return ""
	}
	d1 := binary.LittleEndian.Uint32(b[0:4])
	d2 := binary.LittleEndian.Uint16(b[4:6])
	d3 := binary.LittleEndian.Uint16(b[6:8])
	return strings.ToUpper(fmt.Sprintf("%08X-%04X-%04X-%02X%02X-%02X%02X%02X%02X%02X%02X",
		d1, d2, d3, b[8], b[9], b[10], b[11], b[12], b[13], b[14], b[15]))
}

// probeDiskLayout reads the partition table of a \\.\PhysicalDriveN via
// IOCTL_DISK_GET_DRIVE_LAYOUT_EX. Windows-native clone-disk is raw/VSS only:
// every partition is cloned block-wise and the table itself is reproduced by
// raw-copying the leading sectors (see computeBootRegionsWindows), so no
// sfdisk-style rebuild is needed.
func probeDiskLayout(path string) (*DiskLayout, error) {
	d := &DiskLayout{DiskPath: path, Tools: probeToolsetWindows()}

	info, err := statPath(path)
	if err != nil {
		return nil, err
	}
	if !info.Exists {
		return nil, fmt.Errorf("%s does not exist", path)
	}
	d.DiskSize = info.Size
	d.LogicalSector, _ = logicalSectorSize(path)
	if d.LogicalSector == 0 {
		d.LogicalSector = 512
	}
	d.PhysicalSector = d.LogicalSector

	if !isDevicePath(path) {
		// An image file: we cannot enumerate its table without extra tooling.
		d.Scheme = "raw"
		d.Warnings = append(d.Warnings, "image-file source on Windows: whole-image raw clone only")
		return d, nil
	}

	p16, err := syscall.UTF16PtrFromString(path)
	if err != nil {
		return nil, err
	}
	h, err := syscall.CreateFile(p16, syscall.GENERIC_READ,
		syscall.FILE_SHARE_READ|syscall.FILE_SHARE_WRITE, nil, syscall.OPEN_EXISTING, 0, 0)
	if err != nil {
		return nil, &os.PathError{Op: "open", Path: path, Err: err}
	}
	defer syscall.CloseHandle(h)

	buf := make([]byte, 64*1024)
	var ret uint32
	if err := syscall.DeviceIoControl(h, ioctlDiskGetDriveLayoutEx, nil, 0,
		&buf[0], uint32(len(buf)), &ret, nil); err != nil {
		return nil, fmt.Errorf("IOCTL_DISK_GET_DRIVE_LAYOUT_EX %s: %w", path, err)
	}

	style := binary.LittleEndian.Uint32(buf[0:4])
	count := binary.LittleEndian.Uint32(buf[4:8])
	switch style {
	case 0:
		d.Scheme, d.Label = "mbr", "dos"
		sig := binary.LittleEndian.Uint32(buf[8:12])
		d.LabelID = fmt.Sprintf("0x%08x", sig)
	case 1:
		d.Scheme, d.Label = "gpt", "gpt"
		d.LabelID = guidString(buf[8:24])
	default:
		d.Scheme = "raw"
	}

	const hdr = 48
	const peSize = 144
	sector := d.LogicalSector
	for i := uint32(0); i < count; i++ {
		off := hdr + int(i)*peSize
		if off+peSize > len(buf) {
			break
		}
		e := buf[off : off+peSize]
		start := int64(binary.LittleEndian.Uint64(e[8:16]))
		length := int64(binary.LittleEndian.Uint64(e[16:24]))
		num := int(binary.LittleEndian.Uint32(e[24:28]))
		if length == 0 || num == 0 {
			continue
		}
		p := Partition{
			Num:   num,
			Start: start / sector,
			Size:  length / sector,
		}
		if style == 1 { // GPT union at offset 32
			p.Type = guidString(e[32:48])
			p.UUID = guidString(e[48:64])
			p.Name = decodeUTF16(e[72:144])
		} else if style == 0 { // MBR union at offset 32
			p.Type = fmt.Sprintf("%02x", e[32])
			if e[33] != 0 {
				p.Attrs = "bootable"
			}
		}
		d.Partitions = append(d.Partitions, p)
	}

	d.BootRegions = computeBootRegionsWindows(d)
	return d, nil
}

func decodeUTF16(b []byte) string {
	u := make([]uint16, 0, len(b)/2)
	for i := 0; i+1 < len(b); i += 2 {
		c := binary.LittleEndian.Uint16(b[i : i+2])
		if c == 0 {
			break
		}
		u = append(u, c)
	}
	return strings.TrimRight(syscall.UTF16ToString(u), "\x00")
}

// computeBootRegionsWindows: reproduce the partition table by raw-copying the
// sectors before the first partition, plus the trailing 33 sectors that hold
// the backup GPT.
func computeBootRegionsWindows(d *DiskLayout) []Region {
	var regs []Region
	if len(d.Partitions) == 0 {
		return regs
	}
	first := d.Partitions[0].Start * d.LogicalSector
	for _, p := range d.Partitions {
		if s := p.Start * d.LogicalSector; s > 0 && s < first {
			first = s
		}
	}
	if first > 0 {
		regs = append(regs, Region{Offset: 0, Length: first, Note: "protective MBR + primary GPT / MBR table + gap"})
	}
	if d.Scheme == "gpt" && d.DiskSize > 33*d.LogicalSector {
		regs = append(regs, Region{
			Offset: d.DiskSize - 33*d.LogicalSector,
			Length: 33 * d.LogicalSector,
			Note:   "backup GPT",
		})
	}
	return regs
}

func rereadPartTable(target string) {
	if !isDevicePath(target) {
		return
	}
	p16, err := syscall.UTF16PtrFromString(target)
	if err != nil {
		return
	}
	h, err := syscall.CreateFile(p16, syscall.GENERIC_READ|syscall.GENERIC_WRITE,
		syscall.FILE_SHARE_READ|syscall.FILE_SHARE_WRITE, nil, syscall.OPEN_EXISTING, 0, 0)
	if err != nil {
		return
	}
	defer syscall.CloseHandle(h)
	var ret uint32
	_ = syscall.DeviceIoControl(h, ioctlDiskUpdateProperties, nil, 0, nil, 0, &ret, nil)
}

// applyPartitionTable is a no-op on Windows: the table is reproduced via the
// raw boot regions and rereadPartTable() refreshes the kernel's view.
func applyPartitionTable(target, script string) error {
	rereadPartTable(target)
	return nil
}

func probeToolsetWindows() map[string]bool {
	m := map[string]bool{}
	for _, t := range []string{"vssadmin", "wmic", "bcdboot", "diskpart", "robocopy", "powershell"} {
		if _, err := exec.LookPath(t); err == nil {
			m[t] = true
		}
	}
	return m
}

type diskAttachment struct {
	blockPath string
	cleanup   func()
}

// attachDisk is a no-op for a raw \\.\PhysicalDriveN (partitions are addressed
// by byte offset within the whole-disk handle on Windows) and unsupported for
// an image file.
func attachDisk(path string, writable bool) (*diskAttachment, error) {
	if isDevicePath(path) {
		return &diskAttachment{blockPath: path, cleanup: func() {}}, nil
	}
	return nil, fmt.Errorf("filesystem-aware / file-level clone of an image file is not supported on Windows; use --raw")
}

func partNode(disk string, num int) string { return fmt.Sprintf("%s (partition %d)", disk, num) }

func lsblkFSType(string) string { return "" }

func mountAt(string, bool) (string, func(), error) {
	return "", nil, fmt.Errorf("mount-based file-level clone is not supported on Windows; use --raw")
}