using System;
using System.Collections.Generic;
using System.Collections.ObjectModel;
using System.Linq;
using System.Text;
using System.Threading;
using LibUsbDotNet;
using LibUsbDotNet.Main;

namespace USBMassStorage
{
  public class USBDevice
  {
    private UsbRegistry _reg;
    private int _vendorId;
    private int _productId;
    private byte _configurationId;
    private byte _interfaceId;
    private LibUsbDotNet.LibUsb.LibUsbDevice _device = null;
    private LibUsbDotNet.Info.UsbInterfaceInfo _interface;
    private UsbEndpointReader _dataReader = null;
    private UsbEndpointWriter _dataWriter = null;
    private UsbEndpointReader _statusReader = null;
    private Queue<byte> _dataReadBuffer = new Queue<byte>();
    private Queue<byte> _statusBuffer = new Queue<byte>();

    public USBDevice(UsbRegistry registry, byte configurationId, byte interfaceId)
    {
      _reg = registry;
      _configurationId = configurationId;
      _interfaceId = interfaceId;
      _vendorId = _reg.Vid;
      _productId = _reg.Pid;
    }

    public UsbRegistry Registry
    {
      get
      {
        return _reg;
      }
    }

    public byte ConfigurationId
    {
      get
      {
        return _configurationId;
      }
    }

    public byte InterfaceId
    {
      get
      {
        return _interfaceId;
      }
    }

    public int VendorId
    {
      get
      {
        return _vendorId;
      }
    }

    public int ProductId
    {
      get
      {
        return _productId;
      }
    }

    public bool IsConnected
    {
      get
      {
        return (_device != null && _device.IsOpen);
      }
    }

    public bool UsesStatusInterrupt
    {
      get
      {
        return (_interface != null) && (_interface.Descriptor.Protocol <= 0x01) && (_statusReader != null);
      }
    }

    public static List<USBDevice> GetAllCompatibleDevices()
    {
      UsbDevice device;
      var ret = new List<USBDevice>();

      foreach (UsbRegistry reg in UsbDevice.AllDevices)
      {
        if (reg.Open(out device))
        {
          foreach (LibUsbDotNet.Info.UsbConfigInfo config in device.Configs)
          {
            foreach (LibUsbDotNet.Info.UsbInterfaceInfo inter in config.InterfaceInfoList)
            {
              if (inter.Descriptor.Class == LibUsbDotNet.Descriptors.ClassCodeType.MassStorage)
                ret.Add(new USBDevice(reg, config.Descriptor.ConfigID, inter.Descriptor.InterfaceID));
            }
          }

          device.Close();
        }
      }

      return ret;
    }

    public void Connect()
    {
      if (_device == null)
      {
        //Open the device
        UsbDevice device;
        _reg.Open(out device);
        _device = device as LibUsbDotNet.LibUsb.LibUsbDevice;

        //Set configuration
        _device.SetConfiguration(_configurationId);

        //Claim the interface
        _device.ClaimInterface(_interfaceId);
        var success = _device.SetAltInterface(_interfaceId, 0);

        //Set up the incoming and outgoing pipes
        foreach (LibUsbDotNet.Info.UsbConfigInfo config in _device.Configs)
        {
          if (config.Descriptor.ConfigID == _configurationId)
          {
            foreach (LibUsbDotNet.Info.UsbInterfaceInfo intrface in config.InterfaceInfoList)
            {
              if (intrface.Descriptor.InterfaceID == _interfaceId)
              {
                _interface = intrface;

                foreach (LibUsbDotNet.Info.UsbEndpointInfo endpoint in intrface.EndpointInfoList)
                {
                  if ((endpoint.Descriptor.EndpointID & 0x80) > 0)
                  {
                    //Incoming endpoint
                    if ((endpoint.Descriptor.Attributes & 0x03) == 0x03)
                    {
                      //This is an interrupt endpoint
                      if (_statusReader == null)
                      {
                        _statusReader = _device.OpenEndpointReader((ReadEndpointID)endpoint.Descriptor.EndpointID,
                          endpoint.Descriptor.MaxPacketSize, EndpointType.Interrupt);
                        _statusReader.DataReceived += new EventHandler<EndpointDataEventArgs>(_statusReader_DataReceived);
                        _statusReader.DataReceivedEnabled = true;
                      }
                    }
                    else if ((endpoint.Descriptor.Attributes & 0x03) == 0x02)
                    {
                      //This is a bulk endpoint
                      if (_dataReader == null)
                      {
                        _dataReader = _device.OpenEndpointReader((ReadEndpointID)endpoint.Descriptor.EndpointID,
                          endpoint.Descriptor.MaxPacketSize, EndpointType.Bulk);
                        _dataReader.DataReceived += _dataReader_DataReceived;
                        _dataReader.DataReceivedEnabled = true;
                      }
                    }
                    else
                    {
                      //Uh...nevermind
                    }
                  }
                  else
                  {
                    //Outgoing endpoint
                    if ((endpoint.Descriptor.Attributes & 0x03) == 0x02)
                    {
                      //This is a bulk endpoint
                      if (_dataWriter == null)
                        _dataWriter = _device.OpenEndpointWriter((WriteEndpointID)endpoint.Descriptor.EndpointID,
                          EndpointType.Bulk);
                    }
                  }
                }

                break;
              }
            }
          }
        }

        if (_dataReader == null || _dataWriter == null)
          throw new InvalidOperationException("Endpoints have not been set up for this device.");
      }
    }

