//go:build windows

package usbip

import (
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"log"
	"sync"

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

// USB/IP server for Windows, driving devices through the usbshare filter
// driver (driver/windows).
//
// The shape differs from the Linux server because the interfaces differ:
// usbdevfs submits asynchronously and hands completions back through a reap
// loop, whereas the filter driver's IOCTL blocks until the transfer finishes.
// Concurrency therefore comes from a pool of workers rather than from one
// reaper.
//
// UNTESTED: this depends on the filter driver, which has never been built or
// run. Treat it as a starting point, not as working code.

// transferWorkers bounds how many transfers are in flight at once. USB/IP
// clients keep several outstanding, and serialising them would stall the
// device on every round trip.
const transferWorkers = 8

// Server handles USB/IP protocol on the share side.
type Server struct {
	device *usb.Device
	handle *usb.DriverHandle

	mu      sync.Mutex
	closed  bool
	pending map[uint32]uint64 // USB/IP seqnum -> driver transfer ID

	epTypes map[uint8]uint8

	work     chan *transferJob
	ctrlWork chan *transferJob
	stop     chan struct{}
	workers  sync.WaitGroup
	stopOnce sync.Once
}

// transferJob is one queued USB/IP request.
type transferJob struct {
	hdr         *URBHeader
	body        *CmdSubmitBody
	transferBuf []byte
	retChan     chan<- []byte
}

// NewServer creates a USB/IP server for a specific device.
func NewServer(dev *usb.Device) *Server {
	return &Server{
		device:   dev,
		pending:  make(map[uint32]uint64),
		epTypes:  make(map[uint8]uint8),
		work:     make(chan *transferJob, 64),
		ctrlWork: make(chan *transferJob, 64),
		stop:     make(chan struct{}),
	}
}

// Attach opens the device through the filter driver and claims it.
func (s *Server) Attach() error {
	if s.device.DevPath == "" {
		return fmt.Errorf("device %s has no driver path; is the usbshare filter attached?", s.device.BusID)
	}

	handle, err := usb.OpenDriverDevice(s.device.DevPath)
	if err != nil {
		return fmt.Errorf("claiming %s: %w", s.device.BusID, err)
	}
	s.handle = handle

	s.buildEndpointTypeMap()
	return nil
}

// Detach releases the device back to its class driver.
func (s *Server) Detach() {
	s.mu.Lock()
	if s.closed {
		s.mu.Unlock()
		return
	}
	s.closed = true
	s.mu.Unlock()

	s.stopOnce.Do(func() { close(s.stop) })
	s.workers.Wait()

	if s.handle != nil {
		s.handle.Close()
		s.handle = nil
	}
}

// buildEndpointTypeMap indexes endpoints by full address, including the
// direction bit — a composite device can use the same endpoint number for an
// interrupt IN and a bulk OUT, and conflating them breaks HID devices.
func (s *Server) buildEndpointTypeMap() {
	for _, ep := range s.device.Endpoints {
		s.epTypes[ep.Address] = ep.TransferType
	}
	if len(s.epTypes) > 0 {
		return
	}
	for _, iface := range s.device.Interfaces {
		for _, ep := range iface.Endpoints {
			s.epTypes[ep.Address] = ep.TransferType
		}
	}
}

// getTransferType maps an endpoint to a driver transfer type, falling back to
// the request's interval: only periodic transfers carry one, so an unknown
// endpoint with a non-zero interval is interrupt rather than bulk.
func (s *Server) getTransferType(epAddr uint8, interval uint32) uint8 {
	if epAddr&0x0F == 0 {
		return usb.TransferTypeControl
	}
	if t, ok := s.epTypes[epAddr]; ok {
		return t
	}
	if interval > 0 {
		return usb.TransferTypeInterrupt
	}
	return usb.TransferTypeBulk
}

// BuildDeviceDescriptor creates a USB/IP device descriptor.
func (s *Server) BuildDeviceDescriptor() DeviceDescriptor {
	var desc DeviceDescriptor
	SetPath(&desc.Path, s.device.DevPath)
	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.
func (s *Server) HandleConnection(r io.Reader, w io.Writer) error {
	retChan := make(chan []byte, 256)

	// Control transfers get their own serial worker because endpoint 0 is a
	// single shared pipe; everything else runs on a pool.
	s.workers.Add(1)
	go s.controlWorker(retChan)

	for i := 0; i < transferWorkers; i++ {
		s.workers.Add(1)
		go s.transferWorker(retChan)
	}

	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
			}
		}
	}()

	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)
		}
	}
}

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

	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)
		}
	}

	// Isochronous transfers are not supported by the driver yet; the packet
	// descriptors still have to be consumed or the stream desynchronises.
	if body.NumberOfPackets != 0xFFFFFFFF && body.NumberOfPackets > 0 {
		descs := make([]ISOPacketDescriptor, body.NumberOfPackets)
		binary.Read(r, binary.BigEndian, &descs)

		log.Printf("[usbip-win] isochronous transfer on EP%d rejected (not implemented)", hdr.Endpoint)
		resp, _ := BuildRetSubmit(hdr.SeqNum, hdr.DevID, hdr.Direction, hdr.Endpoint, -32, 0, nil)
		s.send(retChan, resp)
		return nil
	}

	job := &transferJob{hdr: hdr, body: body, transferBuf: transferBuf, retChan: retChan}

	queue := s.work
	if hdr.Endpoint == 0 {
		queue = s.ctrlWork
	}

	select {
	case queue <- job:
	case <-s.stop:
	}
	return nil
}

