package platform

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
)

const installToRAMDir = "/dev/shm/bee-live"

func (s *System) IsLiveMediaInRAM() bool {
	return s.LiveMediaRAMState().InRAM
}

func (s *System) LiveBootSource() LiveBootSource {
	fsType := mountFSType("/run/live/medium")
	source := mountSource("/run/live/medium")
	device := findLiveBootDevice()
	status := LiveBootSource{
		InRAM:  strings.EqualFold(fsType, "tmpfs"),
		Source: source,
		Device: device,
	}
	if fsType == "" && source == "" && device == "" {
		if toramActive() {
			status.InRAM = true
			status.Kind = "ram"
			status.Source = "tmpfs"
			return status
		}
		status.Kind = "unknown"
		return status
	}
	status.Kind = inferLiveBootKind(fsType, source, blockDeviceType(device), blockDeviceTransport(device))
	if status.Kind == "" {
		status.Kind = "unknown"
	}
	if status.InRAM && strings.TrimSpace(status.Source) == "" {
		status.Source = "tmpfs"
	}
	return status
}

func (s *System) LiveMediaRAMState() LiveMediaRAMState {
	return evaluateLiveMediaRAMState(
		s.LiveBootSource(),
		toramActive(),
		globPaths("/run/live/medium/live/*.squashfs"),
		globPaths(filepath.Join(installToRAMDir, "*.squashfs")),
	)
}

func evaluateLiveMediaRAMState(status LiveBootSource, toram bool, sourceSquashfs, copiedSquashfs []string) LiveMediaRAMState {
	state := LiveMediaRAMState{
		LiveBootSource: status,
		ToramActive:    toram,
		CopyPresent:    len(copiedSquashfs) > 0,
	}
	if status.InRAM {
		state.State = "in_ram"
		state.Status = "ok"
		state.CopyComplete = true
		state.Message = "Running from RAM — installation media can be safely disconnected."
		return state
	}

	expected := pathBaseSet(sourceSquashfs)
	copied := pathBaseSet(copiedSquashfs)
	state.CopyComplete = len(expected) > 0 && setContainsAll(copied, expected)

	switch {
	case state.CopyComplete:
		state.State = "partial"
		state.Status = "partial"
		state.CanStartCopy = true
		state.Message = "Live media files were copied to RAM, but the system is still mounted from the original boot source."
	case state.CopyPresent:
		state.State = "partial"
		state.Status = "partial"
		state.CanStartCopy = true
		state.Message = "Partial RAM copy detected. A previous Copy to RAM run was interrupted or cancelled."
	case toram:
		state.State = "toram_failed"
		state.Status = "failed"
		state.CanStartCopy = true
		state.Message = "toram boot parameter is set but the live medium is not mounted from RAM."
	default:
		state.State = "not_in_ram"
		state.Status = "warning"
		state.CanStartCopy = true
		state.Message = "ISO not copied to RAM. Use Copy to RAM to free the boot drive and improve performance."
	}
	return state
}

func globPaths(pattern string) []string {
	matches, _ := filepath.Glob(pattern)
	return matches
}

func pathBaseSet(paths []string) map[string]struct{} {
	out := make(map[string]struct{}, len(paths))
	for _, path := range paths {
		base := strings.TrimSpace(filepath.Base(path))
		if base != "" {
			out[base] = struct{}{}
		}
	}
	return out
}

func setContainsAll(have, want map[string]struct{}) bool {
	if len(want) == 0 {
		return false
	}
	for name := range want {
		if _, ok := have[name]; !ok {
			return false
		}
	}
	return true
}

