package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"time"

	"metis/pkg/plan"
	"metis/pkg/service"
	"metis/pkg/writer"
)

func remoteDevicesCmd(args []string) {
	fs := flag.NewFlagSet("remote-devices", flag.ExitOnError)
	maxBytes := fs.Int64("max-device-bytes", 300000000000, "max real removable device size")
	hostTmpDir := fs.String("host-tmp-dir", "/tmp/metis-flash-test", "host tmp dir for test writes")
	fs.Parse(args)

	devices, err := localFlashDevices(*maxBytes, *hostTmpDir)
	if err != nil {
		log.Fatalf("remote devices: %v", err)
	}
	sort.Slice(devices, func(i, j int) bool {
		left := localDeviceScore(devices[i])
		right := localDeviceScore(devices[j])
		if left != right {
			return left > right
		}
		if devices[i].SizeBytes != devices[j].SizeBytes {
			return devices[i].SizeBytes < devices[j].SizeBytes
		}
		return devices[i].Path < devices[j].Path
	})
	writeStructuredResult(map[string]any{"devices": devices})
}

func remoteBuildCmd(args []string) {
	fs := flag.NewFlagSet("remote-build", flag.ExitOnError)
	invPath := fs.String("inventory", "inventory.yaml", "inventory file")
	node := fs.String("node", "", "target node")
	cacheDir := fs.String("cache", filepath.Join(os.TempDir(), "metis-cache"), "image cache dir")
	workDir := fs.String("work-dir", filepath.Join(os.TempDir(), "metis-work"), "working directory")
	artifactRef := fs.String("artifact-ref", "", "harbor artifact ref without tag")
	buildTag := fs.String("build-tag", "", "artifact build tag")
	harborRegistry := fs.String("harbor-registry", getenvOr("METIS_HARBOR_REGISTRY", "registry.bstein.dev"), "harbor registry host")
	harborUsername := fs.String("harbor-username", getenvOr("METIS_HARBOR_USERNAME", ""), "harbor username")
	harborPassword := fs.String("harbor-password", getenvOr("METIS_HARBOR_PASSWORD", ""), "harbor password")
	fs.Parse(args)
	if *node == "" || *artifactRef == "" || *buildTag == "" {
		log.Fatalf("--node, --artifact-ref, and --build-tag are required")
	}

	if err := os.MkdirAll(*workDir, 0o755); err != nil {
		log.Fatalf("mkdir workdir: %v", err)
	}
	output := filepath.Join(*workDir, fmt.Sprintf("%s.img", *node))
	inv := loadInventory(*invPath)
	if err := plan.BuildImageFile(context.Background(), inv, *node, *cacheDir, output); err != nil {
		log.Fatalf("build image: %v", err)
	}
	if err := exec.Command("xz", "-T0", "-z", "-f", output).Run(); err != nil {
		log.Fatalf("xz compress: %v", err)
	}
	compressedPath := output + ".xz"
	info, err := os.Stat(compressedPath)
	if err != nil {
		log.Fatalf("stat compressed image: %v", err)
	}

	metadataPath := filepath.Join(*workDir, "metadata.json")
	builtAt := time.Now().UTC()
	meta := map[string]any{
		"node":         *node,
		"artifact_ref": *artifactRef,
		"build_tag":    *buildTag,
		"built_at":     builtAt.Format(time.RFC3339),
		"size_bytes":   info.Size(),
		"compressed":   true,
	}
	metaBytes, err := json.MarshalIndent(meta, "", "  ")
	if err != nil {
		log.Fatalf("encode metadata: %v", err)
	}
	if err := os.WriteFile(metadataPath, metaBytes, 0o644); err != nil {
		log.Fatalf("write metadata: %v", err)
	}
	if err := orasLogin(*harborRegistry, *harborUsername, *harborPassword); err != nil {
		log.Fatalf("oras login: %v", err)
	}
	taggedRef := fmt.Sprintf("%s:%s", *artifactRef, *buildTag)
	if err := orasPush(taggedRef, compressedPath, metadataPath); err != nil {
		log.Fatalf("oras push: %v", err)
	}
	if err := orasTag(taggedRef, "latest"); err != nil {
		log.Fatalf("oras tag latest: %v", err)
	}

	summary := service.ArtifactSummary{
		Node:       *node,
		Ref:        fmt.Sprintf("%s:latest", *artifactRef),
		BuildTag:   *buildTag,
		LocalPath:  compressedPath,
		Compressed: true,
		UpdatedAt:  builtAt,
		SizeBytes:  info.Size(),
	}
	writeStructuredResult(summary)
}

