package image

import (
	"archive/zip"
	"crypto/md5"
	"crypto/sha256"
	"encoding/hex"
	"os"
	"os/exec"
	"path/filepath"
	"testing"
)

func TestDownloadDecompressesXZFileURLs(t *testing.T) {
	if _, err := exec.LookPath("xz"); err != nil {
		t.Skip("xz not available")
	}
	dir := t.TempDir()
	raw := filepath.Join(dir, "base.img")
	if err := os.WriteFile(raw, []byte("metis-xz-test"), 0o644); err != nil {
		t.Fatal(err)
	}
	compressed := raw + ".xz"
	cmd := exec.Command("xz", "-zk", raw)
	if out, err := cmd.CombinedOutput(); err != nil {
		t.Fatalf("xz: %v: %s", err, string(out))
	}
	dest := filepath.Join(dir, "copy.img")
	if err := Download("file://"+compressed, dest); err != nil {
		t.Fatalf("Download: %v", err)
	}
	data, err := os.ReadFile(dest)
	if err != nil {
		t.Fatalf("ReadFile: %v", err)
	}
	if string(data) != "metis-xz-test" {
		t.Fatalf("unexpected decompressed content: %q", string(data))
	}
}

func TestDownloadAndVerifyUsesArchiveChecksumForXZ(t *testing.T) {
	if _, err := exec.LookPath("xz"); err != nil {
		t.Skip("xz not available")
	}
	dir := t.TempDir()
	raw := filepath.Join(dir, "base.img")
	if err := os.WriteFile(raw, []byte("metis-xz-test"), 0o644); err != nil {
		t.Fatal(err)
	}
	compressed := raw + ".xz"
	cmd := exec.Command("xz", "-zk", raw)
	if out, err := cmd.CombinedOutput(); err != nil {
		t.Fatalf("xz: %v: %s", err, string(out))
	}
	archiveBytes, err := os.ReadFile(compressed)
	if err != nil {
		t.Fatalf("ReadFile archive: %v", err)
	}
	archiveSum := sha256.Sum256(archiveBytes)
	dest := filepath.Join(dir, "copy.img")
	localPath, err := DownloadAndVerify("file://"+compressed, dest, "sha256:"+hex.EncodeToString(archiveSum[:]))
	if err != nil {
		t.Fatalf("DownloadAndVerify: %v", err)
	}
	if localPath != dest {
		t.Fatalf("expected local path %s, got %s", dest, localPath)
	}
	data, err := os.ReadFile(dest)
	if err != nil {
		t.Fatalf("ReadFile dest: %v", err)
	}
	if string(data) != "metis-xz-test" {
		t.Fatalf("unexpected decompressed content: %q", string(data))
	}
}

func TestDownloadAndVerifyReplacesStaleBadArchiveCache(t *testing.T) {
	if _, err := exec.LookPath("xz"); err != nil {
		t.Skip("xz not available")
	}
	dir := t.TempDir()
	raw := filepath.Join(dir, "base.img")
	if err := os.WriteFile(raw, []byte("metis-xz-test"), 0o644); err != nil {
		t.Fatal(err)
	}
	compressed := raw + ".xz"
	cmd := exec.Command("xz", "-zk", raw)
	if out, err := cmd.CombinedOutput(); err != nil {
		t.Fatalf("xz: %v: %s", err, string(out))
	}
	archiveBytes, err := os.ReadFile(compressed)
	if err != nil {
		t.Fatalf("ReadFile archive: %v", err)
	}
	archiveSum := sha256.Sum256(archiveBytes)
	dest := filepath.Join(dir, "copy.img")
	staleArchive := dest + ".xz"
	if err := os.WriteFile(staleArchive, []byte("bad-cache"), 0o644); err != nil {
		t.Fatalf("WriteFile stale archive: %v", err)
	}
	if _, err := DownloadAndVerify("file://"+compressed, dest, "sha256:"+hex.EncodeToString(archiveSum[:])); err != nil {
		t.Fatalf("DownloadAndVerify with stale archive: %v", err)
	}
	data, err := os.ReadFile(dest)
	if err != nil {
		t.Fatalf("ReadFile dest: %v", err)
	}
	if string(data) != "metis-xz-test" {
		t.Fatalf("unexpected decompressed content: %q", string(data))
	}
}

func TestDownloadDecompressesZIPFileURLs(t *testing.T) {
	dir := t.TempDir()
	archive := filepath.Join(dir, "base.zip")
	if err := writeTestZIP(archive, map[string]string{
		"nested/base.img": "metis-zip-test",
	}); err != nil {
		t.Fatalf("writeTestZIP: %v", err)
	}
	dest := filepath.Join(dir, "copy.img")
	if err := Download("file://"+archive, dest); err != nil {
		t.Fatalf("Download: %v", err)
	}
	data, err := os.ReadFile(dest)
	if err != nil {
		t.Fatalf("ReadFile: %v", err)
	}
	if string(data) != "metis-zip-test" {
		t.Fatalf("unexpected decompressed content: %q", string(data))
	}
}

func TestDownloadAndVerifyUsesArchiveChecksumForZIP(t *testing.T) {
	dir := t.TempDir()
	archive := filepath.Join(dir, "base.zip")
	if err := writeTestZIP(archive, map[string]string{
		"base.img": "metis-zip-test",
	}); err != nil {
		t.Fatalf("writeTestZIP: %v", err)
	}
	archiveBytes, err := os.ReadFile(archive)
	if err != nil {
		t.Fatalf("ReadFile archive: %v", err)
	}
	archiveSum := sha256.Sum256(archiveBytes)
	dest := filepath.Join(dir, "copy.img")
	localPath, err := DownloadAndVerify("file://"+archive, dest, "sha256:"+hex.EncodeToString(archiveSum[:]))
	if err != nil {
		t.Fatalf("DownloadAndVerify: %v", err)
	}
	if localPath != dest {
		t.Fatalf("expected local path %s, got %s", dest, localPath)
	}
	data, err := os.ReadFile(dest)
	if err != nil {
		t.Fatalf("ReadFile dest: %v", err)
	}
	if string(data) != "metis-zip-test" {
		t.Fatalf("unexpected decompressed content: %q", string(data))
	}
}

func TestVerifyChecksumAcceptsMD5(t *testing.T) {
	dir := t.TempDir()
	path := filepath.Join(dir, "sample.img")
	if err := os.WriteFile(path, []byte("metis-md5-test"), 0o644); err != nil {
		t.Fatalf("WriteFile: %v", err)
	}
	sum := md5.Sum([]byte("metis-md5-test"))
	if err := VerifyChecksum(path, "md5:"+hex.EncodeToString(sum[:])); err != nil {
		t.Fatalf("VerifyChecksum md5: %v", err)
	}
}

func writeTestZIP(path string, files map[string]string) error {
	out, err := os.Create(path)
	if err != nil {
		return err
	}
	defer out.Close()
	zw := zip.NewWriter(out)
	for name, content := range files {
		w, err := zw.Create(name)
		if err != nil {
			_ = zw.Close()
			return err
		}
		if _, err := w.Write([]byte(content)); err != nil {
			_ = zw.Close()
			return err
		}
	}
	if err := zw.Close(); err != nil {
		return err
	}
	return out.Sync()
}