func (s *Server) controlWorker(retChan chan<- []byte) {
	defer s.workers.Done()

	for {
		select {
		case <-s.stop:
			return
		case job := <-s.ctrlWork:
			s.runTransfer(job, retChan)
		}
	}
}

func (s *Server) transferWorker(retChan chan<- []byte) {
	defer s.workers.Done()

	for {
		select {
		case <-s.stop:
			return
		case job := <-s.work:
			s.runTransfer(job, retChan)
		}
	}
}

// runTransfer performs one transfer and emits its RET_SUBMIT.
func (s *Server) runTransfer(job *transferJob, retChan chan<- []byte) {
	hdr, body := job.hdr, job.body

	s.mu.Lock()
	handle := s.handle
	closed := s.closed
	s.mu.Unlock()

	if closed || handle == nil {
		return
	}

	epAddr := uint8(hdr.Endpoint)
	if hdr.Direction == DirIn {
		epAddr |= 0x80
	}

	direction := uint8(0) // USBSHARE_DIR_OUT
	if hdr.Direction == DirIn {
		direction = 1
	}

	// Intercept the standard requests that need dedicated driver calls:
	// sending them as raw control transfers changes the device without
	// telling the USB stack, after which later transfers fail.
	if hdr.Endpoint == 0 && hdr.Direction == DirOut {
		bmRequestType := body.Setup[0]
		bRequest := body.Setup[1]
		wValue := binary.LittleEndian.Uint16(body.Setup[2:4])
		wIndex := binary.LittleEndian.Uint16(body.Setup[4:6])

		switch {
		case bmRequestType == 0x01 && bRequest == 0x0B: // SET_INTERFACE
			var status int32
			if err := handle.SetInterface(uint8(wIndex), uint8(wValue)); err != nil {
				log.Printf("[usbip-win] SET_INTERFACE(%d, %d) failed: %v", wIndex, wValue, err)
				status = -32
			}
			resp, _ := BuildRetSubmit(hdr.SeqNum, hdr.DevID, hdr.Direction, hdr.Endpoint, status, 0, nil)
			s.send(retChan, resp)
			return

		case bmRequestType == 0x02 && bRequest == 0x01 && wValue == 0x0000: // CLEAR_FEATURE(HALT)
			var status int32
			if err := handle.ClearHalt(uint8(wIndex)); err != nil {
				log.Printf("[usbip-win] CLEAR_HALT(0x%02x) failed: %v", wIndex, err)
				status = -32
			}
			resp, _ := BuildRetSubmit(hdr.SeqNum, hdr.DevID, hdr.Direction, hdr.Endpoint, status, 0, nil)
			s.send(retChan, resp)
			return

		case bmRequestType == 0x00 && bRequest == 0x09: // SET_CONFIGURATION
			// The device is already configured; forwarding this would reset
			// its endpoint state behind the stack's back.
			log.Printf("[usbip-win] SET_CONFIGURATION(%d) intercepted", wValue)
			resp, _ := BuildRetSubmit(hdr.SeqNum, hdr.DevID, hdr.Direction, hdr.Endpoint, 0, 0, nil)
			s.send(retChan, resp)
			return
		}
	}

	var buf []byte
	if hdr.Direction == DirIn {
		buf = make([]byte, body.TransferBufferLen)
	} else {
		buf = job.transferBuf
		if buf == nil {
			buf = make([]byte, 0)
		}
	}

	params := &usb.TransferParams{
		EndpointAddress: epAddr,
		Type:            s.getTransferType(epAddr, body.Interval),
		Direction:       direction,
		Data:            buf,
		TimeoutMS:       5000,
		Setup:           body.Setup,
	}

	n, err := handle.Transfer(params)

	var status int32
	if err != nil {
		log.Printf("[usbip-win] transfer on EP%d failed: %v", hdr.Endpoint, err)
		status = -32 // -EPIPE
	}

	// actualLength must be reported for both directions: the kernel UVC
	// driver checks it on OUT control transfers.
	var data []byte
	if hdr.Direction == DirIn && n > 0 {
		if n > len(buf) {
			n = len(buf)
		}
		data = buf[:n]
	}

	resp, buildErr := BuildRetSubmit(hdr.SeqNum, hdr.DevID, hdr.Direction, hdr.Endpoint,
		status, uint32(n), data)
	if buildErr != nil {
		return
	}
	s.send(retChan, resp)
}

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

	// Cancellation is best effort here: the driver tracks transfers by its
	// own ID, and a transfer already completing cannot be recalled.
	resp, err := BuildRetUnlink(hdr.SeqNum, hdr.DevID, 0)
	if err != nil {
		return err
	}
	s.send(retChan, resp)

	log.Printf("[usbip-win] UNLINK for seq=%d acknowledged", body.UnlinkSeqNum)
	return nil
}

func (s *Server) send(retChan chan<- []byte, resp []byte) {
	select {
	case retChan <- resp:
	case <-s.stop:
	}
}

// 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)
	}
	desc := s.BuildDeviceDescriptor()
	return BuildImportReply(0, &desc)
}

// ReadManagementRequest reads and dispatches a management phase message.
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
		}
		resp, err := s.HandleImportRequest(GetBusID(busID))
		if err != nil {
			return nil, false, err
		}
		return resp, len(resp) > 8, nil

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