//
// Copyright (c) 2014-2019 Cesanta Software Limited
// All rights reserved
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
package xds110

// Client for the TI XDS110 debug probe.

import (
	"bytes"
	"context"
	"encoding/binary"
	"fmt"
	"strings"
	"time"

	"github.com/google/gousb"
	"github.com/juju/errors"
	"github.com/mongoose-os/mos/cli/flash/common"
	glog "k8s.io/klog/v2"
)

const (
	vendorTI      = 0x0451
	productXDS110 = 0xbef3
	intfNum       = 0x02
	endpointOut   = 0x02
	endpointIn    = 0x83
	maxPacketLen  = 0x1000

	shortTimeout = 200 * time.Millisecond
)

type XDS110Client struct {
	ctx     *gousb.Context
	dev     *gousb.Device
	cfg     *gousb.Config
	intf    *gousb.Interface
	inEndp  *gousb.InEndpoint
	outEndp *gousb.OutEndpoint
}

func NewXDS110Client(serial string) (*XDS110Client, error) {
	ok := false
	xc := &XDS110Client{}
	defer func() {
		if !ok {
			xc.Close()
		}
	}()

	var err error
	xc.ctx, xc.dev, err = common.OpenUSBDevice(vendorTI, productXDS110, serial)
	if err != nil {
		return nil, errors.Annotatef(err, "failed to open XDS110 probe")
	}

	cfgNum, err := xc.dev.ActiveConfigNum()
	if err != nil {
		return nil, errors.Annotatef(err, "failed to get active config number of device %s", xc.dev)
	}
	xc.cfg, err = xc.dev.Config(cfgNum)
	if err != nil {
		return nil, errors.Annotatef(err, "failed to claim config %d of device %s", cfgNum, xc.dev)
	}
	xc.intf, err = xc.cfg.Interface(intfNum, 0)
	if err != nil {
		return nil, errors.Annotatef(err, "couldn't open interface %+v", xc.cfg)
	}

	xc.outEndp, err = xc.intf.OutEndpoint(endpointOut)
	if err != nil {
		return nil, errors.Annotatef(err, "couldn't open output endpoint")
	}

	xc.inEndp, err = xc.intf.InEndpoint(endpointIn)
	if err != nil {
		return nil, errors.Annotatef(err, "couldn't open input endpoint")
	}

	ok = true
	return xc, nil
}

type xds110Command uint8

const (
	headerByte uint8 = 0x2a
)

const (
	cmdEcho           = 0x00
	cmdConnect        = 0x01
	cmdDisconnect     = 0x02
	cmdGetVersionInfo = 0x03
	cmdSetTCLKDelay   = 0x04
	cmdSetTRST        = 0x05
	cmdGetTRST        = 0x06
	cmdSetSRST        = 0x0e
	cmdReadMemory     = 0x13
	cmdWriteMemory    = 0x14
)

func (xc *XDS110Client) GetSerialNumber() string {
	sn, _ := xc.dev.SerialNumber()
	// Sometimes serial contains trailing NULs. Not sure whose fault it is (libusb?), strip them.
	return strings.TrimRight(sn, "\x00")
}

func (xc *XDS110Client) doCommand(cmd xds110Command, argBuf *bytes.Buffer, respPayloadLen int, timeout time.Duration) (*bytes.Buffer, error) {
	if xc.outEndp == nil {
		return nil, errors.Errorf("not connected")
	}
	if argBuf == nil {
		argBuf = bytes.NewBuffer(nil)
	}
	glog.V(2).Infof("=> %s %q", cmd, argBuf.Bytes())
	cmdBuf := bytes.NewBuffer([]byte{headerByte})
	binary.Write(cmdBuf, binary.LittleEndian, uint16(len(argBuf.Bytes()))+1)
	binary.Write(cmdBuf, binary.LittleEndian, &cmd)
	cmdBuf.Write(argBuf.Bytes())
	glog.V(4).Infof("=> (%d) %q", len(cmdBuf.Bytes()), cmdBuf.Bytes())
	_, err := xc.outEndp.Write(cmdBuf.Bytes())
	if err != nil {
		return nil, errors.Annotatef(err, "failed to send command")
	}
	ctx := context.TODO()
	if timeout != 0 {
		ctx2, done := context.WithTimeout(ctx, timeout)
		ctx = ctx2
		defer done()
	}
	resp := make([]byte, maxPacketLen)
	n, err := xc.inEndp.ReadContext(ctx, resp)
	if err != nil {
		return nil, errors.Annotatef(err, "failed to read response")
	}
	glog.V(4).Infof("<= (%d) %q", n, resp[:n])
	if n < 7 {
		return nil, errors.Errorf("response is too short (%d) %q", n, resp[:n])
	}
	if resp[0] != headerByte {
		return nil, errors.Errorf("invalid response header")
	}
	respBuf := bytes.NewBuffer(resp[1:])
	var packetLen uint16
	binary.Read(respBuf, binary.LittleEndian, &packetLen)
	var status uint32
	binary.Read(respBuf, binary.LittleEndian, &status)
	glog.V(2).Infof("<= %d %q", status, resp[7:n])
	if status != 0 {
		return nil, errors.Errorf("error response: %d", status)
	}
	if n-3 != int(packetLen) {
		return nil, errors.Errorf("incomplete packet: expected %d, got %d", packetLen, n-3)
	}
	if int(packetLen) != respPayloadLen+4 {
		return nil, errors.Errorf("incomplete response: expected %d, got %d", respPayloadLen, packetLen-4)
	}
	return bytes.NewBuffer(resp[7:n]), nil
}

