//go:build linux

package usbip

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"io"
	"log"
	"sync"
	"unsafe"

	"github.com/duffy/usb-server/internal/usb"
	"golang.org/x/sys/unix"
)

// 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
	closed     bool
}

type pendingURB struct {
	seqNum    uint32
	devID     uint32
	direction uint32
	endpoint  uint32
	buffer    []byte
	urbPtr    unsafe.Pointer // pointer to submitted usbdevfs_urb
}

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

// 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

	// Disconnect kernel drivers from all interfaces
	for _, iface := range s.device.Interfaces {
		if iface.Driver != "" && iface.Driver != "(none)" {
			// Try to disconnect - ignore errors for already-disconnected interfaces
			handle.DisconnectDriver()
		}
	}

	// Claim all interfaces
	for _, iface := range s.device.Interfaces {
		if err := handle.ClaimInterface(uint32(iface.Number)); err != nil {
			log.Printf("[usbip-server] warning: could not claim interface %d: %v", iface.Number, err)
		}
	}

	return nil
}

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

	if s.handle == nil {
		return
	}

	// Release all interfaces
	for _, iface := range s.device.Interfaces {
		s.handle.ReleaseInterface(uint32(iface.Number))
	}

	// Reconnect kernel driver
	s.handle.ConnectDriver()

	s.handle.Close()
	s.handle = nil
}

// 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 {
	// Start the URB reaper goroutine
	retChan := make(chan []byte, 64)
	done := make(chan struct{})
	defer close(done)

	go s.reapLoop(retChan, done)

	// Forward completed URBs to the writer
	go func() {
		for {
			select {
			case data, ok := <-retChan:
				if !ok {
					return
				}
				if _, err := w.Write(data); err != nil {
					return
				}
			case <-done:
				return
			}
		}
	}()

	// Read and process incoming USB/IP messages
	for {
		// Read the URB header (20 bytes basic + 28 bytes specific = 48 total)
		hdr, err := ReadURBHeader(r)
		if err != nil {
			if err == io.EOF {
				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
	}

	// 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
	if body.NumberOfPackets != 0xFFFFFFFF && body.NumberOfPackets > 0 {
		isoDescs := make([]ISOPacketDescriptor, body.NumberOfPackets)
		if err := binary.Read(r, binary.BigEndian, &isoDescs); err != nil {
			return fmt.Errorf("reading ISO descriptors: %w", err)
		}
		// TODO: handle ISO transfers properly
	}

	// Determine URB type from endpoint
	endpoint := uint8(hdr.Endpoint)
	var urbType uint8
	if endpoint == 0 {
		urbType = 2 // control
	} else {
		urbType = 3 // bulk (most common, we'll detect interrupt from endpoint descriptor later)
	}

	// Handle control transfers specially (endpoint 0)
	if endpoint == 0 && hdr.Direction == DirIn {
		// Control IN: send setup packet, receive data
		buf := make([]byte, body.TransferBufferLen)
		n, err := s.handle.ControlTransfer(
			body.Setup[0], body.Setup[1],
			binary.LittleEndian.Uint16(body.Setup[2:4]),
			binary.LittleEndian.Uint16(body.Setup[4:6]),
			binary.LittleEndian.Uint16(body.Setup[6:8]),
			5000, buf,
		)
		var status int32
		if err != nil {
			status = -32 // -EPIPE
			n = 0
		}
		resp, err := BuildRetSubmit(hdr.SeqNum, hdr.DevID, hdr.Direction, hdr.Endpoint, status, buf[:n])
		if err != nil {
			return err
		}
		retChan <- resp
		return nil
	}

	if endpoint == 0 && hdr.Direction == DirOut {
		// Control OUT
		buf := transferBuf
		if buf == nil {
			buf = make([]byte, 0)
		}
		_, err := s.handle.ControlTransfer(
			body.Setup[0], body.Setup[1],
			binary.LittleEndian.Uint16(body.Setup[2:4]),
			binary.LittleEndian.Uint16(body.Setup[4:6]),
			binary.LittleEndian.Uint16(body.Setup[6:8]),
			5000, buf,
		)
		var status int32
		if err != nil {
			status = -32 // -EPIPE
		}
		resp, err := BuildRetSubmit(hdr.SeqNum, hdr.DevID, hdr.Direction, hdr.Endpoint, status, nil)
		if err != nil {
			return err
		}
		retChan <- resp
		return nil
	}

	// For non-control transfers, submit asynchronously
	var buf []byte
	if hdr.Direction == DirIn {
		buf = make([]byte, body.TransferBufferLen)
	} else {
		buf = transferBuf
	}

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

	urb, err := s.handle.SubmitURB(&usb.SubmitURBParams{
		Type:       urbType,
		Endpoint:   ep,
		Flags:      0,
		Buffer:     buf,
		UserContext: uintptr(hdr.SeqNum),
	})
	if err != nil {
		// Submit failed - send error response immediately
		resp, _ := BuildRetSubmit(hdr.SeqNum, hdr.DevID, hdr.Direction, hdr.Endpoint, -32, nil)
		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
}

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)
	}
	s.mu.Unlock()

	var status int32
	if exists && pending.urbPtr != nil {
		// Try to discard the URB
		// Note: we cast back to the URB type for the ioctl
		urbForDiscard := (*usbDevfsURBForDiscard)(pending.urbPtr)
		_, _, errno := unix.Syscall(unix.SYS_IOCTL, uintptr(s.handle.Fd()),
			uintptr(0x8000550B), // USBDEVFS_DISCARDURB
			uintptr(pending.urbPtr))
		_ = urbForDiscard
		if errno == 0 {
			status = -104 // -ECONNRESET
		}
	}

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

	return nil
}

// usbDevfsURBForDiscard is a placeholder to make the Go compiler happy
type usbDevfsURBForDiscard struct{}

// reapLoop continuously reaps completed URBs and sends responses
func (s *Server) reapLoop(retChan chan<- []byte, done <-chan struct{}) {
	for {
		select {
		case <-done:
			return
		default:
		}

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

		urb, err := s.handle.ReapURB()
		if err != nil {
			// Check if we should stop
			select {
			case <-done:
				return
			default:
				continue
			}
		}

		seqNum := uint32(urb.UserContext)

		s.mu.Lock()
		pending, exists := s.pendingURBs[seqNum]
		if exists {
			delete(s.pendingURBs, seqNum)
		}
		s.mu.Unlock()

		if !exists {
			continue
		}

		var data []byte
		if pending.direction == DirIn && urb.ActualLength > 0 {
			data = pending.buffer[:urb.ActualLength]
		}

		resp, err := BuildRetSubmit(
			pending.seqNum,
			pending.devID,
			pending.direction,
			pending.endpoint,
			urb.Status,
			data,
		)
		if err != nil {
			continue
		}

		select {
		case retChan <- resp:
		case <-done:
			return
		}
	}
}

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