/*
  
  Version: MPL 1.1/GPL 2.0/LGPL 2.1

  The contents of this file are subject to the Mozilla Public License Version
  1.1 (the "License"); you may not use this file except in compliance with
  the License. You may obtain a copy of the License at
 
  http://www.mozilla.org/MPL/

  Software distributed under the License is distributed on an "AS IS" basis,
  WITHOUT WARRANTY OF ANY KIND, either express or implied. See the License
  for the specific language governing rights and limitations under the License.

  The Original Code is the Open Hardware Monitor code.

  The Initial Developer of the Original Code is 
  Michael Möller <m.moeller@gmx.ch>.
  Portions created by the Initial Developer are Copyright (C) 2009-2010
  the Initial Developer. All Rights Reserved.

  Contributor(s): Paul Werelds

  Alternatively, the contents of this file may be used under the terms of
  either the GNU General Public License Version 2 or later (the "GPL"), or
  the GNU Lesser General Public License Version 2.1 or later (the "LGPL"),
  in which case the provisions of the GPL or the LGPL are applicable instead
  of those above. If you wish to allow use of your version of this file only
  under the terms of either the GPL or the LGPL, and not to allow others to
  use your version of this file under the terms of the MPL, indicate your
  decision by deleting the provisions above and replace them with the notice
  and other provisions required by the GPL or the LGPL. If you do not delete
  the provisions above, a recipient may use your version of this file under
  the terms of any one of the MPL, the GPL or the LGPL.
 
*/

using System;
using System.Collections.Generic;
using System.Runtime.InteropServices;

namespace OpenHardwareMonitor.Hardware.HDD {

  internal class SMART {

    [Flags]
    public enum Status : ushort {
      PreFailureWarranty = 0x01,
      OnLineCollection = 0x02,
      Performance = 0x04,
      ErrorRate = 0x08,
      EventCount = 0x10,
      SelfPreserving = 0x20
    }

    [StructLayout(LayoutKind.Sequential, Pack = 1)]
    public struct AttributeID {
      private byte value;

      public AttributeID(byte value) {
        this.value = value;
      }

      public override bool Equals(Object obj) {
        return obj is AttributeID && this == (AttributeID)obj;
      }
      public override int GetHashCode() {
        return value.GetHashCode() ^ value.GetHashCode();
      }
      public static bool operator ==(AttributeID a, AttributeID b) {
        return a.value == b.value;
      }
      public static bool operator !=(AttributeID a, AttributeID b) {
        return !(a == b);
      }

      public string ToString(string format) {
        return value.ToString(format);
      }

      public static readonly AttributeID None = new AttributeID(0x00);
    }