    public void Disconnect()
    {
      if (_device != null)
      {
        if (_dataReader != null) _dataReader.Abort();
        if (_dataWriter != null) _dataWriter.Abort();
        if (_statusReader != null) _statusReader.Abort();
        _device.ReleaseAllInterfaces();

        if (_device.IsOpen)
          _device.Close();

        _device = null;
      }
    }

    public void Dispose()
    {
      try
      {
        Disconnect();
      }
      catch
      {
        //Don't care...
      }
    }

    public byte GetMaxLUN()
    {
      var setup = new UsbSetupPacket(0xA1, 0xFE, 0, 0, 1);
      var ret = new byte[1];
      int transferred;

      _device.ControlTransfer(ref setup, ret, 128, out transferred);

      return ret[0];
    }

    public byte[] GetStringDescriptor(byte index)
    {
      var ret = new byte[128];
      int transferred;

      _device.GetDescriptor(0x03, index, 0x0409, ret, 128, out transferred);

      return ret;
    }

    public byte[] SendCommand(byte[] request)
    {
      return _SendCommand(request);
    }

    public byte[] SendCommand(byte[] request, byte[] outgoingData)
    {
      return _SendCommand(request, outgoingData);
    }

    public byte[] SendCommand(byte[] request, ref byte[] incomingData)
    {
      return _SendCommand(request, ref incomingData);
    }

    public byte[] SendRequestSenseCommand()
    {
      var buffer = new byte[] {0x03, 0x00, 0x00, 0x00, 0x12, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00};
      var ret = new byte[18];

      _SendCommand(buffer, ref ret);

      return ret;
    }

    public byte[] SendReadCapacityCommand()
    {
      var buffer = new byte[] {0x25, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00};
      var ret = new byte[8];

      _SendCommand(buffer, ref ret);

      return ret;
    }

    public byte[] SendReadSectorCommand(int LBA, int length)
    {
      //TODO: Technically shouldn't hard-code 512 here
      var buffer = new byte[] {0x28, 0x00, (byte)((LBA >> 24) & 0xFF), (byte)((LBA >> 16) & 0xFF),
        (byte)((LBA >> 8) & 0xFF), (byte)(LBA & 0xFF), 0x00, (byte)((length / 512 >> 8) & 0xFF), (byte)(length / 512 & 0xFF),
        0x00, 0x00, 0x00};
      var ret = new byte[length];

      _SendCommand(buffer, ref ret);

      return ret;
    }

    public void SendWriteSectorCommand(int LBA, byte[] data, int index, int length)
    {
      //TODO: Technically shouldn't hard-code 512 here
      var buffer = new byte[] {0x2A, 0x00, (byte)((LBA >> 24) & 0xFF), (byte)((LBA >> 16) & 0xFF),
        (byte)((LBA >> 8) & 0xFF), (byte)(LBA & 0xFF), 0x00, (byte)((length / 512 >> 8) & 0xFF), (byte)(length / 512 & 0xFF),
        0x00, 0x00, 0x00};
      var dataToWrite = new byte[length];

      Array.Copy(data, index, dataToWrite, 0, length);
      _SendCommand(buffer, dataToWrite);
    }

    private byte[] _SendCommand(byte[] request)
    {
      byte[] outgoingData = null;

      return _SendCommand(request, outgoingData);
    }

    private byte[] _SendCommand(byte[] request, ref byte[] incomingData)
    {
      byte[] outgoingData = null;

      return _SendCommand(request, outgoingData, ref incomingData);
    }

    private byte[] _SendCommand(byte[] request, byte[] outgoingData)
    {
      byte[] incomingData = null;

      return _SendCommand(request, outgoingData, ref incomingData);
    }

