package utils

import (
	"fmt"
	"io"
	"os"
	"strconv"
	"sync"
	"syscall"
	"up/utils/mtd"
)

var uploadState struct {
	sync.RWMutex
	Progress int
}

func WriteUbootToDevices(data []byte, filepath string) error {
	UBOOT_MAX_SIZE, err := strconv.Atoi(os.Getenv("UBOOT_MAX_SIZE"))
	if err != nil {
		fmt.Println("error: ", err)
		return err
	}
	MMC_BOOT_OFFSET, err := strconv.Atoi(os.Getenv("MMC_BOOT_OFFSET"))
	if err != nil {
		fmt.Println("error: ", err)
		return err
	}
	const chunkSize = 4096
	dataLen := int64(len(data))

	fmt.Printf("Writing U-Boot to devices (%d bytes)...\n", dataLen)

	if dataLen > int64(UBOOT_MAX_SIZE*1024*1024) {
		return fmt.Errorf("U-Boot image too large for SPI NOR (%d > %d bytes)", dataLen, UBOOT_MAX_SIZE)
	}

	f, err := os.Create(filepath)
	if err != nil {
		return fmt.Errorf("failed to create file %s: %w", filepath, err)
	}

	var written int64
	for written < dataLen {
		toWrite := min(dataLen-written, chunkSize)
		n, err := f.Write(data[written : written+toWrite])
		if err != nil {
			f.Close()
			return fmt.Errorf("failed to write to file: %w", err)
		}
		written += int64(n)
		uploadState.Lock()
		uploadState.Progress = int(written * 30 / dataLen)
		uploadState.Unlock()
	}
	f.Close()
	syscall.Sync()

	fmt.Println("Writing to /dev/mtd0 (SPI NOR)...")
	mtdFd, err := os.OpenFile("/dev/mtd0", os.O_RDWR|os.O_SYNC, 0660)
	if err == nil {
		defer mtdFd.Close()

		mtdSize, err := mtd.ReadMTDSize(0)
		if err != nil {
			fmt.Printf("Cannot get MTD size from sysfs, proceeding without erase: %v\n", err)
		} else {
			fmt.Printf("Erasing SPI NOR (size %d bytes)...\n", mtdSize)
			if err := mtd.EraseMTD(mtdFd, 0, mtdSize); err != nil {
				fmt.Printf("Warning: erase failed (continuing anyway): %v\n", err)
			} else {
				fmt.Println("SPI NOR erased successfully")
			}
		}

		written = 0
		for written < dataLen {
			toWrite := min(dataLen-written, chunkSize)
			n, err := mtdFd.Write(data[written : written+toWrite])
			if err != nil {
				fmt.Printf("Error writing to /dev/mtd0: %v\n", err)
				break
			}
			written += int64(n)
			uploadState.Lock()
			uploadState.Progress = 30 + int(written*35/dataLen)
			uploadState.Unlock()
		}
		fmt.Printf("Written %d bytes to /dev/mtd0\n", written)
	} else {
		fmt.Printf("/dev/mtd0 not available: %v (skipping SPI NOR)\n", err)
	}

	fmt.Printf("Writing to /dev/mmcblk1 at offset %d (SD card boot area)...\n", MMC_BOOT_OFFSET)
	mmcFd, err := os.OpenFile("/dev/mmcblk1", os.O_RDWR|os.O_SYNC, 0660)
	if err == nil {
		defer mmcFd.Close()

		if _, err := mmcFd.Seek(int64(MMC_BOOT_OFFSET), io.SeekStart); err != nil {
			fmt.Printf("Failed to seek on /dev/mmcblk1: %v\n", err)
		} else {
			written = 0
			for written < dataLen {
				toWrite := min(dataLen-written, chunkSize)
				n, err := mmcFd.Write(data[written : written+toWrite])
				if err != nil {
					fmt.Printf("Error writing to /dev/mmcblk1: %v\n", err)
					break
				}
				written += int64(n)
				uploadState.Lock()
				uploadState.Progress = 65 + int(written*35/dataLen)
				uploadState.Unlock()
			}
			fmt.Printf("Written %d bytes to /dev/mmcblk1 at offset %d\n", written, MMC_BOOT_OFFSET)
		}
	} else {
		fmt.Printf("/dev/mmcblk1 not available: %v (skipping SD card)\n", err)
	}

	syscall.Sync()

	uploadState.Lock()
	uploadState.Progress = 100
	uploadState.Unlock()

	return nil
}