//go:build linux

package usbip

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"sync"
	"time"
	"unsafe"

	"github.com/duffy/usb-server/internal/usb"
)

// controlQueueDepth bounds how many control transfers may be waiting for the
// control worker. Control transfers are serialised because endpoint 0 is a
// single shared pipe; the queue only exists so that a slow transfer does not
// stall the URB read loop behind it.
const controlQueueDepth = 64

// controlTimeout is the per-transfer timeout passed to USBDEVFS_CONTROL.
const controlTimeout = 5000 // ms

// reapPollInterval is how long the reap loop waits for a completed URB before
// re-checking whether it should shut down.
const reapPollInterval = 100 * time.Millisecond

// Server handles USB/IP protocol on the share side.
// It manages a single USB device and forwards URBs between
// the USB/IP client (via tunnel) and the physical device (via usbdevfs).
type Server struct {
	device      *usb.Device
	handle      *usb.DeviceHandle
	mu          sync.Mutex
	pendingURBs map[uint32]*pendingURB // seqnum -> pending URB

	// unlinkedURBs holds URBs that were discarded but not yet reaped.
	//
	// USBDEVFS_DISCARDURB is asynchronous: the kernel still owns the URB
	// struct and its transfer buffer, and will write the completion status
	// into them. Dropping the last Go reference at unlink time would let the
	// garbage collector reclaim memory the kernel is about to write to.
	unlinkedURBs map[uint32]*pendingURB

	closed bool

	// epTypes maps a full bEndpointAddress (direction bit included) to a
	// usbdevfs URB type. Indexing by address rather than endpoint number
	// matters: a composite device can have endpoint 1 as interrupt IN (0x81)
	// and endpoint 1 as bulk OUT (0x01), and submitting an interrupt URB with
	// the bulk type makes the kernel reject it.
	epTypes map[uint8]uint8

	// ctrlQueue serialises control transfers on a dedicated worker so that a
	// blocking USBDEVFS_CONTROL ioctl never stalls the protocol read loop.
	ctrlQueue chan *ctrlRequest

	// stop is closed to shut down the reap loop and control worker; workers
	// signals when both have exited so Detach can safely close the fd.
	stop    chan struct{}
	workers sync.WaitGroup
	runOnce sync.Once
}

type pendingURB struct {
	seqNum     uint32
	devID      uint32
	direction  uint32
	endpoint   uint32
	buffer     []byte
	urbPtr     unsafe.Pointer // pointer to submitted usbdevfs_urb
	isISO      bool
	numPackets int32
	isoMem     []byte   // keeps ISO URB+descriptors memory alive for GC
	packetLens []uint32 // original request lengths per ISO packet (for offset computation)
}

// ctrlRequest is one queued control transfer.
type ctrlRequest struct {
	hdr         *URBHeader
	body        *CmdSubmitBody
	transferBuf []byte
}

// NewServer creates a USB/IP server for a specific device
func NewServer(dev *usb.Device) *Server {
	return &Server{
		device:       dev,
		pendingURBs:  make(map[uint32]*pendingURB),
		unlinkedURBs: make(map[uint32]*pendingURB),
		epTypes:      make(map[uint8]uint8),
		ctrlQueue:    make(chan *ctrlRequest, controlQueueDepth),
		stop:         make(chan struct{}),
	}
}

// Attach opens the device, disconnects the kernel driver, and claims all interfaces
func (s *Server) Attach() error {
	handle, err := usb.OpenDevice(s.device.DevPath, s.device.BusID)
	if err != nil {
		return fmt.Errorf("opening device: %w", err)
	}
	s.handle = handle

	// Build endpoint type map from device info (must be done before driver detach)
	s.buildEndpointTypeMap()

	// For each interface: disconnect kernel driver and claim it
	for _, iface := range s.device.Interfaces {
		ifnum := uint32(iface.Number)

		// Try atomic disconnect+claim first (USBDEVFS_DISCONNECT_CLAIM, Linux 3.7+)
		err := handle.DisconnectClaimInterface(ifnum)
		if err == nil {
			log.Printf("[usbip-server] interface %d: disconnect+claim OK", ifnum)
			continue
		}

		// Fallback: disconnect driver per interface, then claim
		if disconnErr := handle.DisconnectDriverForInterface(ifnum); disconnErr != nil {
			log.Printf("[usbip-server] interface %d: disconnect warning: %v", ifnum, disconnErr)
		}

		if claimErr := handle.ClaimInterface(ifnum); claimErr != nil {
			log.Printf("[usbip-server] error: could not claim interface %d: %v", ifnum, claimErr)
			// This is a critical error - clean up and fail
			for _, prev := range s.device.Interfaces {
				if uint32(prev.Number) < ifnum {
					handle.ReleaseInterface(uint32(prev.Number))
				}
			}
			handle.ConnectDriver()
			handle.Close()
			s.handle = nil
			return fmt.Errorf("claiming interface %d: %w", ifnum, claimErr)
		}
		log.Printf("[usbip-server] interface %d: fallback disconnect+claim OK", ifnum)
	}

	return nil
}