func remoteFlashCmd(args []string) {
	fs := flag.NewFlagSet("remote-flash", flag.ExitOnError)
	node := fs.String("node", "", "target node")
	device := fs.String("device", "", "target device path or test sink")
	artifactRef := fs.String("artifact-ref", "", "harbor artifact ref without tag")
	workDir := fs.String("work-dir", filepath.Join(os.TempDir(), "metis-flash"), "working directory")
	harborRegistry := fs.String("harbor-registry", getenvOr("METIS_HARBOR_REGISTRY", "registry.bstein.dev"), "harbor registry host")
	harborUsername := fs.String("harbor-username", getenvOr("METIS_HARBOR_USERNAME", ""), "harbor username")
	harborPassword := fs.String("harbor-password", getenvOr("METIS_HARBOR_PASSWORD", ""), "harbor password")
	hostTmpDir := fs.String("host-tmp-dir", "/host-tmp/metis-flash-test", "mounted host tmp dir for test writes")
	fs.Parse(args)
	if *node == "" || *device == "" || *artifactRef == "" {
		log.Fatalf("--node, --device, and --artifact-ref are required")
	}

	if err := os.MkdirAll(*workDir, 0o755); err != nil {
		log.Fatalf("mkdir workdir: %v", err)
	}
	if err := orasLogin(*harborRegistry, *harborUsername, *harborPassword); err != nil {
		log.Fatalf("oras login: %v", err)
	}
	if err := orasPull(fmt.Sprintf("%s:latest", *artifactRef), *workDir); err != nil {
		log.Fatalf("oras pull: %v", err)
	}
	imagePath, compressed, err := resolvePulledArtifact(*workDir)
	if err != nil {
		log.Fatalf("resolve artifact: %v", err)
	}
	rawImage := imagePath
	if compressed {
		rawImage = filepath.Join(*workDir, fmt.Sprintf("%s.img", *node))
		cmd := exec.Command("sh", "-lc", fmt.Sprintf("xz -dc '%s' > '%s'", imagePath, rawImage))
		if out, err := cmd.CombinedOutput(); err != nil {
			log.Fatalf("xz stream decompress: %v: %s", err, strings.TrimSpace(string(out)))
		}
	}

	destPath := *device
	if strings.HasPrefix(destPath, "hosttmp://") {
		if err := os.MkdirAll(*hostTmpDir, 0o755); err != nil {
			log.Fatalf("mkdir host tmp dir: %v", err)
		}
		destPath = filepath.Join(*hostTmpDir, fmt.Sprintf("%s.img", *node))
	}
	if err := writer.WriteImage(context.Background(), rawImage, destPath); err != nil {
		log.Fatalf("write image: %v", err)
	}
	_ = exec.Command("sync").Run()
	if strings.HasPrefix(destPath, "/dev/") {
		_ = exec.Command("blockdev", "--flushbufs", destPath).Run()
	}

	info, err := os.Stat(destPath)
	if err != nil {
		log.Fatalf("stat destination: %v", err)
	}
	writeStructuredResult(map[string]any{
		"node":       *node,
		"device":     *device,
		"dest_path":  destPath,
		"size_bytes": info.Size(),
	})
}

func writeStructuredResult(payload any) {
	data, err := json.Marshal(payload)
	if err != nil {
		log.Fatalf("encode result: %v", err)
	}
	if _, err := os.Stdout.Write(append(data, '\n')); err != nil {
		log.Fatalf("write stdout result: %v", err)
	}
	// Keep the result available in pod status so Metis does not depend on the
	// kubelet log endpoint for successful worker runs.
	_ = os.WriteFile("/dev/termination-log", data, 0o644)
}

func localFlashDevices(maxBytes int64, hostTmpDir string) ([]service.Device, error) {
	cmd := exec.Command("lsblk", "-J", "-b", "-o", "NAME,PATH,RM,HOTPLUG,SIZE,MODEL,TRAN,TYPE")
	out, err := cmd.Output()
	if err != nil {
		return nil, err
	}
	var payload struct {
		Blockdevices []struct {
			Name       string `json:"name"`
			Path       string `json:"path"`
			RM         bool   `json:"rm"`
			Hotplug    bool   `json:"hotplug"`
			Size       any    `json:"size"`
			Model      string `json:"model"`
			Tran       string `json:"tran"`
			Type       string `json:"type"`
			Mountpoint string `json:"mountpoint"`
			Children   []struct {
				Mountpoint string `json:"mountpoint"`
			} `json:"children"`
		} `json:"blockdevices"`
	}
	if err := json.Unmarshal(out, &payload); err != nil {
		return nil, err
	}
	devices := make([]service.Device, 0, len(payload.Blockdevices)+1)
	for _, dev := range payload.Blockdevices {
		if dev.Type != "disk" {
			continue
		}
		size := int64(0)
		switch value := dev.Size.(type) {
		case string:
			size, _ = strconv.ParseInt(value, 10, 64)
		case float64:
			size = int64(value)
		}
		if size <= 0 || size > maxBytes {
			continue
		}
		if dev.Tran != "usb" && !dev.RM {
			continue
		}
		if strings.TrimSpace(dev.Mountpoint) != "" {
			continue
		}
		if hasMountedChildren(dev.Children) {
			continue
		}
		devices = append(devices, service.Device{
			Name:      dev.Name,
			Path:      dev.Path,
			Model:     strings.TrimSpace(dev.Model),
			Transport: dev.Tran,
			Type:      dev.Type,
			Removable: dev.RM,
			Hotplug:   dev.Hotplug,
			SizeBytes: size,
		})
	}
	devices = append(devices, service.Device{
		Name:      "host-tmp",
		Path:      "hosttmp:///tmp",
		Model:     "Host /tmp",
		Transport: "test",
		Type:      "file",
		Note:      fmt.Sprintf("Test-only host write target under %s", humanHostPath(hostTmpDir)),
		Removable: false,
		Hotplug:   false,
		SizeBytes: 1,
	})
	return devices, nil
}

