package kvm2

import (
	"bufio"
	"bytes"
	"io"
	"mime"
	"mime/multipart"
	"net/textproto"
	"os"
	"strings"
	"testing"

	"github.com/kdomanski/iso9660"
)

// readSeed builds a seed and returns its files by name plus the label.
func readSeed(t *testing.T, name, userData, vendorData string) (map[string]string, string) {
	t.Helper()
	var buf bytes.Buffer
	if err := buildSeedISO(name, userData, vendorData, &buf); err != nil {
		t.Fatalf("buildSeedISO: %v", err)
	}
	img, err := iso9660.OpenImage(bytes.NewReader(buf.Bytes()))
	if err != nil {
		t.Fatalf("OpenImage: %v", err)
	}
	root, err := img.RootDir()
	if err != nil {
		t.Fatalf("RootDir: %v", err)
	}
	children, err := root.GetChildren()
	if err != nil {
		t.Fatalf("GetChildren: %v", err)
	}
	files := make(map[string]string, len(children))
	for _, c := range children {
		if c.IsDir() {
			t.Errorf("unexpected directory %q in seed", c.Name())
			continue
		}
		data, err := io.ReadAll(c.Reader())
		if err != nil {
			t.Fatalf("reading %s: %v", c.Name(), err)
		}
		files[c.Name()] = string(data)
	}
	label, err := img.Label()
	if err != nil {
		t.Fatalf("Label: %v", err)
	}
	return files, label
}

func TestSeedISO_filesAndLabel(t *testing.T) {
	t.Parallel()
	userData := "#cloud-config\npackages: [squid]\n"
	vendorData := "#!/bin/sh\necho hi\n"
	files, label := readSeed(t, "proxy-abc", userData, vendorData)
	if label != seedLabel {
		t.Errorf("label = %q, want %q", label, seedLabel)
	}
	want := map[string]string{
		"user-data":   userData,
		"meta-data":   "instance-id: proxy-abc\nlocal-hostname: proxy-abc\n",
		"vendor-data": vendorData,
	}
	if len(files) != len(want) {
		t.Errorf("seed has %d files %v, want exactly %d", len(files), keys(files), len(want))
	}
	for name, content := range want {
		if got, ok := files[name]; !ok {
			t.Errorf("seed lacks %s", name)
		} else if got != content {
			t.Errorf("%s = %q, want %q", name, got, content)
		}
	}
}

func keys(m map[string]string) []string {
	out := make([]string, 0, len(m))
	for k := range m {
		out = append(out, k)
	}
	return out
}

func TestSeedISO_userDataByteForByte(t *testing.T) {
	t.Parallel()
	tests := []struct {
		name     string
		userData string
	}{
		{name: "mime multipart with CRLF", userData: "MIME-Version: 1.0\r\nContent-Type: multipart/mixed; boundary=\"b\"\r\n\r\n--b\r\nContent-Type: text/cloud-config\r\n\r\n#cloud-config\r\npackages: [squid]\r\n--b--\r\n"},
		{name: "jinja header", userData: "## template: jinja\n#cloud-config\nhostname: {{ v1.local_hostname }}\n"},
		{name: "no trailing newline", userData: "#cloud-config\nruncmd:\n  - echo done"},
		{name: "bytes above 0x7f", userData: "#cloud-config\nwrite_files:\n  - content: \"caf\xc3\xa9 \xff\xfe\x80\"\n"},
		{name: "empty", userData: ""},
		{name: "trailing whitespace preserved", userData: "#!/bin/sh\necho x \n\n\n"},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()
			files, _ := readSeed(t, "proxy-abc", tc.userData, "")
			if got := files["user-data"]; got != tc.userData {
				t.Errorf("user-data changed:\n got %q\nwant %q", got, tc.userData)
			}
		})
	}
}

// Not parallel: it changes TMPDIR for the process.
func TestSeedISO_tempDirRemoved(t *testing.T) {
	tmp := t.TempDir()
	t.Setenv("TMPDIR", tmp)
	if got := os.TempDir(); got != tmp {
		t.Fatalf("os.TempDir() = %q, want %q", got, tmp)
	}
	var buf bytes.Buffer
	if err := buildSeedISO("proxy-abc", "#cloud-config\n", guestAgentScript, &buf); err != nil {
		t.Fatalf("buildSeedISO: %v", err)
	}
	if buf.Len() == 0 {
		t.Fatal("seed ISO is empty")
	}
	entries, err := os.ReadDir(tmp)
	if err != nil {
		t.Fatal(err)
	}
	if len(entries) != 0 {
		names := make([]string, 0, len(entries))
		for _, e := range entries {
			names = append(names, e.Name())
		}
		t.Errorf("staging left behind in TMPDIR: %v", names)
	}
}

