package main

import (
	"bytes"
	"embed"
	"encoding/binary"
	"errors"
	"fmt"
	"hash/crc32"
	"io/ioutil"
	"os"
	"strconv"
	"text/tabwriter"
	"time"

	"gitee.com/wmdng/jlink"
)

func LoadProgram(jk *jlink.Jlink, prog []byte) error {

	/* halt, write CPSR */
	iret := jk.Halt()
	if iret != 0 {
		return errors.New("halt target fail")
	}

	/* read halted PC reg, set cpsr to xxx */
	_ = jk.ReadReg(9)
	jk.WriteReg(8, 0x400000D3)
	// fmt.Printf("Halt, PC = %X\n", pc)

	/* write prog, set pc = 0x100000 */
	jk.WriteMem(0x100000, prog)
	jk.WriteReg(9, 0x100000)
	jk.GoIntDis()
	return nil
}

//go:embed  aflash.bin
var efs embed.FS

func ReadAgentFile(afname string) ([]byte, error) {

	/**/
	afp, err := efs.Open("aflash.bin")
	if err != nil {
		return nil, err
	}

	/**/
	adat, err := ioutil.ReadAll(afp)
	if err != nil {
		return nil, err
	}

	/**/
	if len(adat) < 32 {
		return nil, errors.New("agent file is to small ")
	}

	/* check format, check crc */
	tmagic := [8]byte{0x73, 0x6a, 0x79, 0x5f, 0x61, 0x70, 0x70, 0x69}
	if 0 != bytes.Compare(adat[0:8], tmagic[:]) {
		return nil, errors.New("agent fmt err, magic err")
	}

	temp := binary.LittleEndian.Uint32(adat[8:])
	if int(temp)+16 > len(adat) {
		return nil, errors.New("agent fmt err, length err")
	}

	crc := crc32.Update(0, crc32.IEEETable, adat[16:16+int(temp)])
	tcrc := binary.LittleEndian.Uint32(adat[12:])
	if tcrc != crc {
		return nil, errors.New("agent fmt err, crc err")
	}

	return adat[16 : 16+int(temp)], nil
}

type burnArgs struct {
	offs  uint32 // offset in flash
	speed uint32 // jlink swd clock freq
	fbin  string // burn target file name
	rbck  uint32 // read back, 0 - for prog.
}

func printUsage() {
	w := tabwriter.NewWriter(os.Stdout, 0, 0, 8, ' ', tabwriter.TabIndent)
	fmt.Fprintln(w, "usage: ", os.Args[0], "-f file -s freq")
	fmt.Fprintln(w, "  -f <bin>\tburn target file")
	fmt.Fprintln(w, "  -s <freq>\tswd clock freq kHz, default 8000")
	fmt.Fprintln(w, "  -a <offs>\toffset in flash, default 0")
	fmt.Fprintln(w, "  --help\tdisplay this infomation")
	w.Flush()
}

func parseArgs(args []string) (*burnArgs, error) {

	var i int
	var v string
	var brg burnArgs

	// default
	brg.offs = 0
	brg.speed = 8000
	brg.fbin = ""
	brg.rbck = 0

	// open program, check format.
	for i = 1; i < len(args); i++ {

		/**/
		v = os.Args[i]

		switch {
		case v == "-h" || v == "-help" || v == "--help":
			printUsage()
			return nil, errors.New("help end")

		case v == "-f":
			i += 1
			if i >= len(args) {
				return nil, errors.New("-f target file name")
			}
			brg.fbin = args[i]

		case v == "-s":
			i += 1
			if i >= len(args) {
				return nil, errors.New("-s swd clock freq")
			}
			temp, err := strconv.ParseUint(args[i], 0, 32)
			if err != nil {
				return nil, errors.New("speed fmt err, " + err.Error())
			}

			if temp > 12000 {
				return nil, errors.New("speed value range, max 12000 kHz")
			}

			brg.speed = uint32(temp)

		case v == "-a":
			i += 1
			if i >= len(args) {
				return nil, errors.New("-a offset")
			}

			temp, err := strconv.ParseUint(args[i], 0, 32)
			if err != nil {
				return nil, errors.New("offset fmt err, " + err.Error())
			}

			brg.offs = uint32(temp)

		case v == "-rb":
			i += 1
			if i >= len(args) {
				return nil, errors.New("-rb size")
			}

			temp, err := strconv.ParseUint(args[i], 0, 32)
			if err != nil {
				return nil, errors.New("size fmt err, " + err.Error())
			}

			if temp == 0 {
				return nil, errors.New("size zero err, " + err.Error())
			}

			brg.rbck = uint32(temp)

		default:
			printUsage()
			return nil, errors.New("not support arg, " + v)
		}
	}

	if len(brg.fbin) <= 0 {
		printUsage()
		return nil, errors.New("target file name must set")
	}

	return &brg, nil
}

