package usb

import (
	"encoding/binary"
	"fmt"
)

// USB descriptor types
const (
	DescTypeDevice        = 0x01
	DescTypeConfiguration = 0x02
	DescTypeInterface     = 0x04
	DescTypeEndpoint      = 0x05
)

// ParsedDescriptors holds everything we extract from a device's raw
// descriptor blob (device descriptor followed by all configuration
// descriptors, as returned by reading a usbdevfs device file).
type ParsedDescriptors struct {
	VendorID       uint16
	ProductID      uint16
	BcdDevice      uint16
	DeviceClass    uint8
	DeviceSubClass uint8
	DeviceProtocol uint8
	NumConfigs     uint8

	// Configs holds every configuration, each with every interface
	// alternate setting and its endpoints.
	Configs []ConfigDescriptor
}

// ConfigDescriptor is one USB configuration
type ConfigDescriptor struct {
	Value      uint8       // bConfigurationValue
	Interfaces []Interface // every alternate setting, in descriptor order
}

// ParseDescriptors parses a raw descriptor blob: an 18-byte device
// descriptor followed by one or more complete configuration descriptors.
//
// Reading a usbdevfs file (/dev/bus/usb/BBB/DDD) from offset 0 yields
// exactly this layout, which is the only way to see interface alternate
// settings — sysfs only exposes the currently active one.
func ParseDescriptors(data []byte) (*ParsedDescriptors, error) {
	if len(data) < 18 {
		return nil, fmt.Errorf("descriptor blob too short: %d bytes", len(data))
	}
	if data[1] != DescTypeDevice {
		return nil, fmt.Errorf("first descriptor is type 0x%02x, expected device (0x01)", data[1])
	}

	pd := &ParsedDescriptors{
		DeviceClass:    data[4],
		DeviceSubClass: data[5],
		DeviceProtocol: data[6],
		VendorID:       binary.LittleEndian.Uint16(data[8:10]),
		ProductID:      binary.LittleEndian.Uint16(data[10:12]),
		BcdDevice:      binary.LittleEndian.Uint16(data[12:14]),
		NumConfigs:     data[17],
	}

	// Walk the remaining descriptors. Configuration descriptors start a new
	// config; interface descriptors start a new alternate setting; endpoint
	// descriptors attach to the most recent interface. Class-specific
	// descriptors (HID, UVC, audio) are skipped by their bLength.
	pos := int(data[0]) // skip the device descriptor using its own bLength
	if pos < 18 {
		pos = 18
	}

	var curConfig *ConfigDescriptor
	var curIface *Interface

	for pos+2 <= len(data) {
		bLength := int(data[pos])
		bType := data[pos+1]

		// A zero-length descriptor would loop forever; a descriptor running
		// past the end of the blob means the device returned garbage.
		if bLength < 2 || pos+bLength > len(data) {
			break
		}

		switch bType {
		case DescTypeConfiguration:
			if bLength >= 9 {
				pd.Configs = append(pd.Configs, ConfigDescriptor{Value: data[pos+5]})
				curConfig = &pd.Configs[len(pd.Configs)-1]
				curIface = nil
			}

		case DescTypeInterface:
			if bLength >= 9 && curConfig != nil {
				curConfig.Interfaces = append(curConfig.Interfaces, Interface{
					Number:     data[pos+2],
					AltSetting: data[pos+3],
					Class:      data[pos+5],
					SubClass:   data[pos+6],
					Protocol:   data[pos+7],
				})
				curIface = &curConfig.Interfaces[len(curConfig.Interfaces)-1]
			}

		case DescTypeEndpoint:
			if bLength >= 7 && curIface != nil {
				curIface.Endpoints = append(curIface.Endpoints, Endpoint{
					Address:       data[pos+2],
					TransferType:  data[pos+3] & 0x03,
					MaxPacketSize: binary.LittleEndian.Uint16(data[pos+4 : pos+6]),
					Interval:      data[pos+6],
				})
			}
		}

		pos += bLength
	}

	if len(pd.Configs) == 0 {
		return nil, fmt.Errorf("no configuration descriptor found")
	}

	return pd, nil
}

// FindConfig returns the configuration with the given bConfigurationValue,
// or nil if the device has no such configuration.
func (pd *ParsedDescriptors) FindConfig(value uint8) *ConfigDescriptor {
	for i := range pd.Configs {
		if pd.Configs[i].Value == value {
			return &pd.Configs[i]
		}
	}
	return nil
}

// AllEndpoints returns every endpoint across every alternate setting of the
// given configuration, keyed by full bEndpointAddress (direction bit
// included). Endpoints only present in a non-zero alternate setting — the
// isochronous endpoints of webcams, for example — are included, which is
// what makes the endpoint type map correct after a SET_INTERFACE.
func (c *ConfigDescriptor) AllEndpoints() map[uint8]Endpoint {
	eps := make(map[uint8]Endpoint)
	for _, iface := range c.Interfaces {
		for _, ep := range iface.Endpoints {
			// Alternate settings reuse addresses with identical transfer
			// types in practice; keep the first one we see so alt 0 wins.
			if _, seen := eps[ep.Address]; !seen {
				eps[ep.Address] = ep
			}
		}
	}
	return eps
}

// ActiveInterfaces returns one Interface per interface number, using
// alternate setting 0 — the set of interfaces that must be claimed.
func (c *ConfigDescriptor) ActiveInterfaces() []Interface {
	var result []Interface
	seen := make(map[uint8]bool)
	for _, iface := range c.Interfaces {
		if iface.AltSetting != 0 || seen[iface.Number] {
			continue
		}
		seen[iface.Number] = true
		result = append(result, iface)
	}
	return result
}