func TestSeedISO_emptyVendorDataStillHasFile(t *testing.T) {
	t.Parallel()
	files, _ := readSeed(t, "proxy-abc", "#cloud-config\n", "")
	got, ok := files["vendor-data"]
	if !ok {
		t.Fatalf("seed lacks vendor-data; files: %v", keys(files))
	}
	if got != "" {
		t.Errorf("vendor-data = %q, want empty", got)
	}
}

func TestSeedISO_rejectsEmptyName(t *testing.T) {
	t.Parallel()
	var buf bytes.Buffer
	if err := buildSeedISO("", "#cloud-config\n", "", &buf); err == nil {
		t.Fatal("buildSeedISO with empty name succeeded")
	}
	if buf.Len() != 0 {
		t.Errorf("wrote %d bytes alongside the error", buf.Len())
	}
}

func TestVendorData_isShellScriptNotCloudConfig(t *testing.T) {
	t.Parallel()
	got, err := buildVendorData("", true)
	if err != nil {
		t.Fatal(err)
	}
	if !strings.HasPrefix(got, "#!/bin/sh\n") {
		t.Errorf("vendor-data starts with %q, want #!/bin/sh", firstLine(got))
	}
	if strings.HasPrefix(got, "#cloud-config") {
		t.Error("vendor-data is a cloud-config; cloud-init would merge it away")
	}
	for _, want := range []string{"qemu-guest-agent", "set -e", "apt-get", "dnf", "yum", "zypper", "apk", "systemctl enable --now qemu-guest-agent", "rc-update", "installGuestAgent: false", "DEBIAN_FRONTEND=noninteractive"} {
		if !strings.Contains(got, want) {
			t.Errorf("script lacks %q", want)
		}
	}
	if strings.Contains(got, "MIME-Version") {
		t.Error("script alone should not be wrapped in a multipart")
	}
}

func firstLine(s string) string {
	line, _, _ := strings.Cut(s, "\n")
	return line
}

func bufioReader(s string) *bufio.Reader {
	return bufio.NewReader(strings.NewReader(s))
}

func TestVendorData_multipartWhenExtraProvided(t *testing.T) {
	t.Parallel()
	configured := "#cloud-config\nntp:\n  enabled: true\n"
	got, err := buildVendorData(configured, true)
	if err != nil {
		t.Fatal(err)
	}
	if !strings.HasPrefix(got, "MIME-Version: 1.0\r\n") {
		t.Errorf("vendor-data starts with %q, want MIME-Version: 1.0", firstLine(got))
	}
	r := textproto.NewReader(bufioReader(got))
	hdr, err := r.ReadMIMEHeader()
	if err != nil {
		t.Fatalf("reading top-level headers: %v", err)
	}
	mediaType, params, err := mime.ParseMediaType(hdr.Get("Content-Type"))
	if err != nil {
		t.Fatalf("parsing Content-Type %q: %v", hdr.Get("Content-Type"), err)
	}
	if mediaType != "multipart/mixed" {
		t.Errorf("media type = %q, want multipart/mixed", mediaType)
	}
	mr := multipart.NewReader(r.R, params["boundary"])
	type part struct{ contentType, body string }
	var parts []part
	for {
		p, err := mr.NextPart()
		if err == io.EOF {
			break
		}
		if err != nil {
			t.Fatalf("NextPart: %v", err)
		}
		body, err := io.ReadAll(p)
		if err != nil {
			t.Fatalf("reading part: %v", err)
		}
		parts = append(parts, part{p.Header.Get("Content-Type"), string(body)})
	}
	want := []part{{"text/plain", configured}, {"text/x-shellscript", guestAgentScript}}
	if len(parts) != len(want) {
		t.Fatalf("got %d parts, want %d", len(parts), len(want))
	}
	for i := range want {
		if parts[i] != want[i] {
			t.Errorf("part %d = {%q, %q}, want {%q, %q}", i, parts[i].contentType, firstLine(parts[i].body), want[i].contentType, firstLine(want[i].body))
		}
	}
}

func TestVendorData_offWithoutConfiguredIsEmpty(t *testing.T) {
	t.Parallel()
	got, err := buildVendorData("", false)
	if err != nil {
		t.Fatal(err)
	}
	if got != "" {
		t.Errorf("vendor-data = %q, want empty", got)
	}
}

func TestVendorData_offWithConfiguredPassesThrough(t *testing.T) {
	t.Parallel()
	configured := "#cloud-config\nntp:\n  enabled: true"
	got, err := buildVendorData(configured, false)
	if err != nil {
		t.Fatal(err)
	}
	if got != configured {
		t.Errorf("vendor-data = %q, want the configured document unchanged", got)
	}
}