package uf2

import (
	"encoding/binary"
	"fmt"
	"os"
	"sort"
)

const (
	blockSize = 512
	magic0    = 0x0A324655
	magic1    = 0x9E5D5157
	magicEnd  = 0x0AB16F30

	flagNotMainFlash       = 0x00000001
	flagFamilyIDPresent    = 0x00002000
	absoluteFamilyID       = 0xe48bff57
)

type Block struct {
	Addr uint32
	Data []byte
}

func Parse(path string) ([]Block, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return nil, err
	}
	if len(data)%blockSize != 0 {
		return nil, fmt.Errorf("file size %d not multiple of %d", len(data), blockSize)
	}

	var blocks []Block
	for i := 0; i < len(data); i += blockSize {
		b := data[i : i+blockSize]
		m0 := binary.LittleEndian.Uint32(b[0:4])
		m1 := binary.LittleEndian.Uint32(b[4:8])
		me := binary.LittleEndian.Uint32(b[508:512])
		if m0 != magic0 || m1 != magic1 || me != magicEnd {
			return nil, fmt.Errorf("block %d: bad magic", i/blockSize)
		}
		flags := binary.LittleEndian.Uint32(b[8:12])
		if flags&flagNotMainFlash != 0 {
			continue
		}
		if flags&flagFamilyIDPresent != 0 {
			familyID := binary.LittleEndian.Uint32(b[28:32])
			if familyID == absoluteFamilyID {
				continue
			}
		}
		addr := binary.LittleEndian.Uint32(b[12:16])
		size := binary.LittleEndian.Uint32(b[16:20])
		if size > 256 {
			return nil, fmt.Errorf("block %d: data size %d > 256", i/blockSize, size)
		}
		blocks = append(blocks, Block{
			Addr: addr,
			Data: make([]byte, size),
		})
		copy(blocks[len(blocks)-1].Data, b[32:32+size])
	}

	sort.Slice(blocks, func(i, j int) bool {
		return blocks[i].Addr < blocks[j].Addr
	})

	return blocks, nil
}