// Detach releases all interfaces, closes the device, and rebinds kernel drivers.
func (s *Server) Detach() {
	s.mu.Lock()
	alreadyClosed := s.closed
	s.closed = true
	s.mu.Unlock()

	if alreadyClosed || s.handle == nil {
		return
	}

	// 1. Stop the reap loop and control worker, then wait for them to exit.
	// This must happen before closing the fd: a worker mid-ioctl on a closed
	// fd would either fail confusingly or, worse, operate on a recycled fd.
	s.stopWorkers()

	// 2. Discard all pending URBs to clean up device state.
	// The maps stay populated on purpose: they are what keeps the URB
	// structs and transfer buffers reachable while the kernel still owns
	// them. They are only cleared after the fd is closed, which is the point
	// at which the kernel definitively drops its references.
	s.mu.Lock()
	for _, pending := range s.pendingURBs {
		if pending.urbPtr != nil {
			s.handle.DiscardURBByPtr(pending.urbPtr)
		}
	}
	s.mu.Unlock()

	// 3. Release all claimed interfaces
	for _, iface := range s.device.Interfaces {
		if err := s.handle.ReleaseInterface(uint32(iface.Number)); err != nil {
			log.Printf("[usbip-server] release interface %d: %v", iface.Number, err)
		}
	}

	// 4. Close the device file descriptor.
	// The kernel auto-cancels remaining URBs on close.
	s.handle.Close()
	s.handle = nil

	s.mu.Lock()
	s.pendingURBs = make(map[uint32]*pendingURB)
	s.unlinkedURBs = make(map[uint32]*pendingURB)
	s.mu.Unlock()

	// 5. Force kernel driver re-binding via sysfs authorized toggle.
	// After USBDEVFS_DISCONNECT_CLAIM, the kernel sets privileges_dropped=true.
	// This means closing the fd does NOT auto-rebind drivers.
	// Also USBDEVFS_RESET after ReleaseInterface doesn't rebind because
	// the kernel sets needs_binding=false on release.
	// The reliable solution: toggle authorized 0->1 which forces complete
	// re-enumeration and driver binding.
	s.rebindDrivers()
}

// stopWorkers signals the reap loop and control worker to exit and waits for
// them. Safe to call more than once.
func (s *Server) stopWorkers() {
	s.runOnce.Do(func() { close(s.stop) })
	s.workers.Wait()
}

// rebindDrivers forces the kernel to re-bind drivers to the device
// by toggling the sysfs authorized attribute.
func (s *Server) rebindDrivers() {
	authPath := filepath.Join(s.device.SysPath, "authorized")

	// Deauthorize: kernel disconnects device, unbinds all drivers
	if err := os.WriteFile(authPath, []byte("0"), 0644); err != nil {
		log.Printf("[usbip-server] sysfs deauthorize failed: %v, trying fallback", err)
		s.rebindDriversFallback()
		return
	}

	// Brief delay for the kernel to process the deauthorization
	time.Sleep(100 * time.Millisecond)

	// Re-authorize: kernel re-enumerates device, binds drivers
	if err := os.WriteFile(authPath, []byte("1"), 0644); err != nil {
		log.Printf("[usbip-server] sysfs re-authorize failed: %v", err)
		return
	}

	log.Printf("[usbip-server] device re-authorized, kernel drivers re-bound")
}

