package image

import (
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"

	"github.com/kdomanski/iso9660"
)

type ExtractResult struct {
	KernelFilename string
	InitrdFilename string
	ISOFilename    string
}

type ProgressFunc func(stage string, detail string)

var kernelCandidates = []string{"linux26", "vmlinuz", "bzImage"}
var initrdCandidates = []string{"initrd.img", "initrd", "initrd.gz"}

func ExtractFromISO(r io.Reader, destDir string) (*ExtractResult, error) {
	return ExtractFromISOWithProgress(r, destDir, nil)
}

func ExtractFromISOWithProgress(r io.Reader, destDir string, progress ProgressFunc) (*ExtractResult, error) {
	report := func(stage, detail string) {
		if progress != nil {
			progress(stage, detail)
		}
	}

	report("receiving", "Writing ISO to disk...")

	isoPath := filepath.Join(destDir, "original.iso")
	tmp, err := os.Create(isoPath)
	if err != nil {
		return nil, fmt.Errorf("create ISO file: %w", err)
	}

	if _, err := io.Copy(tmp, r); err != nil {
		tmp.Close()
		return nil, fmt.Errorf("write ISO to disk: %w", err)
	}
	tmp.Close()

	report("parsing", "Parsing ISO image...")

	f, err := os.Open(isoPath)
	if err != nil {
		return nil, fmt.Errorf("open temp ISO: %w", err)
	}
	defer f.Close()

	img, err := iso9660.OpenImage(f)
	if err != nil {
		return nil, fmt.Errorf("parse ISO: %w", err)
	}

	root, err := img.RootDir()
	if err != nil {
		return nil, fmt.Errorf("read ISO root: %w", err)
	}

	candidateSet := make(map[string]bool)
	for _, c := range kernelCandidates {
		candidateSet[strings.ToLower(c)] = true
	}
	for _, c := range initrdCandidates {
		candidateSet[strings.ToLower(c)] = true
	}

	found := make(map[string]*iso9660.File)
	walkDir(root, candidateSet, found)

	kernelFile, kernelName := matchFirst(found, kernelCandidates)
	initrdFile, initrdName := matchFirst(found, initrdCandidates)

	if kernelFile == nil {
		return nil, fmt.Errorf("no kernel found in ISO (looked for %s)", strings.Join(kernelCandidates, ", "))
	}
	if initrdFile == nil {
		return nil, fmt.Errorf("no initrd found in ISO (looked for %s)", strings.Join(initrdCandidates, ", "))
	}

	report("extracting", "Extracting kernel...")

	if err := extractFile(kernelFile, filepath.Join(destDir, kernelName)); err != nil {
		return nil, fmt.Errorf("extract kernel: %w", err)
	}

	report("extracting", "Extracting initrd...")

	if err := extractFile(initrdFile, filepath.Join(destDir, initrdName)); err != nil {
		return nil, fmt.Errorf("extract initrd: %w", err)
	}

	report("complete", "Extraction complete")

	return &ExtractResult{
		KernelFilename: kernelName,
		InitrdFilename: initrdName,
		ISOFilename:    "original.iso",
	}, nil
}

func walkDir(dir *iso9660.File, candidates map[string]bool, found map[string]*iso9660.File) {
	children, err := dir.GetChildren()
	if err != nil {
		return
	}
	for _, child := range children {
		name := strings.ToLower(child.Name())
		if child.IsDir() {
			walkDir(child, candidates, found)
		} else if candidates[name] {
			if _, exists := found[name]; !exists {
				found[name] = child
			}
		}
	}
}

func matchFirst(found map[string]*iso9660.File, candidates []string) (*iso9660.File, string) {
	for _, c := range candidates {
		lower := strings.ToLower(c)
		if f, ok := found[lower]; ok {
			return f, f.Name()
		}
	}
	return nil, ""
}

func extractFile(isoFile *iso9660.File, destPath string) error {
	out, err := os.Create(destPath)
	if err != nil {
		return err
	}
	defer out.Close()

	reader := isoFile.Reader()
	if reader == nil {
		return fmt.Errorf("cannot read %s from ISO", isoFile.Name())
	}
	_, err = io.Copy(out, reader)
	return err
}