    // These are the more-or-less standard S.M.A.R.T attributes
    // TODO: Filter out unused/obscure ones; some are interpreted differently
    // between manufacturers
    public static class CommonAttributes {
      public static readonly AttributeID
        ReadErrorRate = new AttributeID(0x01),
        ThroughputPerformance = new AttributeID(0x02),
        SpinUpTime = new AttributeID(0x03),
        StartStopCount = new AttributeID(0x04),
        ReallocatedSectorsCount = new AttributeID(0x05),
        ReadChannelMargin = new AttributeID(0x06),
        SeekErrorRate = new AttributeID(0x07),
        SeekTimePerformance = new AttributeID(0x08),
        PowerOnHours = new AttributeID(0x09),
        SpinRetryCount = new AttributeID(0x0A),
        RecalibrationRetries = new AttributeID(0x0B),
        PowerCycleCount = new AttributeID(0x0C),
        SoftReadErrorRate = new AttributeID(0x0D),
        SataDownshiftErrorCount = new AttributeID(0xB7),
        EndToEndError = new AttributeID(0xB8),
        HeadStability = new AttributeID(0xB9),
        InducedOpVibrationDetection = new AttributeID(0xBA),
        ReportedUncorrectableErrors = new AttributeID(0xBB),
        CommandTimeout = new AttributeID(0xBC),
        HighFlyWrites = new AttributeID(0xBD),
        AirflowTemperature = new AttributeID(0xBE),
        GSenseErrorRate = new AttributeID(0xBF),
        PowerOffRetractCount = new AttributeID(0xC0),
        LoadCycleCount = new AttributeID(0xC1),
        Temperature = new AttributeID(0xC2),
        HardwareEccRecovered = new AttributeID(0xC3),
        ReallocationEventCount = new AttributeID(0xC4),
        CurrentPendingSectorCount = new AttributeID(0xC5),
        UncorrectableSectorCount = new AttributeID(0xC6),
        UltraDmaCrcErrorCount = new AttributeID(0xC7),
        WriteErrorRate = new AttributeID(0xC8),
        DataAddressMarkerrors = new AttributeID(0xCA),
        RunOutCancel = new AttributeID(0xCB),
        SoftEccCorrection = new AttributeID(0xCC),
        ThermalAsperityRate = new AttributeID(0xCD),
        FlyingHeight = new AttributeID(0xCE),
        SpinHighCurrent = new AttributeID(0xCF),
        SpinBuzz = new AttributeID(0xD0),
        OfflineSeekPerformance = new AttributeID(0xD1),
        VibrationDuringWrite = new AttributeID(0xD3),
        ShockDuringWrite = new AttributeID(0xD4),
        DiskShift = new AttributeID(0xDC),
        GSenseErrorRateAlt = new AttributeID(0xDD), // Alternative to 0xBF
        LoadedHours = new AttributeID(0xDE),
        LoadUnloadRetryCount = new AttributeID(0xDF),
        LoadFriction = new AttributeID(0xE0),
        LoadUnloadCycleCount = new AttributeID(0xE1),
        LoadInTime = new AttributeID(0xE2),
        TorqueAmplificationCount = new AttributeID(0xE3),
        PowerOffRetractCycle = new AttributeID(0xE4),
        GMRHeadAmplitude = new AttributeID(0xE6),
        DriveTemperature = new AttributeID(0xE7),
        HeadFlyingHours = new AttributeID(0xF0),
        LBAsWrittenTotal = new AttributeID(0xF1),
        LBAsReadTotal = new AttributeID(0xF2),
        ReadErrorRetryRate = new AttributeID(0xFA),
        FreeFallProtection = new AttributeID(0xFE)
      ;
    }

    // Indilinx SSD SMART attributes
    // TODO: Find out the purpose of attribute 0xD2
    // Seems to be unique to Indilinx drives, hence its name of UnknownUnique.
    public static class IndilinxAttributes {
      public static readonly AttributeID
        ReadErrorRate = CommonAttributes.ReadErrorRate,
        PowerOnHours = CommonAttributes.PowerOnHours,
        PowerCycleCount = CommonAttributes.PowerCycleCount,
        InitialBadBlockCount = new AttributeID(0xB8),
        RemainingLife = new AttributeID(0xD1),
        ProgramFailure = new AttributeID(0xC3),
        EraseFailure = new AttributeID(0xC4),
        ReadFailure = new AttributeID(0xC5),
        SectorsRead = new AttributeID(0xC6),
        SectorsWritten = new AttributeID(0xC7),
        ReadCommands = new AttributeID(0xC8),
        WriteCommands = new AttributeID(0xC9),
        BitErrors = new AttributeID(0xCA),
        CorrectedErrors = new AttributeID(0xCB),
        BadBlockFullFlag = new AttributeID(0xCC),
        MaxCellcycles = new AttributeID(0xCD),
        MinErase = new AttributeID(0xCE),
        MaxErase = new AttributeID(0xCF),
        AverageEraseCount = new AttributeID(0xD0),
        UnknownUnique = new AttributeID(0xD2),
        SataErrorCountCRC = new AttributeID(0xD3),
        SataErrorCountHandshake = new AttributeID(0xD4)
      ;
    }

    // Intel SSD SMART attributes
    // TODO: Find out the meaning behind 0xE2, 0xE3 and 0xE4
    public static class IntelAttributes {
      public static readonly AttributeID
        ReadErrorRate = CommonAttributes.ReadErrorRate,
        SpinUpTime = CommonAttributes.SpinUpTime,
        StartStopCount = CommonAttributes.StartStopCount,
        ReallocatedSectorsCount = CommonAttributes.ReallocatedSectorsCount,
        PowerOnHours = CommonAttributes.PowerOnHours,
        PowerCycleCount = CommonAttributes.PowerCycleCount,
        EndToEndError = CommonAttributes.EndToEndError, // Only on G2 drives!