// rebindDriversFallback tries alternative methods to rebind drivers
func (s *Server) rebindDriversFallback() {
	// Try writing bus_id to each original driver's bind file
	for _, iface := range s.device.Interfaces {
		if iface.Driver == "" {
			continue
		}
		ifaceName := fmt.Sprintf("%s:%d.%d", s.device.BusID, s.device.ConfigValue, iface.Number)
		bindPath := filepath.Join("/sys/bus/usb/drivers", iface.Driver, "bind")
		if err := os.WriteFile(bindPath, []byte(ifaceName), 0644); err != nil {
			log.Printf("[usbip-server] bind %s to %s failed: %v", ifaceName, iface.Driver, err)
		} else {
			log.Printf("[usbip-server] re-bound %s to driver %s", ifaceName, iface.Driver)
		}
	}
}

// urbTypeName maps a usbdevfs URB type to a short label for logging.
var urbTypeName = map[uint8]string{
	usbdevfsTypeISO:       "ISO",
	usbdevfsTypeInterrupt: "INT",
	usbdevfsTypeControl:   "CTRL",
	usbdevfsTypeBulk:      "BULK",
}

// usbdevfs URB types (mirrors the constants in the usb package)
const (
	usbdevfsTypeISO       = 0
	usbdevfsTypeInterrupt = 1
	usbdevfsTypeControl   = 2
	usbdevfsTypeBulk      = 3
)

// buildEndpointTypeMap builds the endpoint address -> URB type map.
//
// It prefers Device.Endpoints, which is parsed from the raw descriptors and
// therefore covers every alternate setting. Endpoints that only appear in a
// non-zero alternate setting — the isochronous endpoints of webcams, which
// only activate after SET_INTERFACE — would be missing otherwise.
func (s *Server) buildEndpointTypeMap() {
	record := func(ep usb.Endpoint) {
		var urbType uint8
		switch ep.TransferType {
		case usb.TransferTypeControl:
			urbType = usbdevfsTypeControl
		case usb.TransferTypeIsochronous:
			urbType = usbdevfsTypeISO
		case usb.TransferTypeBulk:
			urbType = usbdevfsTypeBulk
		case usb.TransferTypeInterrupt:
			urbType = usbdevfsTypeInterrupt
		default:
			urbType = usbdevfsTypeBulk
		}
		s.epTypes[ep.Address] = urbType

		dir := "OUT"
		if ep.IsIn() {
			dir = "IN"
		}
		log.Printf("[usbip-server] endpoint 0x%02x (EP%d %s): %s maxpkt=%d interval=%d",
			ep.Address, ep.Number(), dir, urbTypeName[urbType], ep.MaxPacketSize, ep.Interval)
	}

	if len(s.device.Endpoints) > 0 {
		for _, ep := range s.device.Endpoints {
			record(ep)
		}
		return
	}

	// Fallback for devices whose raw descriptors could not be read.
	log.Printf("[usbip-server] warning: no parsed descriptors, falling back to sysfs endpoints")
	for _, iface := range s.device.Interfaces {
		for _, ep := range iface.Endpoints {
			record(ep)
		}
	}
}

// getURBType returns the usbdevfs URB type for a full endpoint address.
//
// interval and numPackets come from the incoming CMD_SUBMIT and act as a
// fallback for endpoints missing from the descriptor map: only periodic
// transfers carry a non-zero interval, so an unknown endpoint with one is an
// interrupt endpoint rather than a bulk endpoint. Guessing bulk there is what
// breaks HID devices, whose interrupt URBs the kernel then rejects.
func (s *Server) getURBType(epAddr uint8, interval uint32, numPackets int32) uint8 {
	if epAddr&0x0F == 0 {
		return usbdevfsTypeControl
	}
	if numPackets > 0 {
		return usbdevfsTypeISO
	}
	if t, ok := s.epTypes[epAddr]; ok {
		return t
	}
	if interval > 0 {
		return usbdevfsTypeInterrupt
	}
	return usbdevfsTypeBulk
}

// BuildDeviceDescriptor creates a USB/IP device descriptor from our device info
func (s *Server) BuildDeviceDescriptor() DeviceDescriptor {
	var desc DeviceDescriptor
	SetPath(&desc.Path, s.device.SysPath)
	SetBusID(&desc.BusID, s.device.BusID)
	desc.BusNum = s.device.BusNum
	desc.DevNum = s.device.DevNum
	desc.Speed = s.device.Speed
	desc.IDVendor = s.device.VendorID
	desc.IDProduct = s.device.ProductID
	desc.BcdDevice = s.device.BcdDevice
	desc.BDeviceClass = s.device.DeviceClass
	desc.BDeviceSubClass = s.device.DeviceSubClass
	desc.BDeviceProtocol = s.device.DeviceProtocol
	desc.BConfigurationValue = s.device.ConfigValue
	desc.BNumConfigurations = s.device.NumConfigs
	desc.BNumInterfaces = uint8(len(s.device.Interfaces))
	return desc
}