func localDeviceScore(device service.Device) int {
	score := 0
	if strings.HasPrefix(device.Path, "hosttmp://") {
		return -100
	}
	if device.Transport == "usb" {
		score += 50
	}
	if device.Removable {
		score += 30
	}
	if device.Hotplug {
		score += 20
	}
	if strings.Contains(strings.ToLower(device.Model), "sd") {
		score += 10
	}
	return score
}

func orasLogin(registry, username, password string) error {
	if strings.TrimSpace(username) == "" || strings.TrimSpace(password) == "" {
		return fmt.Errorf("harbor credentials missing")
	}
	cmd := exec.Command("oras", "login", registry, "-u", username, "-p", password)
	if out, err := cmd.CombinedOutput(); err != nil {
		return fmt.Errorf("%w: %s", err, strings.TrimSpace(string(out)))
	}
	return nil
}

func orasPush(ref, imagePath, metadataPath string) error {
	dir, args, err := orasPushInvocation(ref, imagePath, metadataPath)
	if err != nil {
		return err
	}
	cmd := exec.Command("oras", args...)
	cmd.Dir = dir
	if out, err := cmd.CombinedOutput(); err != nil {
		return fmt.Errorf("%w: %s", err, strings.TrimSpace(string(out)))
	}
	return nil
}

func orasPushInvocation(ref, imagePath, metadataPath string) (string, []string, error) {
	imageDir := filepath.Dir(imagePath)
	metadataDir := filepath.Dir(metadataPath)
	if imageDir != metadataDir {
		return "", nil, fmt.Errorf("oras push requires artifacts in one directory: %s vs %s", imageDir, metadataDir)
	}
	return imageDir, []string{
		"push",
		ref,
		fmt.Sprintf("%s:application/x-raw-disk-image", filepath.Base(imagePath)),
		fmt.Sprintf("%s:application/json", filepath.Base(metadataPath)),
	}, nil
}

func orasTag(ref string, tags ...string) error {
	args := append([]string{"tag", ref}, tags...)
	cmd := exec.Command("oras", args...)
	if out, err := cmd.CombinedOutput(); err != nil {
		return fmt.Errorf("%w: %s", err, strings.TrimSpace(string(out)))
	}
	return nil
}

func orasPull(ref, outDir string) error {
	cmd := exec.Command("oras", "pull", ref, "-o", outDir)
	if out, err := cmd.CombinedOutput(); err != nil {
		return fmt.Errorf("%w: %s", err, strings.TrimSpace(string(out)))
	}
	return nil
}

func resolvePulledArtifact(dir string) (string, bool, error) {
	var rawPath string
	var compressedPath string
	err := filepath.WalkDir(dir, func(path string, d os.DirEntry, walkErr error) error {
		if walkErr != nil {
			return walkErr
		}
		if d.IsDir() {
			return nil
		}
		switch {
		case strings.HasSuffix(path, ".img.xz"):
			compressedPath = path
		case strings.HasSuffix(path, ".img"):
			rawPath = path
		}
		return nil
	})
	if err != nil {
		return "", false, err
	}
	if compressedPath != "" {
		return compressedPath, true, nil
	}
	if rawPath != "" {
		return rawPath, false, nil
	}
	return "", false, fmt.Errorf("no .img or .img.xz artifact found in %s", dir)
}

func hasMountedChildren(children []struct {
	Mountpoint string `json:"mountpoint"`
}) bool {
	for _, child := range children {
		if strings.TrimSpace(child.Mountpoint) != "" {
			return true
		}
	}
	return false
}

func humanHostPath(path string) string {
	path = strings.TrimSpace(path)
	if strings.HasPrefix(path, "/host-tmp/") {
		trimmed := strings.TrimPrefix(path, "/host-tmp/")
		if trimmed == "" {
			return "/tmp"
		}
		return "/tmp/" + trimmed
	}
	if path == "/host-tmp" {
		return "/tmp"
	}
	return path
}

func getenvOr(key, fallback string) string {
	value := strings.TrimSpace(os.Getenv(key))
	if value == "" {
		return fallback
	}
	return value
}
