package main

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"io"
	"log"
	"sort"
	"time"

	"github.com/karalabe/hid"
	"github.com/ulikunitz/xz"
)

const (
	CMD_VIA_VIAL_PREFIX      = 0xFE
	CMD_VIAL_GET_KEYBOARD_ID = 0x00
	CMD_VIAL_GET_SIZE        = 0x01
	CMD_VIAL_GET_DEFINITION  = 0x02
	MSG_LEN                  = 32 // 適切なMSG_LENを設定してください
)

type Device struct {
	dev           *hid.Device
	vial_protocol uint32
	keyboard_id   uint64
}

func (d *Device) hidSend(dev *hid.Device, msg []byte, retries int) ([]byte, error) {
	if len(msg) > MSG_LEN {
		return nil, fmt.Errorf("message must be less than 32 bytes")
	}
	msg = append(msg, make([]byte, MSG_LEN-len(msg))...)

	data := make([]byte, MSG_LEN)
	first := true

	for retries > 0 {
		retries--
		if !first {
			time.Sleep(500 * time.Millisecond)
		}
		first = false
		if _, err := dev.Write(append([]byte{0x00}, msg...)); err != nil {
			continue
		}
		n, err := dev.Read(data)
		if err != nil || n == 0 {
			continue
		}
		break
	}

	if len(data) == 0 {
		return nil, fmt.Errorf("failed to communicate with the device")
	}
	return data, nil
}

func (d *Device) reloadLayout() ([]byte, error) {

	// get keyboard identification
	data, err := d.hidSend(d.dev, []byte{CMD_VIA_VIAL_PREFIX, CMD_VIAL_GET_KEYBOARD_ID}, 20)
	if err != nil {
		return nil, err
	}
	log.Println("get keyboard identification", data)
	buf := bytes.NewReader(data[:12])
	if err := binary.Read(buf, binary.LittleEndian, &d.vial_protocol); err != nil {
		return nil, err
	}
	if err := binary.Read(buf, binary.LittleEndian, &d.keyboard_id); err != nil {
		return nil, err
	}

	// get the size
	data, err = d.hidSend(d.dev, []byte{CMD_VIA_VIAL_PREFIX, CMD_VIAL_GET_SIZE}, 20)
	if err != nil {
		return nil, err
	}
	var sz int32
	log.Println("get the size", data)
	buf = bytes.NewReader(data[:4])
	if err := binary.Read(buf, binary.LittleEndian, &sz); err != nil {
		return nil, err
	}
	log.Println("size:", sz)

	// get the payload
	var lzmaPayload []byte
	var block uint32
	for sz > 0 {
		data, err = d.hidSend(d.dev, append([]byte{CMD_VIA_VIAL_PREFIX, CMD_VIAL_GET_DEFINITION}, byte(block)), 20)
		if err != nil {
			return nil, err
		}
		log.Println("current:", sz)
		if sz < MSG_LEN {
			data = data[:sz]
		}
		lzmaPayload = append(lzmaPayload, data...)
		block++
		sz -= MSG_LEN
	}

	log.Println("received payload")

	// lzma解凍
	dec, err := xz.NewReader(bytes.NewReader(lzmaPayload))
	if err != nil {
		return nil, err
	}
	decoded, err := io.ReadAll(dec)
	if err != nil {
		return nil, fmt.Errorf("failed to decode: %w", err)
	}

	return decoded, nil
}

func main() {
	var (
		vendorID  uint16 = 0x5957
		productID uint16 = 0x0100
	)
	devices := hid.Enumerate(vendorID, productID) // 適切なベンダーIDとプロダクトIDを設定してください
	if len(devices) == 0 {
		log.Fatal("No HID devices found")
	}

	sort.Slice(devices, func(i, j int) bool {
		return devices[i].UsagePage > devices[j].UsagePage
	})

	deviceInfo := devices[0]
	dev, err := deviceInfo.Open()
	if err != nil {
		log.Fatalf("Failed to open device: %v", err)
	}
	defer dev.Close()

	device := &Device{
		dev: dev,
	}
	payload, err := device.reloadLayout()
	if err != nil {
		log.Fatalf("Error: %v", err)
	}
	fmt.Println(string(payload))
}