// BuildInterfaceDescriptors creates USB/IP interface descriptors
func (s *Server) BuildInterfaceDescriptors() []InterfaceDescriptor {
	var descs []InterfaceDescriptor
	for _, iface := range s.device.Interfaces {
		descs = append(descs, InterfaceDescriptor{
			BInterfaceClass:    iface.Class,
			BInterfaceSubClass: iface.SubClass,
			BInterfaceProtocol: iface.Protocol,
		})
	}
	return descs
}

// HandleConnection processes USB/IP protocol on a bidirectional stream.
// It reads USB/IP requests from the reader, processes them, and writes responses to the writer.
// This is the main loop for handling a connected USB/IP client.
func (s *Server) HandleConnection(r io.Reader, w io.Writer) error {
	retChan := make(chan []byte, 256)

	// The reap loop and control worker outlive this function only until
	// Detach stops them; both feed retChan.
	s.workers.Add(2)
	go s.reapLoop(retChan)
	go s.controlWorker(retChan)

	// Forward completed URBs to the writer. This goroutine belongs to the
	// connection, not to the server, so it ends when the connection does.
	connDone := make(chan struct{})
	defer close(connDone)
	go func() {
		for {
			select {
			case data := <-retChan:
				if _, err := w.Write(data); err != nil {
					return
				}
			case <-connDone:
				return
			case <-s.stop:
				return
			}
		}
	}()

	// Read and process incoming USB/IP messages
	for {
		hdr, err := ReadURBHeader(r)
		if err != nil {
			if err == io.EOF || errors.Is(err, io.ErrUnexpectedEOF) || errors.Is(err, io.ErrClosedPipe) {
				return nil
			}
			return fmt.Errorf("reading URB header: %w", err)
		}

		switch hdr.Command {
		case CmdSubmit:
			if err := s.handleCmdSubmit(r, hdr, retChan); err != nil {
				return fmt.Errorf("handling CMD_SUBMIT: %w", err)
			}
		case CmdUnlink:
			if err := s.handleCmdUnlink(r, hdr, retChan); err != nil {
				return fmt.Errorf("handling CMD_UNLINK: %w", err)
			}
		default:
			return fmt.Errorf("unknown URB command: 0x%08x", hdr.Command)
		}
	}
}

// send queues a response, dropping it if the server is shutting down rather
// than blocking forever on a channel nobody is draining.
func (s *Server) send(retChan chan<- []byte, resp []byte) {
	select {
	case retChan <- resp:
	case <-s.stop:
	}
}

