package image

import (
	"archive/zip"
	"crypto/md5"
	"crypto/sha256"
	"encoding/hex"
	"net/http"
	"net/http/httptest"
	"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 TestDownloadAndVerifyUsesHTTPAndCachedFile(t *testing.T) {
	body := []byte("metis-http-test")
	sum := sha256.Sum256(body)
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/image.img" {
			http.NotFound(w, r)
			return
		}
		_, _ = w.Write(body)
	}))
	defer srv.Close()

	dir := t.TempDir()
	dest := filepath.Join(dir, "image.img")
	path, err := DownloadAndVerify(srv.URL+"/image.img", dest, "sha256:"+hex.EncodeToString(sum[:]))
	if err != nil {
		t.Fatalf("DownloadAndVerify: %v", err)
	}
	if path != dest {
		t.Fatalf("path = %q, want %q", path, dest)
	}
	if got, _ := os.ReadFile(dest); string(got) != string(body) {
		t.Fatalf("downloaded body = %q", string(got))
	}

	if err := os.WriteFile(dest, body, 0o644); err != nil {
		t.Fatal(err)
	}
	if _, err := DownloadAndVerify(srv.URL+"/image.img", dest, "sha256:"+hex.EncodeToString(sum[:])); err != nil {
		t.Fatalf("cached DownloadAndVerify: %v", err)
	}
}

func TestDownloadAndVerifyRejectsBadChecksum(t *testing.T) {
	dir := t.TempDir()
	src := filepath.Join(dir, "src.img")
	if err := os.WriteFile(src, []byte("bad"), 0o644); err != nil {
		t.Fatal(err)
	}
	if _, err := DownloadAndVerify("file://"+src, filepath.Join(dir, "dest.img"), "sha256:deadbeef"); err == nil {
		t.Fatal("expected checksum mismatch")
	}
}

func TestDownloadAndVerifyRawAndErrorBranches(t *testing.T) {
	dir := t.TempDir()
	src := filepath.Join(dir, "src.img")
	if err := os.WriteFile(src, []byte("raw"), 0o644); err != nil {
		t.Fatal(err)
	}
	dest := filepath.Join(dir, "dest.img")
	if _, err := DownloadAndVerify("file://"+src, dest, ""); err != nil {
		t.Fatalf("DownloadAndVerify raw: %v", err)
	}
	if err := VerifyChecksum(dest, "bogus"); err == nil {
		t.Fatal("expected invalid checksum format error")
	}

	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		http.Error(w, "boom", http.StatusInternalServerError)
	}))
	defer srv.Close()
	if err := downloadRaw(srv.URL, filepath.Join(dir, "bad.img")); err == nil {
		t.Fatal("expected HTTP error from downloadRaw")
	}

	archive := filepath.Join(dir, "empty.zip")
	if err := writeTestZIP(archive, map[string]string{}); err != nil {
		t.Fatalf("writeTestZIP: %v", err)
	}
	if err := decompressZIP(archive, filepath.Join(dir, "out.img")); err == nil {
		t.Fatal("expected empty zip error")
	}
}

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()
}