func (s *System) RunInstallToRAM(ctx context.Context, logFunc func(string)) (retErr error) {
	log := func(msg string) {
		if logFunc != nil {
			logFunc(msg)
		}
	}

	state := s.LiveMediaRAMState()
	if state.InRAM {
		log("Already running from RAM — installation media can be safely disconnected.")
		return nil
	}

	squashfsFiles, err := filepath.Glob("/run/live/medium/live/*.squashfs")
	if err != nil || len(squashfsFiles) == 0 {
		return fmt.Errorf("no squashfs files found in /run/live/medium/live/")
	}

	free := freeMemBytes()
	var needed int64
	for _, sf := range squashfsFiles {
		fi, err2 := os.Stat(sf)
		if err2 != nil {
			return fmt.Errorf("stat %s: %v", sf, err2)
		}
		needed += fi.Size()
	}
	const headroom = 256 * 1024 * 1024
	if free > 0 && needed+headroom > free {
		return fmt.Errorf("insufficient RAM: need %s, available %s",
			humanBytes(needed+headroom), humanBytes(free))
	}

	dstDir := installToRAMDir
	if state.CopyPresent {
		log("Removing stale partial RAM copy before retry...")
	}
	_ = os.RemoveAll(dstDir)
	if err := os.MkdirAll(dstDir, 0755); err != nil {
		return fmt.Errorf("create tmpfs dir: %v", err)
	}
	defer func() {
		if retErr == nil {
			return
		}
		_ = os.RemoveAll(dstDir)
		log("Removed incomplete RAM copy.")
	}()

	for _, sf := range squashfsFiles {
		if err := ctx.Err(); err != nil {
			return err
		}
		base := filepath.Base(sf)
		dst := filepath.Join(dstDir, base)
		log(fmt.Sprintf("Copying %s to RAM...", base))
		if err := copyFileLarge(ctx, sf, dst, log); err != nil {
			return fmt.Errorf("copy %s: %v", base, err)
		}
		log(fmt.Sprintf("Copied %s.", base))

		loopDev, err := findLoopForFile(sf)
		if err != nil {
			log(fmt.Sprintf("Loop device for %s not found (%v) — skipping re-association.", base, err))
			continue
		}
		if err := reassociateLoopDevice(loopDev, dst); err != nil {
			log(fmt.Sprintf("Warning: could not re-associate %s → %s: %v", loopDev, dst, err))
		} else {
			log(fmt.Sprintf("Loop device %s now backed by RAM copy.", loopDev))
		}
	}

	log("Copying remaining medium files...")
	if err := cpDir(ctx, "/run/live/medium", dstDir, log); err != nil {
		log(fmt.Sprintf("Warning: partial copy: %v", err))
	}
	if err := ctx.Err(); err != nil {
		return err
	}

	mediumRebound := false
	if err := bindMount(dstDir, "/run/live/medium"); err != nil {
		log(fmt.Sprintf("Warning: rebind /run/live/medium → %s failed: %v", dstDir, err))
	} else {
		mediumRebound = true
	}

	log("Verifying live medium now served from RAM...")
	status := s.LiveBootSource()
	if err := verifyInstallToRAMStatus(status, dstDir, mediumRebound, log); err != nil {
		return err
	}
	if status.InRAM {
		log(fmt.Sprintf("Verification passed: live medium now served from %s.", describeLiveBootSource(status)))
	}
	log("Done. Squashfs files are in RAM. Installation media can be safely disconnected.")
	return nil
}

func verifyInstallToRAMStatus(status LiveBootSource, dstDir string, mediumRebound bool, log func(string)) error {
	if status.InRAM {
		return nil
	}

	// The live medium mount was not redirected to RAM. This is expected when
	// booting from an ISO/CD-ROM: the squashfs loop device has a non-zero
	// offset and LOOP_CHANGE_FD cannot be used; the bind mount also fails
	// because the CD-ROM mount is in use. Check whether files were at least
	// copied to the tmpfs directory — that is sufficient for safe disconnection
	// once the kernel has paged in all actively-used data.
	files, _ := filepath.Glob(filepath.Join(dstDir, "*.squashfs"))
	if len(files) > 0 {
		if !mediumRebound {
			log(fmt.Sprintf("Note: squashfs copied to RAM (%s) but /run/live/medium still shows the original source.", dstDir))
			log("This is normal for CD-ROM boots. For a fully transparent RAM boot, add 'toram' to the kernel parameters.")
		}
		return nil
	}

	return fmt.Errorf("install to RAM verification failed: live medium still mounted from %s and no squashfs found in %s", describeLiveBootSource(status), dstDir)
}