func (s *Server) handleCmdSubmit(r io.Reader, hdr *URBHeader, retChan chan<- []byte) error {
	body, err := ReadCmdSubmit(r)
	if err != nil {
		return err
	}

	// Read transfer buffer for OUT direction
	var transferBuf []byte
	if hdr.Direction == DirOut && body.TransferBufferLen > 0 {
		transferBuf = make([]byte, body.TransferBufferLen)
		if _, err := io.ReadFull(r, transferBuf); err != nil {
			return fmt.Errorf("reading transfer buffer: %w", err)
		}
	}

	// Read ISO packet descriptors if present
	var isoDescs []ISOPacketDescriptor
	numPackets := int32(0)
	if body.NumberOfPackets != 0xFFFFFFFF && body.NumberOfPackets > 0 {
		numPackets = int32(body.NumberOfPackets)
		isoDescs = make([]ISOPacketDescriptor, numPackets)
		if err := binary.Read(r, binary.BigEndian, &isoDescs); err != nil {
			return fmt.Errorf("reading ISO descriptors: %w", err)
		}
	}

	endpoint := uint8(hdr.Endpoint)

	// Control transfers go to the dedicated worker: USBDEVFS_CONTROL is a
	// blocking ioctl and running it inline would stall every later URB behind
	// a transfer that can take up to controlTimeout milliseconds.
	if endpoint == 0 {
		req := &ctrlRequest{hdr: hdr, body: body, transferBuf: transferBuf}
		select {
		case s.ctrlQueue <- req:
		case <-s.stop:
		default:
			log.Printf("[usbip-server] control queue full, stalling on seq=%d", hdr.SeqNum)
			select {
			case s.ctrlQueue <- req:
			case <-s.stop:
			}
		}
		return nil
	}

	ep := endpoint
	if hdr.Direction == DirIn {
		ep |= 0x80
	}

	urbType := s.getURBType(ep, body.Interval, numPackets)

	// Handle isochronous transfers.
	// Trust the USB/IP NumberOfPackets field rather than our endpoint type map,
	// because a webcam only activates its ISO endpoints after SET_INTERFACE.
	if numPackets > 0 {
		return s.handleISOSubmit(hdr, body, transferBuf, isoDescs, numPackets, ep, retChan)
	}

	// Bulk and interrupt transfers
	var buf []byte
	if hdr.Direction == DirIn {
		buf = make([]byte, body.TransferBufferLen)
	} else {
		buf = transferBuf
	}

	urb, err := s.handle.SubmitURB(&usb.SubmitURBParams{
		Type:        urbType,
		Endpoint:    ep,
		Flags:       0,
		Buffer:      buf,
		UserContext: uintptr(hdr.SeqNum),
	})
	if err != nil {
		log.Printf("[usbip-server] SubmitURB(ep=0x%02x, type=%s, len=%d, interval=%d) FAILED: %v",
			ep, urbTypeName[urbType], len(buf), body.Interval, err)
		resp, _ := BuildRetSubmit(hdr.SeqNum, hdr.DevID, hdr.Direction, hdr.Endpoint, -32, 0, nil)
		s.send(retChan, resp)
		return nil
	}

	s.mu.Lock()
	s.pendingURBs[hdr.SeqNum] = &pendingURB{
		seqNum:    hdr.SeqNum,
		devID:     hdr.DevID,
		direction: hdr.Direction,
		endpoint:  hdr.Endpoint,
		buffer:    buf,
		urbPtr:    unsafe.Pointer(urb),
	}
	s.mu.Unlock()

	return nil
}

// controlWorker executes queued control transfers one at a time.
// Endpoint 0 is a single shared pipe, so serialising is both correct and
// what the device expects; the queue exists purely to decouple these
// blocking ioctls from the protocol read loop.
func (s *Server) controlWorker(retChan chan<- []byte) {
	defer s.workers.Done()

	for {
		select {
		case <-s.stop:
			return
		case req := <-s.ctrlQueue:
			resp := s.doControlTransfer(req)
			if resp != nil {
				s.send(retChan, resp)
			}
		}
	}
}

