package image

import (
	"bytes"
	"fmt"
	"io"
	"log"
	"os"
	"strings"

	"github.com/kdomanski/iso9660"
)

func ModifyISOForAutoInstall(isoPath, answerURL string) error {
	origConfig, err := readGrubCfgFromISO(isoPath)
	if err != nil {
		return err
	}
	if origConfig == nil {
		log.Printf("image: no grub.cfg found in ISO, skipping auto-install modification")
		return nil
	}

	newConfig := rewriteGrubConfig(string(origConfig), answerURL)

	if len(newConfig) > len(origConfig) {
		return fmt.Errorf("modified GRUB config (%d bytes) exceeds original (%d bytes)", len(newConfig), len(origConfig))
	}

	padded := make([]byte, len(origConfig))
	copy(padded, []byte(newConfig))
	for i := len(newConfig); i < len(padded); i++ {
		padded[i] = '\n'
	}

	offset, err := findContentInISO(isoPath, origConfig)
	if err != nil {
		return fmt.Errorf("locate grub.cfg in ISO: %w", err)
	}

	f, err := os.OpenFile(isoPath, os.O_WRONLY, 0)
	if err != nil {
		return fmt.Errorf("open ISO for writing: %w", err)
	}
	defer f.Close()

	if _, err := f.WriteAt(padded, offset); err != nil {
		return fmt.Errorf("write modified grub.cfg: %w", err)
	}

	log.Printf("image: modified grub.cfg at offset %d for auto-install", offset)
	return nil
}

func readGrubCfgFromISO(isoPath string) ([]byte, error) {
	f, err := os.Open(isoPath)
	if err != nil {
		return nil, fmt.Errorf("open 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)
	}

	grubFile := findFileByName(root, "grub.cfg")
	if grubFile == nil {
		return nil, nil
	}

	reader := grubFile.Reader()
	data, err := io.ReadAll(reader)
	if err != nil {
		return nil, fmt.Errorf("read grub.cfg: %w", err)
	}
	return data, nil
}

func findFileByName(dir *iso9660.File, target string) *iso9660.File {
	children, err := dir.GetChildren()
	if err != nil {
		return nil
	}
	for _, child := range children {
		if child.IsDir() {
			if result := findFileByName(child, target); result != nil {
				return result
			}
		} else if strings.EqualFold(child.Name(), target) {
			return child
		}
	}
	return nil
}

func rewriteGrubConfig(original, answerURL string) string {
	lines := strings.Split(original, "\n")
	var result []string
	depth := 0

	for _, line := range lines {
		trimmed := strings.TrimSpace(line)

		if strings.HasPrefix(trimmed, "set timeout=") {
			result = append(result, "set timeout=0")
			continue
		}
		if strings.HasPrefix(trimmed, "set default=") {
			result = append(result, "set default=0")
			continue
		}

		if strings.HasPrefix(trimmed, "menuentry") {
			depth++
		}
		if trimmed == "}" && depth > 0 {
			depth--
		}

		if depth == 1 && strings.Contains(trimmed, "linux ") && strings.Contains(trimmed, "/boot/linux") {
			if !strings.Contains(line, "proxmox-start-auto-installer") {
				line = strings.TrimRight(line, " \t") + " proxmox-start-auto-installer"
			}
			if answerURL != "" && !strings.Contains(line, "proxmox-auto-installer-answer-url") {
				line += " proxmox-auto-installer-answer-url=" + answerURL
			}
		}

		result = append(result, line)
	}

	return strings.Join(result, "\n")
}

func findContentInISO(isoPath string, content []byte) (int64, error) {
	f, err := os.Open(isoPath)
	if err != nil {
		return 0, err
	}
	defer f.Close()

	prefix := content
	if len(prefix) > 256 {
		prefix = prefix[:256]
	}

	buf := make([]byte, 2048)
	stat, _ := f.Stat()

	for offset := int64(0); offset < stat.Size(); offset += 2048 {
		n, err := f.ReadAt(buf, offset)
		if n == 0 {
			break
		}
		if err != nil && err != io.EOF {
			return 0, err
		}
		if bytes.HasPrefix(buf[:n], prefix) {
			return offset, nil
		}
	}
	return 0, fmt.Errorf("grub.cfg content not found in ISO raw data")
}