func main() {

	var tcrc uint32
	var pdat []byte

	brg, err := parseArgs(os.Args)
	if err != nil {
		fmt.Println("parse args,", err)
		return
	}

	/**/
	fmt.Println("jlink speed:", brg.speed)
	fmt.Println("target file:", brg.fbin)
	fmt.Println(" flash offs:", brg.offs)

	//
	jk, err := jlink.TryNewJlink()
	if err != nil {
		fmt.Println("Find jlink dll,", err)
		return
	}

	//
	err = jk.Open()
	if err != nil {
		fmt.Println("Open JLINK,", err)
		return
	}

	defer jk.Close()

	//
	err = jk.Connect(brg.speed)
	if err != nil {
		fmt.Println("Connect target,", err)
		return
	}

	if brg.rbck == 0 {

		/* open target file */
		fmt.Println(">> read target file >>")
		tfp, err := os.Open(brg.fbin)
		if err != nil {
			fmt.Println("Open file,", err)
			return
		}

		//
		defer tfp.Close()

		pdat, err := ioutil.ReadAll(tfp)
		if err != nil {
			fmt.Println("Read file,", err)
			return
		}

		if len(pdat) <= 0 {
			fmt.Println("file len zero ??")
			return
		}

		tcrc = crc32.Update(0, crc32.IEEETable, pdat[:])
		// fmt.Printf("target file, crc = %x\n", tcrc)

	} else {
		pdat = make([]byte, brg.rbck)
	}

	/* load agent, start agent */
	fmt.Println(">> load agent prog >>")
	adat, err := ReadAgentFile("aflash.bin")
	if err != nil {
		fmt.Println("Read agent,", err)
		return
	}

	err = LoadProgram(jk, adat)
	if err != nil {
		fmt.Println("Load agent,", err)
		return
	}

	/* clean DCC, (target --> jlink) */
	jk.ReadDCC(4, 0)
	/* wait target clean DCC */
	time.Sleep(100 * time.Millisecond)

	/* probe flash through agent */
	fmt.Println(">> probe flash id >>")
	sfh, err := jlink.NewFlash(jk)
	if err != nil {
		fmt.Println("probe flash fail,", err)
		return
	}

	if brg.rbck == 0 {
		/* start program flash */
		fmt.Println(">> write to flash >>")
		_, err = sfh.WriteAt(pdat[:], int64(brg.offs))
		if err != nil {
			fmt.Println("write flash fail,", err)
			return
		}

		/* read back, check */
		fmt.Println(">> read back, for check >>")
		nret, err := sfh.ReadAt(pdat[:], int64(brg.offs))
		if err != nil {
			fmt.Println("read back from flash fail,", err)
			return
		}

		if nret != len(pdat) {
			fmt.Println("read back from flash size err,", nret)
			return
		}

		/**/
		rcrc := crc32.Update(0, crc32.IEEETable, pdat[:])
		if rcrc != tcrc {
			fmt.Println("read back from flash, crc err", rcrc)
			return
		}

		/**/
		fmt.Println(">> write to flash succ >>")

	} else {

		/* read back, check */
		fmt.Println(">> read back, for check >>")
		nret, err := sfh.ReadAt(pdat[:], int64(brg.offs))
		if err != nil {
			fmt.Println("read back from flash fail,", err)
			return
		}

		if nret != len(pdat) {
			fmt.Println("read back from flash size err,", nret)
			return
		}

		/* write to file */
		err = ioutil.WriteFile(brg.fbin, pdat, 0644)
		if err != nil {
			fmt.Println("write to file err,", err)
		}

		/**/
		fmt.Println(">> read from flash succ >>")
	}

	return
}