// doControlTransfer performs one control transfer and builds its RET_SUBMIT.
func (s *Server) doControlTransfer(req *ctrlRequest) []byte {
	hdr, body := req.hdr, req.body

	bmRequestType := body.Setup[0]
	bRequest := body.Setup[1]
	wValue := binary.LittleEndian.Uint16(body.Setup[2:4])
	wIndex := binary.LittleEndian.Uint16(body.Setup[4:6])
	wLength := binary.LittleEndian.Uint16(body.Setup[6:8])

	s.mu.Lock()
	handle := s.handle
	closed := s.closed
	s.mu.Unlock()
	if closed || handle == nil {
		return nil
	}

	if hdr.Direction == DirIn {
		buf := make([]byte, body.TransferBufferLen)
		n, err := handle.ControlTransfer(bmRequestType, bRequest, wValue, wIndex, wLength, controlTimeout, buf)
		var status int32
		if err != nil {
			log.Printf("[usbip-server] CTRL IN seq=%d bmReqType=0x%02x bReq=0x%02x wVal=0x%04x failed: %v",
				hdr.SeqNum, bmRequestType, bRequest, wValue, err)
			status = -32 // -EPIPE
			n = 0
		}
		resp, err := BuildRetSubmit(hdr.SeqNum, hdr.DevID, hdr.Direction, hdr.Endpoint, status, uint32(n), buf[:n])
		if err != nil {
			return nil
		}
		return resp
	}

	var status int32
	var actualLength uint32

	// Intercept standard USB requests that require special usbdevfs ioctls.
	// Raw control transfers via USBDEVFS_CONTROL don't update kernel state.
	switch {
	case bmRequestType == 0x01 && bRequest == 0x0B:
		// SET_INTERFACE (Standard, Interface recipient)
		// MUST use USBDEVFS_SETINTERFACE so the kernel updates endpoint state
		// and allocates bandwidth for ISO endpoints (critical for webcams).
		if err := handle.SetInterface(uint32(wIndex), uint32(wValue)); err != nil {
			log.Printf("[usbip-server] SET_INTERFACE(iface=%d, alt=%d) failed: %v", wIndex, wValue, err)
			status = -32 // -EPIPE
		} else {
			log.Printf("[usbip-server] SET_INTERFACE(iface=%d, alt=%d) OK", wIndex, wValue)
		}

	case bmRequestType == 0x02 && bRequest == 0x01 && wValue == 0x0000:
		// CLEAR_FEATURE(ENDPOINT_HALT) (Standard, Endpoint recipient)
		if err := handle.ClearHalt(uint32(wIndex)); err != nil {
			log.Printf("[usbip-server] CLEAR_HALT(ep=0x%02x) failed: %v", wIndex, err)
			status = -32
		}

	case bmRequestType == 0x00 && bRequest == 0x09:
		// SET_CONFIGURATION — do NOT forward to the physical device.
		// The device is already configured (we claimed interfaces during Attach).
		// Sending SET_CONFIGURATION via raw USBDEVFS_CONTROL would reset the
		// device's endpoint state without updating the kernel's internal USB
		// subsystem, breaking all subsequent SETINTERFACE and SUBMITURB calls
		// (ESRCH / EHOSTUNREACH).
		// Do NOT reset host-side data toggles either: after DisconnectClaimInterface
		// the host and device toggles are already in sync. Resetting host-side
		// toggles to DATA0 would create a mismatch (device still at its current
		// toggle), causing the first interrupt packet to be silently discarded.
		log.Printf("[usbip-server] SET_CONFIGURATION(%d) intercepted (device already configured)", wValue)

	default:
		// Generic OUT control transfer
		buf := req.transferBuf
		if buf == nil {
			buf = make([]byte, 0)
		}
		n, err := handle.ControlTransfer(bmRequestType, bRequest, wValue, wIndex, wLength, controlTimeout, buf)
		if err != nil {
			log.Printf("[usbip-server] CTRL OUT seq=%d bmReqType=0x%02x bReq=0x%02x wVal=0x%04x failed: %v",
				hdr.SeqNum, bmRequestType, bRequest, wValue, err)
			status = -32 // -EPIPE
		} else {
			// actualLength must be reported for OUT transfers too: the kernel
			// UVC driver checks it (a VS_PROBE SET_CUR expects 26).
			actualLength = uint32(n)
		}
	}

	resp, err := BuildRetSubmit(hdr.SeqNum, hdr.DevID, hdr.Direction, hdr.Endpoint, status, actualLength, nil)
	if err != nil {
		return nil
	}
	return resp
}

// handleISOSubmit handles isochronous URB submission
func (s *Server) handleISOSubmit(hdr *URBHeader, body *CmdSubmitBody, transferBuf []byte,
	isoDescs []ISOPacketDescriptor, numPackets int32, ep uint8, retChan chan<- []byte) error {

	// Collect packet lengths and compute total buffer size
	packetLens := make([]uint32, numPackets)
	var totalBufLen uint32
	for i := int32(0); i < numPackets; i++ {
		packetLens[i] = isoDescs[i].Length
		totalBufLen += isoDescs[i].Length
	}

	// Prepare buffer
	var buf []byte
	if hdr.Direction == DirIn {
		buf = make([]byte, totalBufLen)
	} else {
		// For OUT: the incoming data is packed (compact), expand to sequential layout
		buf = make([]byte, totalBufLen)
		if transferBuf != nil {
			srcOff := uint32(0)
			dstOff := uint32(0)
			for i := int32(0); i < numPackets; i++ {
				pktLen := packetLens[i]
				if srcOff+pktLen <= uint32(len(transferBuf)) {
					copy(buf[dstOff:dstOff+pktLen], transferBuf[srcOff:srcOff+pktLen])
				}
				srcOff += pktLen
				dstOff += pktLen
			}
		}
	}

	urb, isoMem, err := s.handle.SubmitISOURB(&usb.SubmitISOURBParams{
		Endpoint:        ep,
		Flags:           0x02, // URB_ISO_ASAP
		Buffer:          buf,
		NumberOfPackets: numPackets,
		PacketLengths:   packetLens,
		UserContext:     uintptr(hdr.SeqNum),
	})
	if err != nil {
		log.Printf("[usbip-server] ISO submit FAILED (ep=0x%02x pkts=%d buf=%d): %v",
			ep, numPackets, totalBufLen, err)
		resp, _ := BuildRetSubmit(hdr.SeqNum, hdr.DevID, hdr.Direction, hdr.Endpoint, -32, 0, nil)
		s.send(retChan, resp)
		return nil
	}

	s.mu.Lock()
	s.pendingURBs[hdr.SeqNum] = &pendingURB{
		seqNum:     hdr.SeqNum,
		devID:      hdr.DevID,
		direction:  hdr.Direction,
		endpoint:   hdr.Endpoint,
		buffer:     buf,
		urbPtr:     unsafe.Pointer(urb),
		isISO:      true,
		numPackets: numPackets,
		isoMem:     isoMem,
		packetLens: packetLens,
	}
	s.mu.Unlock()

	return nil
}