        // Different from the common attribute PowerOffRetractCount, same ID
        UnsafeShutdownCount = new AttributeID(0xC0),
        HostWrites = new AttributeID(0xE1),
        RemainingLife = new AttributeID(0xE8),
        MediaWearOutIndicator = new AttributeID(0xE9)
      ;
    }

    // Samsung SSD SMART attributes
    // TODO: AF, B0, B1, B5, B6, BB, C3, C6, C7, E8, E9
    public static class SamsungAttributes {
      public static readonly AttributeID
        PowerOnHours = CommonAttributes.PowerOnHours,
        PowerCycleCount = CommonAttributes.PowerCycleCount,
        UsedReservedBlockCountChip = new AttributeID(0xB2), // Unique
        UsedReservedBlockCountTotal = new AttributeID(0xB3), // Unique
        RemainingLife = new AttributeID(0xB4), // Unique
        RuntimeBadBlockTotal = new AttributeID(0xB7)
      ;
    }

    // SandForce SSD SMART attributes
    // Note: 0xE9 and 0xEA are reserved attributes and unique
    public static class SandForceAttributes {
      public static readonly AttributeID
        ReadErrorRate = CommonAttributes.ReadErrorRate,
        RetiredBlockCount = new AttributeID(0x05),
        PowerOnHours = CommonAttributes.PowerOnHours,
        PowerCycleCount = CommonAttributes.PowerCycleCount,
        ProgramFailCount = new AttributeID(0xAB), // Unique
        EraseFailCount = new AttributeID(0xAC), // Unique
        UnexpectedPowerLossCount = new AttributeID(0xAE), // Unique
        WearRangeDelta = new AttributeID(0xB1), // Unique
        ProgramFailCountAlt = new AttributeID(0xB5), // Same as 0xAB
        EraseFailCountAlt = new AttributeID(0xB6), // Same as 0xAC
        ReportedUncorrectableErrors =
          CommonAttributes.ReportedUncorrectableErrors,
        
        Temperature = CommonAttributes.Temperature, // SF-1500 only!
        
        // Opposite of the common attribute HardwareECCRecovered
        UnrecoverableECC = new AttributeID(0xC3),
        ReallocationEventCount = new AttributeID(0xC4),
        RemainingLife = new AttributeID(0xE7),
        LifetimeWrites = new AttributeID(0xF1),
        LifetimeReads = new AttributeID(0xF2)
      ;
    }

    [StructLayout(LayoutKind.Sequential, Pack = 1)]
    public struct DriveAttribute {
      public AttributeID ID;
      public Status StatusFlags;
      public byte AttrValue;
      public byte WorstValue;
      [MarshalAs(UnmanagedType.ByValArray, SizeConst = 6)]
      public byte[] RawValue;
      public byte Reserved;
    };

    [Flags]
    protected enum AccessMode : uint {     
      Read = 0x80000000,    
      Write = 0x40000000,     
      Execute = 0x20000000,     
      All = 0x10000000
    }

    [Flags]
    protected enum ShareMode : uint {
      None = 0,     
      Read = 1,     
      Write = 2,    
      Delete = 4
    }

    protected enum CreationMode : uint {
      New = 1,
      CreateAlways = 2,    
      OpenExisting = 3,    
      OpenAlways = 4,    
      TruncateExisting = 5
    }

    [Flags]
    protected enum FileAttribute : uint {
      Readonly = 0x00000001,
      Hidden = 0x00000002,
      System = 0x00000004,
      Directory = 0x00000010,
      Archive = 0x00000020,
      Device = 0x00000040,
      Normal = 0x00000080,
      Temporary = 0x00000100,
      SparseFile = 0x00000200,
      ReparsePoint = 0x00000400,
      Compressed = 0x00000800,
      Offline = 0x00001000,
      NotContentIndexed = 0x00002000,
      Encrypted = 0x00004000,
    }

    protected enum DriveCommand : uint {
      GetVersion = 0x00074080,
      SendDriveCommand = 0x0007c084,
      ReceiveDriveData = 0x0007c088
    }

    [StructLayout(LayoutKind.Sequential, Pack = 1)]
    protected struct CommandBlockRegisters {
      public byte Features;         
      public byte SectorCount;      
      public byte LBALow;       
      public byte LBAMid;           
      public byte LBAHigh;        
      public byte Device;       
      public byte Command;           
      public byte Reserved;                  
    }