    //Sends command block with optional incoming/outgoing data buffers, and returns status data buffer.
    private byte[] _SendCommand(byte[] request, byte[] outgoingData, ref byte[] incomingData)
    {
      const int REQUEST_TIMEOUT = 5000;
      int transferred;
      byte[] ret = null;

      if (this.UsesStatusInterrupt)
      {
        lock (_statusBuffer)
        {
          //Clear the buffer
          _statusBuffer.Clear();

          //Send the control request
          var setup = new UsbSetupPacket(0x21, 0x00, 0x0000, 0x0000, (short)request.Length);
          _device.ControlTransfer(ref setup, request, request.Length, out transferred);

          //Send/wait on any data
          if (outgoingData != null && outgoingData.Length > 0)
            _dataWriter.Write(outgoingData, REQUEST_TIMEOUT, out transferred);
          if (incomingData != null && incomingData.Length > 0)
            _WaitOnData(_dataReadBuffer, ref incomingData, REQUEST_TIMEOUT);
        }

        //Wait on status
        ret = new byte[2];
        _WaitOnData(_statusBuffer, ref ret, REQUEST_TIMEOUT);
      }
      else
      {
        lock (_dataReadBuffer)
        {
          //Clear the buffer
          _dataReadBuffer.Clear();

          //Send the request
          var cbw = new byte[31];
          cbw[0] = 0x55; cbw[1] = 0x53; cbw[2] = 0x42; cbw[3] = 0x43;
          cbw[4] = 0x41; cbw[5] = 0x42; cbw[6] = 0x43; cbw[7] = 0x44;
          int dataLength = 0;
          if (incomingData != null && incomingData.Length > 0) dataLength = incomingData.Length;
          if (outgoingData != null && outgoingData.Length > 0) dataLength = outgoingData.Length;
          cbw[8] = (byte)(dataLength & 0xFF); cbw[9] = (byte)((dataLength >> 8) & 0xFF);
          cbw[10] = (byte)((dataLength >> 16) & 0xFF); cbw[11] = (byte)((dataLength >> 24) & 0xFF);
          cbw[12] = (byte)((incomingData != null && incomingData.Length > 0) ? 0x80 : 0);
          cbw[14] = (byte)request.Length;
          Array.Copy(request, 0, cbw, 15, request.Length);
          _dataWriter.Write(cbw, REQUEST_TIMEOUT, out transferred);

          //Send any data
          if (outgoingData != null && outgoingData.Length > 0)
            _dataWriter.Write(outgoingData, REQUEST_TIMEOUT, out transferred);
        }

        //Wait on any data
        if (incomingData != null && incomingData.Length > 0)
          _WaitOnData(_dataReadBuffer, ref incomingData, REQUEST_TIMEOUT);

        //Wait on status
        ret = new byte[13];
        _WaitOnData(_dataReadBuffer, ref ret, REQUEST_TIMEOUT);
      }

      return ret;
    }

    private void _WaitOnData(Queue<byte> buffer, ref byte[] incomingData, int requestTimeout)
    {
      const int SLEEP_TIME = 10;
      int count = 0;
      var start = DateTime.Now;

      do
      {
        lock (buffer)
        {
          count = buffer.Count;
        }

        Thread.Sleep(SLEEP_TIME);

        if (start.AddMilliseconds(requestTimeout) < DateTime.Now)
          break; //Uh...freak out here
      } while (count < incomingData.Length);

      //Return the two status bytes
      lock (buffer)
      {
        for (int i = 0; i < incomingData.Length && buffer.Count > 0; i++)
          incomingData[i] = buffer.Dequeue();
      }
    }

    private void _statusReader_DataReceived(object sender, EndpointDataEventArgs e)
    {
      string msg = "Status: ";
      for (int i = 0; i < e.Count; i++)
        msg += e.Buffer[i].ToString("X2") + " ";
      Logger.WriteLine(msg, Logger.Type.Information);

      lock (_statusBuffer)
      {
        for (int i = 0; i < e.Count; i++)
          _statusBuffer.Enqueue(e.Buffer[i]);
      }
    }

    private void _dataReader_DataReceived(object sender, EndpointDataEventArgs e)
    {
      string msg = "Incoming data: ";
      for (int i = 0; i < e.Count; i++)
        msg += e.Buffer[i].ToString("X2") + " ";
      Logger.WriteLine(msg, Logger.Type.Information);

      lock (_dataReadBuffer)
      {
        for (int i = 0; i < e.Count; i++)
          _dataReadBuffer.Enqueue(e.Buffer[i]);
      }
    }
  }
}