func (s *Server) handleCmdUnlink(r io.Reader, hdr *URBHeader, retChan chan<- []byte) error {
	body, err := ReadCmdUnlink(r)
	if err != nil {
		return err
	}

	s.mu.Lock()
	pending, exists := s.pendingURBs[body.UnlinkSeqNum]
	if exists {
		delete(s.pendingURBs, body.UnlinkSeqNum)
		// Keep the URB and its buffer reachable until the kernel hands it
		// back through the reap loop; discarding is asynchronous.
		s.unlinkedURBs[body.UnlinkSeqNum] = pending
	}
	handle := s.handle
	s.mu.Unlock()

	// -ECONNRESET tells the client the URB was actually cancelled; 0 means it
	// had already completed, which is also a valid outcome.
	var status int32
	if exists && pending.urbPtr != nil && handle != nil {
		if err := handle.DiscardURBByPtr(pending.urbPtr); err == nil {
			status = -104 // -ECONNRESET
		}
	}

	resp, err := BuildRetUnlink(hdr.SeqNum, hdr.DevID, status)
	if err != nil {
		return err
	}
	s.send(retChan, resp)

	return nil
}

// reapLoop collects completed URBs and turns them into RET_SUBMIT responses.
func (s *Server) reapLoop(retChan chan<- []byte) {
	defer s.workers.Done()

	for {
		select {
		case <-s.stop:
			return
		default:
		}

		s.mu.Lock()
		closed := s.closed
		handle := s.handle
		s.mu.Unlock()
		if closed || handle == nil {
			return
		}

		// Wait for a completion rather than spinning on a non-blocking reap.
		ready, err := handle.WaitForURB(reapPollInterval)
		if err != nil {
			if errors.Is(err, usb.ErrDeviceGone) {
				log.Printf("[usbip-server] device gone, stopping reap loop")
				return
			}
			log.Printf("[usbip-server] reap poll error: %v", err)
			// Back off so a persistent poll error cannot become a busy loop.
			select {
			case <-time.After(reapPollInterval):
			case <-s.stop:
				return
			}
			continue
		}
		if !ready {
			continue // timeout, re-check shutdown
		}

		urbInfo, err := handle.ReapURBInfoNonBlock()
		if err != nil {
			if errors.Is(err, usb.ErrNoURBReady) {
				continue
			}
			if errors.Is(err, usb.ErrDeviceGone) {
				log.Printf("[usbip-server] device gone, stopping reap loop")
				return
			}
			log.Printf("[usbip-server] reap error: %v", err)
			continue
		}

		seqNum := uint32(urbInfo.UserContext)

		s.mu.Lock()
		pending, exists := s.pendingURBs[seqNum]
		if exists {
			delete(s.pendingURBs, seqNum)
		} else if _, wasUnlinked := s.unlinkedURBs[seqNum]; wasUnlinked {
			// The kernel is done with it; the memory may now be reclaimed.
			delete(s.unlinkedURBs, seqNum)
		}
		s.mu.Unlock()

		if !exists {
			// Already unlinked; the client is not expecting a RET_SUBMIT.
			continue
		}

		if urbInfo.Status != 0 {
			dirStr := "OUT"
			if pending.direction == DirIn {
				dirStr = "IN"
			}
			log.Printf("[usbip-server] URB error: seq=%d EP%d %s status=%d actual=%d",
				pending.seqNum, pending.endpoint, dirStr, urbInfo.Status, urbInfo.ActualLength)
		}

		var resp []byte
		if pending.isISO {
			resp, err = s.buildISOResponse(urbInfo, pending)
		} else {
			var data []byte
			if pending.direction == DirIn && urbInfo.ActualLength > 0 {
				n := int(urbInfo.ActualLength)
				if n > len(pending.buffer) {
					n = len(pending.buffer)
				}
				data = pending.buffer[:n]
			}
			resp, err = BuildRetSubmit(
				pending.seqNum,
				pending.devID,
				pending.direction,
				pending.endpoint,
				urbInfo.Status,
				uint32(urbInfo.ActualLength),
				data,
			)
		}
		if err != nil {
			continue
		}

		select {
		case retChan <- resp:
		case <-s.stop:
			return
		}
	}
}