    [StructLayout(LayoutKind.Sequential, Pack = 1)]
    protected struct DriveCommandParameter {
      public uint BufferSize;           
      public CommandBlockRegisters Registers;           
      public byte DriveNumber;   
      [MarshalAs(UnmanagedType.ByValArray, SizeConst = 20)]
      public byte[] Reserved;                                
    }

    [StructLayout(LayoutKind.Sequential, Pack = 1)]
    protected struct DriverStatus {
      public byte DriverError;   
      public byte IDEError;             
      [MarshalAs(UnmanagedType.ByValArray, SizeConst = 10)]
      public byte[] Reserved;               
    }

    [StructLayout(LayoutKind.Sequential, Pack = 1)]
    protected struct DriveCommandResult {
      public uint BufferSize;
      public DriverStatus DriverStatus;
    } 

    [StructLayout(LayoutKind.Sequential, Pack = 1)]
    protected struct DriveSmartReadResult {
      public uint BufferSize;           
      public DriverStatus DriverStatus;
      public byte Version;
      public byte Reserved;
      [MarshalAs(UnmanagedType.ByValArray, SizeConst = MAX_DRIVE_ATTRIBUTES)]
      public DriveAttribute[] Attributes;                                                                                       
    }

    [StructLayout(LayoutKind.Sequential, Pack = 1)]
    protected struct Identify {
      public ushort GeneralConfiguration;
      public ushort NumberOfCylinders;
      public ushort Reserved;
      public ushort NumberOfHeads;
      public ushort UnformattedBytesPerTrack;
      public ushort UnformattedBytesPerSector;
      public ushort SectorsPerTrack;
      [MarshalAs(UnmanagedType.ByValArray, SizeConst = 3)]
      public ushort[] VendorUnique;
      [MarshalAs(UnmanagedType.ByValArray, SizeConst = 20)]
      public byte[] SerialNumber;
      public ushort BufferType;
      public ushort BufferSectorSize;
      public ushort NumberOfEccBytes;
      [MarshalAs(UnmanagedType.ByValArray, SizeConst = 8)]
      public byte[] FirmwareRevision;
      [MarshalAs(UnmanagedType.ByValArray, SizeConst = 40)]
      public byte[] ModelNumber;
      public ushort MoreVendorUnique;
      public ushort DoubleWordIo;
      public ushort Capabilities;
      public ushort MoreReserved;
      public ushort PioCycleTimingMode;
      public ushort DmaCycleTimingMode;
      [MarshalAs(UnmanagedType.ByValArray, SizeConst = 406)]
      public byte[] More;
    }

    [StructLayout(LayoutKind.Sequential, Pack = 1)]
    protected struct DriveIdentifyResult {
      public uint BufferSize;
      public DriverStatus DriverStatus;
      public Identify Identify;
    } 

    public static readonly IntPtr INVALID_HANDLE_VALUE = (IntPtr)(-1);

    private const byte SMART_CMD = 0xB0;
    private const byte ID_CMD = 0xEC;
    
    private const byte SMART_READ_DATA = 0xD0;
    private const byte SMART_ENABLE_OPERATIONS = 0xD8;
    
    private const byte SMART_LBA_MID = 0x4F;
    private const byte SMART_LBA_HI = 0xC2;

    private const int MAX_DRIVE_ATTRIBUTES = 512;

    private SMART() { }

    public static IntPtr OpenPhysicalDrive(int driveNumber) {
      return NativeMethods.CreateFile(@"\\.\PhysicalDrive" + driveNumber,
        AccessMode.Read | AccessMode.Write, ShareMode.Read | ShareMode.Write,
        IntPtr.Zero, CreationMode.OpenExisting, FileAttribute.Device,
        IntPtr.Zero);
    }

    public static bool EnableSmart(IntPtr handle, int driveNumber) {
      DriveCommandParameter parameter = new DriveCommandParameter();
      DriveCommandResult result;
      uint bytesReturned;

      parameter.DriveNumber = (byte)driveNumber;
      parameter.Registers.Features = SMART_ENABLE_OPERATIONS;
      parameter.Registers.LBAMid = SMART_LBA_MID;
      parameter.Registers.LBAHigh = SMART_LBA_HI;
      parameter.Registers.Command = SMART_CMD;

      return NativeMethods.DeviceIoControl(handle, DriveCommand.SendDriveCommand, 
        ref parameter, Marshal.SizeOf(typeof(DriveCommandParameter)), out result,
        Marshal.SizeOf(typeof(DriveCommandResult)), out bytesReturned, 
        IntPtr.Zero);
    }