func describeLiveBootSource(status LiveBootSource) string {
	source := strings.TrimSpace(status.Device)
	if source == "" {
		source = strings.TrimSpace(status.Source)
	}
	if source == "" {
		source = "unknown source"
	}
	switch strings.TrimSpace(status.Kind) {
	case "ram":
		return "RAM"
	case "usb":
		return "USB (" + source + ")"
	case "cdrom":
		return "CD-ROM (" + source + ")"
	case "disk":
		return "disk (" + source + ")"
	default:
		return source
	}
}

func copyFileLarge(ctx context.Context, src, dst string, logFunc func(string)) error {
	in, err := os.Open(src)
	if err != nil {
		return err
	}
	defer in.Close()
	fi, err := in.Stat()
	if err != nil {
		return err
	}
	out, err := os.Create(dst)
	if err != nil {
		return err
	}
	defer out.Close()
	total := fi.Size()
	var copied int64
	buf := make([]byte, 4*1024*1024)
	for {
		if err := ctx.Err(); err != nil {
			return err
		}
		n, err := in.Read(buf)
		if n > 0 {
			if _, werr := out.Write(buf[:n]); werr != nil {
				return werr
			}
			copied += int64(n)
			if logFunc != nil && total > 0 {
				pct := int(float64(copied) / float64(total) * 100)
				logFunc(fmt.Sprintf("  %s / %s (%d%%)", humanBytes(copied), humanBytes(total), pct))
			}
		}
		if err == io.EOF {
			break
		}
		if err != nil {
			return err
		}
	}
	return out.Sync()
}

func cpDir(ctx context.Context, src, dst string, logFunc func(string)) error {
	return filepath.Walk(src, func(path string, fi os.FileInfo, err error) error {
		if ctx.Err() != nil {
			return ctx.Err()
		}
		if err != nil {
			return nil
		}
		rel, _ := filepath.Rel(src, path)
		target := filepath.Join(dst, rel)
		if fi.IsDir() {
			return os.MkdirAll(target, fi.Mode())
		}
		if strings.HasSuffix(path, ".squashfs") {
			return nil
		}
		if _, err := os.Stat(target); err == nil {
			return nil
		}
		return copyFileLarge(ctx, path, target, nil)
	})
}

func findLoopForFile(backingFile string) (string, error) {
	out, err := exec.Command("losetup", "--list", "--json").Output()
	if err != nil {
		return "", err
	}
	var result struct {
		Loopdevices []struct {
			Name     string `json:"name"`
			BackFile string `json:"back-file"`
		} `json:"loopdevices"`
	}
	if err := json.Unmarshal(out, &result); err != nil {
		return "", err
	}
	for _, dev := range result.Loopdevices {
		if dev.BackFile == backingFile {
			return dev.Name, nil
		}
	}
	return "", fmt.Errorf("no loop device found for %s", backingFile)
}

// loopDeviceOffset returns the byte offset configured for the loop device,
// or -1 if it cannot be determined.
func loopDeviceOffset(loopDev string) int64 {
	out, err := exec.Command("losetup", "--json", loopDev).Output()
	if err != nil {
		return -1
	}
	var result struct {
		Loopdevices []struct {
			Offset int64 `json:"offset"`
		} `json:"loopdevices"`
	}
	if err := json.Unmarshal(out, &result); err != nil || len(result.Loopdevices) == 0 {
		return -1
	}
	return result.Loopdevices[0].Offset
}

func reassociateLoopDevice(loopDev, newFile string) error {
	// LOOP_CHANGE_FD requires lo_offset == 0. ISO/CD-ROM loop devices are
	// typically set up with a non-zero offset (squashfs lives inside the ISO),
	// so the ioctl returns EINVAL. Detect this early for a clear error message.
	if off := loopDeviceOffset(loopDev); off > 0 {
		return fmt.Errorf("loop device has non-zero offset (%d bytes, typical for ISO/CD-ROM) — LOOP_CHANGE_FD not supported; use 'toram' kernel parameter for RAM boot", off)
	}
	if err := exec.Command("losetup", "--replace", loopDev, newFile).Run(); err == nil {
		return nil
	}
	return loopChangeFD(loopDev, newFile)
}