package image

import (
	"archive/zip"
	"crypto/md5"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
)

// Download fetches url into dest if dest does not exist.
func Download(url, dest string) error {
	_, err := DownloadAndVerify(url, dest, "")
	return err
}

// DownloadAndVerify fetches the source image, verifies it when a checksum is provided,
// and returns the local raw image path ready for copying or injection.
func DownloadAndVerify(url, dest, checksum string) (string, error) {
	if err := os.MkdirAll(filepath.Dir(dest), 0o755); err != nil {
		return "", err
	}
	if strings.HasSuffix(url, ".xz") {
		archive := dest + ".xz"
		if err := ensureVerifiedFile(url, archive, checksum); err != nil {
			return "", err
		}
		if err := decompressXZ(archive, dest); err != nil {
			return "", err
		}
		return dest, nil
	}
	if strings.HasSuffix(url, ".zip") {
		archive := dest + ".zip"
		if err := ensureVerifiedFile(url, archive, checksum); err != nil {
			return "", err
		}
		if err := decompressZIP(archive, dest); err != nil {
			return "", err
		}
		return dest, nil
	}
	if err := ensureVerifiedFile(url, dest, checksum); err != nil {
		return "", err
	}
	return dest, nil
}

func ensureVerifiedFile(url, dest, checksum string) error {
	if _, err := os.Stat(dest); err == nil {
		if err := VerifyChecksum(dest, checksum); err == nil || checksum == "" {
			return nil
		}
		if removeErr := os.Remove(dest); removeErr != nil && !errors.Is(removeErr, os.ErrNotExist) {
			return removeErr
		}
	} else if !errors.Is(err, os.ErrNotExist) {
		return err
	}
	if err := downloadRaw(url, dest); err != nil {
		return err
	}
	return VerifyChecksum(dest, checksum)
}

func downloadRaw(url, dest string) error {
	if strings.HasPrefix(url, "file://") {
		src := strings.TrimPrefix(url, "file://")
		in, err := os.Open(src)
		if err != nil {
			return err
		}
		defer in.Close()
		out, err := os.Create(dest)
		if err != nil {
			return err
		}
		defer out.Close()
		_, err = io.Copy(out, in)
		return err
	}
	resp, err := http.Get(url)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("download failed: %s", resp.Status)
	}
	out, err := os.Create(dest)
	if err != nil {
		return err
	}
	defer out.Close()
	_, err = io.Copy(out, resp.Body)
	return err
}

func decompressXZ(src, dest string) error {
	out, err := os.Create(dest)
	if err != nil {
		return err
	}
	defer out.Close()
	cmd := exec.Command("xz", "-dc", src)
	cmd.Stdout = out
	var stderr strings.Builder
	cmd.Stderr = &stderr
	if err := cmd.Run(); err != nil {
		return fmt.Errorf("xz decompress %s: %w: %s", src, err, stderr.String())
	}
	return out.Sync()
}

func decompressZIP(src, dest string) error {
	reader, err := zip.OpenReader(src)
	if err != nil {
		return err
	}
	defer reader.Close()

	var target *zip.File
	for _, file := range reader.File {
		if file.FileInfo().IsDir() {
			continue
		}
		if strings.HasSuffix(strings.ToLower(file.Name), ".img") {
			target = file
			break
		}
		if target == nil {
			target = file
		}
	}
	if target == nil {
		return fmt.Errorf("zip archive %s does not contain a file entry", src)
	}

	in, err := target.Open()
	if err != nil {
		return err
	}
	defer in.Close()

	out, err := os.Create(dest)
	if err != nil {
		return err
	}
	defer out.Close()
	if _, err := io.Copy(out, in); err != nil {
		return err
	}
	return out.Sync()
}

// VerifyChecksum checks hashes in the form "sha256:<hex>" or "md5:<hex>".
func VerifyChecksum(path, checksum string) error {
	if checksum == "" {
		return nil
	}
	parts := strings.SplitN(checksum, ":", 2)
	if len(parts) != 2 {
		return errors.New("unsupported checksum format; use sha256:<hex> or md5:<hex>")
	}
	algo := strings.ToLower(strings.TrimSpace(parts[0]))
	expected := strings.ToLower(parts[1])
	f, err := os.Open(path)
	if err != nil {
		return err
	}
	defer f.Close()
	var sum string
	switch algo {
	case "sha256":
		h := sha256.New()
		if _, err := io.Copy(h, f); err != nil {
			return err
		}
		sum = hex.EncodeToString(h.Sum(nil))
	case "md5":
		h := md5.New()
		if _, err := io.Copy(h, f); err != nil {
			return err
		}
		sum = hex.EncodeToString(h.Sum(nil))
	default:
		return errors.New("unsupported checksum format; use sha256:<hex> or md5:<hex>")
	}
	if sum != expected {
		return fmt.Errorf("checksum mismatch: expected %s got %s", expected, sum)
	}
	return nil
}