    public static DriveAttribute[] ReadSmart(IntPtr handle,
      int driveNumber)
    {
      DriveCommandParameter parameter = new DriveCommandParameter();
      DriveSmartReadResult result;
      uint bytesReturned;

      parameter.DriveNumber = (byte)driveNumber;
      parameter.Registers.Features = SMART_READ_DATA;
      parameter.Registers.LBAMid = SMART_LBA_MID;
      parameter.Registers.LBAHigh = SMART_LBA_HI;
      parameter.Registers.Command = SMART_CMD;

      bool isValid = NativeMethods.DeviceIoControl(handle, 
        DriveCommand.ReceiveDriveData, ref parameter, Marshal.SizeOf(parameter), 
        out result, Marshal.SizeOf(typeof(DriveSmartReadResult)), 
        out bytesReturned, IntPtr.Zero);

      return (isValid) ? result.Attributes : new DriveAttribute[0];
    }

    public static string ReadName(IntPtr handle, int driveNumber) {
      DriveCommandParameter parameter = new DriveCommandParameter();
      DriveIdentifyResult result;
      uint bytesReturned;

      parameter.DriveNumber = (byte)driveNumber;
      parameter.Registers.Command = ID_CMD;

      bool valid = NativeMethods.DeviceIoControl(handle, 
        DriveCommand.ReceiveDriveData, ref parameter, Marshal.SizeOf(parameter), 
        out result, Marshal.SizeOf(typeof(DriveIdentifyResult)), 
        out bytesReturned, IntPtr.Zero);

      if (!valid)
        return null;
      else {

        byte[] bytes = result.Identify.ModelNumber;
        char[] chars = new char[bytes.Length];
        for (int i = 0; i < bytes.Length; i += 2) {
          chars[i] = (char)bytes[i + 1];
          chars[i + 1] = (char)bytes[i];
        }

        return new string(chars).Trim();
      }
    }

    public static int CloseHandle(IntPtr handle) {
      return NativeMethods.CloseHandle(handle);
    }

    protected static class NativeMethods {
      private const string KERNEL = "kernel32.dll";

      [DllImport(KERNEL, CallingConvention = CallingConvention.Winapi,
        CharSet = CharSet.Unicode)]
      public static extern IntPtr CreateFile(string fileName,
        AccessMode desiredAccess, ShareMode shareMode, IntPtr securityAttributes,
        CreationMode creationDisposition, FileAttribute flagsAndAttributes,
        IntPtr templateFilehandle);

      [DllImport(KERNEL, CallingConvention = CallingConvention.Winapi)]
      public static extern int CloseHandle(IntPtr handle);

      [DllImport(KERNEL, CallingConvention = CallingConvention.Winapi)]
      [return: MarshalAsAttribute(UnmanagedType.Bool)]
      public static extern bool DeviceIoControl(IntPtr handle,
        DriveCommand command, ref DriveCommandParameter parameter,
        int parameterSize, out DriveSmartReadResult result, int resultSize,
        out uint bytesReturned, IntPtr overlapped);

      [DllImport(KERNEL, CallingConvention = CallingConvention.Winapi)]
      [return: MarshalAsAttribute(UnmanagedType.Bool)]
      public static extern bool DeviceIoControl(IntPtr handle,
        DriveCommand command, ref DriveCommandParameter parameter,
        int parameterSize, out DriveCommandResult result, int resultSize,
        out uint bytesReturned, IntPtr overlapped);

      [DllImport(KERNEL, CallingConvention = CallingConvention.Winapi)]
      [return: MarshalAsAttribute(UnmanagedType.Bool)]
      public static extern bool DeviceIoControl(IntPtr handle,
        DriveCommand command, ref DriveCommandParameter parameter,
        int parameterSize, out DriveIdentifyResult result, int resultSize,
        out uint bytesReturned, IntPtr overlapped);
    }    
  }
}