func (xc *XDS110Client) Echo(data []byte) error {
	if len(data) > 255 {
		return errors.Errorf("Echo: max 255 bytes, got %d", len(data))
	}
	ab := bytes.NewBuffer(nil)
	binary.Write(ab, binary.LittleEndian, uint32(len(data)&0xff))
	ab.Write(data)
	echoData, err := xc.doCommand(cmdEcho, ab, len(data), shortTimeout)
	if err == nil && !bytes.Equal(echoData.Bytes(), data) {
		return errors.Errorf("invalid echo response: send %q, got %q", data, echoData.Bytes())
	}
	return errors.Annotatef(err, "Echo")
}

func (xc *XDS110Client) Connect() error {
	_, err := xc.doCommand(cmdConnect, nil, 0, shortTimeout)
	return errors.Annotatef(err, "Connect")
}

func (xc *XDS110Client) Disconnect() error {
	_, err := xc.doCommand(cmdDisconnect, nil, 7, shortTimeout)
	return errors.Annotatef(err, "Disconnect")
}

type XDS110VersionInfo struct {
	Version        uint32
	V1, V2, V3, V4 uint8
	HWVersion      uint16
}

func (xc *XDS110Client) GetVersionInfo() (*XDS110VersionInfo, error) {
	rb, err := xc.doCommand(cmdGetVersionInfo, nil, 6, shortTimeout)
	if err != nil {
		return nil, errors.Annotatef(err, "GetVersionInfo")
	}
	vi := &XDS110VersionInfo{}
	binary.Read(rb, binary.LittleEndian, &vi.Version)
	vi.V1 = uint8(vi.Version >> 24)
	vi.V2 = uint8(vi.Version >> 16)
	vi.V3 = uint8(vi.Version >> 8)
	vi.V4 = uint8(vi.Version & 0xff)
	binary.Read(rb, binary.LittleEndian, &vi.HWVersion)
	return vi, nil
}

func (xc *XDS110Client) SetTCLKDelay(delay uint8) error {
	ab := bytes.NewBuffer(nil)
	binary.Write(ab, binary.LittleEndian, uint32(delay))
	_, err := xc.doCommand(cmdSetTCLKDelay, ab, 0, shortTimeout)
	return errors.Annotatef(err, "SetTCLKDelay")
}

func boolToByte(b bool) byte {
	if b {
		return 1
	} else {
		return 0
	}
}

func (xc *XDS110Client) SetTRST(rstOn bool) error {
	ab := bytes.NewBuffer([]byte{boolToByte(rstOn)})
	_, err := xc.doCommand(cmdSetTRST, ab, 0, shortTimeout)
	return errors.Annotatef(err, "SetTRST")
}

func (xc *XDS110Client) SetSRST(rstOn bool) error {
	ab := bytes.NewBuffer([]byte{boolToByte(!rstOn)})
	_, err := xc.doCommand(cmdSetSRST, ab, 0, shortTimeout)
	return errors.Annotatef(err, "SetSRST")
}

func (xc *XDS110Client) Close() error {
	if xc.intf != nil {
		xc.inEndp = nil
		xc.outEndp = nil
		xc.intf.Close()
		xc.intf = nil
	}
	if xc.cfg != nil {
		xc.cfg.Close()
		xc.cfg = nil
	}
	if xc.dev != nil {
		xc.dev.Close()
		xc.dev = nil
	}
	if xc.ctx != nil {
		xc.ctx.Close()
		xc.ctx = nil
	}
	return nil
}

// TODO(rojer): Use stringer when it actually works.
func (cmd xds110Command) String() string {
	switch cmd {
	case cmdEcho:
		return fmt.Sprintf("Echo(0x%x)", uint8(cmd))
	case cmdConnect:
		return fmt.Sprintf("Connect(0x%x)", uint8(cmd))
	case cmdDisconnect:
		return fmt.Sprintf("Disconnect(0x%x)", uint8(cmd))
	case cmdGetVersionInfo:
		return fmt.Sprintf("GetVersionInfo(0x%x)", uint8(cmd))
	case cmdSetTCLKDelay:
		return fmt.Sprintf("SetTCLKDelay(0x%x)", uint8(cmd))
	case cmdSetTRST:
		return fmt.Sprintf("SetTRST(0x%x)", uint8(cmd))
	case cmdGetTRST:
		return fmt.Sprintf("GetTRST(0x%x)", uint8(cmd))
	case cmdSetSRST:
		return fmt.Sprintf("SetSRST(0x%x)", uint8(cmd))
	case cmdReadMemory:
		return fmt.Sprintf("ReadMemory(0x%x)", uint8(cmd))
	case cmdWriteMemory:
		return fmt.Sprintf("WriteMemory(0x%x)", uint8(cmd))
	default:
		return fmt.Sprintf("?(0x%x)", uint8(cmd))
	}
}