// buildISOResponse builds a RET_SUBMIT for a completed ISO URB
func (s *Server) buildISOResponse(urbInfo *usb.ReapedURBInfo, pending *pendingURB) ([]byte, error) {
	// Read ISO packet results from the URB memory
	isoResults := usb.ReadISOResults(pending.isoMem, pending.numPackets)

	// Build USB/IP ISO descriptors and pack transfer data
	var usbipDescs []ISOPacketDescriptor
	var packedData []byte
	bufOffset := uint32(0) // offset in our sequential buffer

	for i := int32(0); i < pending.numPackets; i++ {
		pktLen := pending.packetLens[i]
		actualLen := isoResults[i].ActualLength
		status := isoResults[i].Status

		usbipDescs = append(usbipDescs, ISOPacketDescriptor{
			Offset:       bufOffset,
			Length:       pktLen,
			ActualLength: actualLen,
			Status:       status,
		})

		// Pack actual data (compact, no gaps) for IN direction
		if pending.direction == DirIn && actualLen > 0 {
			end := bufOffset + actualLen
			if end > uint32(len(pending.buffer)) {
				end = uint32(len(pending.buffer))
			}
			if bufOffset < end {
				packedData = append(packedData, pending.buffer[bufOffset:end]...)
			}
		}

		bufOffset += pktLen
	}

	return BuildRetSubmitISO(
		pending.seqNum,
		pending.devID,
		pending.direction,
		pending.endpoint,
		urbInfo.Status,
		uint32(urbInfo.ActualLength),
		packedData,
		uint32(urbInfo.StartFrame),
		pending.numPackets,
		urbInfo.ErrorCount,
		usbipDescs,
	)
}

// HandleDevlistRequest handles an OP_REQ_DEVLIST for this device
func (s *Server) HandleDevlistRequest() ([]byte, error) {
	desc := s.BuildDeviceDescriptor()
	ifaceDescs := s.BuildInterfaceDescriptors()
	return BuildDevlistReply([]DeviceDescriptor{desc}, [][]InterfaceDescriptor{ifaceDescs})
}

// HandleImportRequest handles an OP_REQ_IMPORT for this device
func (s *Server) HandleImportRequest(requestedBusID string) ([]byte, error) {
	if requestedBusID != s.device.BusID {
		return BuildImportReply(1, nil) // device not found
	}
	desc := s.BuildDeviceDescriptor()
	return BuildImportReply(0, &desc)
}

// ReadManagementRequest reads and dispatches a management phase message.
// Returns the response bytes and whether we should transition to transfer phase.
func (s *Server) ReadManagementRequest(r io.Reader) (response []byte, startTransfer bool, err error) {
	hdr, err := ReadOpHeader(r)
	if err != nil {
		return nil, false, err
	}

	switch hdr.Command {
	case OpReqDevlist:
		resp, err := s.HandleDevlistRequest()
		return resp, false, err

	case OpReqImport:
		var busID [32]byte
		if _, err := io.ReadFull(r, busID[:]); err != nil {
			return nil, false, err
		}
		reqBusID := GetBusID(busID)
		resp, err := s.HandleImportRequest(reqBusID)
		if err != nil {
			return nil, false, err
		}

		// Check if import was successful (status in response)
		var checkBuf bytes.Buffer
		checkBuf.Write(resp)
		checkHdr, _ := ReadOpHeader(&checkBuf)
		if checkHdr != nil && checkHdr.Status == 0 {
			return resp, true, nil // successful import -> transfer phase
		}
		return resp, false, nil

	default:
		return nil, false, fmt.Errorf("unknown management command: 0x%04x", hdr.Command)
	}
}
