dorisoy_HardwareMonitor/LibreHardwareMonitorLib/Hardware/Storage/ATAStorage.cs
// This Source Code Form is subject to the terms of the Mozilla Public License, v. 2.0.
// If a copy of the MPL was not distributed with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
// Copyright (C) LibreHardwareMonitor and Contributors.
// Partial Copyright (C) Michael Möller <mmoeller@openhardwaremonitor.org> and Contributors.
// All Rights Reserved.

using System;
using System.Collections.Generic;
using System.Globalization;
using System.IO;
using System.Linq;
using System.Reflection;
using System.Text;
using LibreHardwareMonitor.Interop;

namespace LibreHardwareMonitor.Hardware.Storage
{
    public abstract class AtaStorage : AbstractStorage
    {
        // array of all hard drive types, matching type is searched in this order
        private static readonly Type[] _hddTypes = { typeof(SsdPlextor), typeof(SsdIntel), typeof(SsdSandforce), typeof(SsdIndilinx), typeof(SsdSamsung), typeof(SsdMicron), typeof(GenericHardDisk) };

        private IDictionary<SmartAttribute, Sensor> _sensors;

        /// <summary>
        /// Gets the SMART data.
        /// </summary>
        public ISmart Smart { get; }

        /// <summary>
        /// Gets the SMART attributes.
        /// </summary>
        public IReadOnlyList<SmartAttribute> SmartAttributes { get; }

        internal AtaStorage(StorageInfo storageInfo, ISmart smart, string name, string firmwareRevision, string id, int index, IReadOnlyList<SmartAttribute> smartAttributes, ISettings settings)
          : base(storageInfo, name, firmwareRevision, id, index, settings)
        {
            Smart = smart;
            if (smart.IsValid)
                smart.EnableSmart();
            
            SmartAttributes = smartAttributes;
            CreateSensors();
        }

        internal static AbstractStorage CreateInstance(StorageInfo storageInfo, ISettings settings)
        {
            ISmart smart = new WindowsSmart(storageInfo.Index);
            string name = null;
            string firmwareRevision = null;
            Kernel32.SMART_ATTRIBUTE[] smartAttributes = { };

            if (smart.IsValid)
            {
                bool nameValid = smart.ReadNameAndFirmwareRevision(out name, out firmwareRevision);
                bool smartEnabled = smart.EnableSmart();

                if (smartEnabled)
                    smartAttributes = smart.ReadSmartData();
                
                if (!nameValid)
                {
                    name = null;
                    firmwareRevision = null;
                }
            }
            else
            {
                string[] logicalDrives = WindowsStorage.GetLogicalDrives(storageInfo.Index);
                if (logicalDrives == null || logicalDrives.Length == 0)
                {
                    smart.Close();
                    return null;
                }

                bool hasNonZeroSizeDrive = false;
                foreach (string logicalDrive in logicalDrives)
                {
                    try
                    {
                        var driveInfo = new DriveInfo(logicalDrive);
                        if (driveInfo.TotalSize > 0)
                        {
                            hasNonZeroSizeDrive = true;
                            break;
                        }
                    }
                    catch (ArgumentException) { }
                    catch (IOException) { }
                    catch (UnauthorizedAccessException) { }
                }

                if (!hasNonZeroSizeDrive)
                {
                    smart.Close();
                    return null;
                }
            }

            if (string.IsNullOrEmpty(name))
                name = string.IsNullOrEmpty(storageInfo.Name) ? "Generic Hard Disk" : storageInfo.Name;

            if (string.IsNullOrEmpty(firmwareRevision))
                firmwareRevision = string.IsNullOrEmpty(storageInfo.Revision) ? "Unknown" : storageInfo.Revision;

            foreach (Type type in _hddTypes)
            {
                // get the array of the required SMART attributes for the current type

                // check if all required attributes are present
                bool allAttributesFound = true;

                if (type.GetCustomAttributes(typeof(RequireSmartAttribute), true) is RequireSmartAttribute[] requiredAttributes)
                {
                    foreach (RequireSmartAttribute requireAttribute in requiredAttributes)
                    {
                        bool attributeFound = false;

                        foreach (Kernel32.SMART_ATTRIBUTE value in smartAttributes)
                        {
                            if (value.Id == requireAttribute.AttributeId)
                            {
                                attributeFound = true;
                                break;
                            }
                        }

                        if (!attributeFound)
                        {
                            allAttributesFound = false;
                            break;
                        }
                    }
                }

                // if an attribute is missing, then try the next type
                if (!allAttributesFound)
                    continue;


                // check if there is a matching name prefix for this type
                if (type.GetCustomAttributes(typeof(NamePrefixAttribute), true) is NamePrefixAttribute[] namePrefixes)
                {
                    foreach (NamePrefixAttribute prefix in namePrefixes)
                    {
                        if (name.StartsWith(prefix.Prefix, StringComparison.InvariantCulture))
                        {
                            const BindingFlags flags = BindingFlags.NonPublic | BindingFlags.Public | BindingFlags.Instance;

                            return Activator.CreateInstance(type, flags, null, new object[] { storageInfo, smart, name, firmwareRevision, storageInfo.Index, settings }, null) as AtaStorage;
                        }
                    }
                }
            }

            // no matching type has been found
            smart.Close();
            return null;
        }

        protected sealed override void CreateSensors()
        {
            _sensors = new Dictionary<SmartAttribute, Sensor>();

            if (Smart.IsValid)
            {
                byte[] smartIds = Smart.ReadSmartData().Select(x => x.Id).ToArray();

                // unique attributes by SensorType and SensorChannel.
                IEnumerable<SmartAttribute> smartAttributes = SmartAttributes
                                      .Where(x => x.SensorType.HasValue && smartIds.Contains(x.Id))
                                      .GroupBy(x => new { x.SensorType.Value, x.SensorChannel })
                                      .Select(x => x.First());

                _sensors = smartAttributes.ToDictionary(attribute => attribute,
                                                        attribute => new Sensor(attribute.SensorName,
                                                                           attribute.SensorChannel,
                                                                           attribute.DefaultHiddenSensor,
                                                                           attribute.SensorType.GetValueOrDefault(),
                                                                           this,
                                                                           attribute.ParameterDescriptions,
                                                                           _settings));

                foreach (KeyValuePair<SmartAttribute, Sensor> sensor in _sensors)
                    ActivateSensor(sensor.Value);
            }

            base.CreateSensors();
        }

        protected virtual void UpdateAdditionalSensors(Kernel32.SMART_ATTRIBUTE[] values) { }

        protected override void UpdateSensors()
        {
            if (Smart.IsValid)
            {
                Kernel32.SMART_ATTRIBUTE[] smartAttributes = Smart.ReadSmartData();

                foreach (KeyValuePair<SmartAttribute, Sensor> keyValuePair in _sensors)
                {
                    SmartAttribute attribute = keyValuePair.Key;
                    foreach (Kernel32.SMART_ATTRIBUTE value in smartAttributes)
                    {
                        if (value.Id == attribute.Id)
                        {
                            Sensor sensor = keyValuePair.Value;
                            sensor.Value = attribute.ConvertValue(value, sensor.Parameters);
                        }
                    }
                }

                UpdateAdditionalSensors(smartAttributes);
            }
        }

        protected override void GetReport(StringBuilder r)
        {
            if (Smart.IsValid)
            {
                Kernel32.SMART_ATTRIBUTE[] values = Smart.ReadSmartData();
                Kernel32.SMART_THRESHOLD[] thresholds = Smart.ReadSmartThresholds();
                if (values.Length > 0)
                {
                    r.AppendFormat(CultureInfo.InvariantCulture,
                                   " {0}{1}{2}{3}{4}{5}{6}{7}",
                                   "Id".PadRight(3),
                                   "Description".PadRight(35),
                                   "Raw Value".PadRight(13),
                                   "Worst".PadRight(6),
                                   "Value".PadRight(6),
                                   "Threshold".PadRight(6),
                                   "Physical".PadRight(8),
                                   Environment.NewLine);

                    foreach (Kernel32.SMART_ATTRIBUTE value in values)
                    {
                        if (value.Id == 0x00)
                            break;

                        byte? threshold = null;
                        foreach (Kernel32.SMART_THRESHOLD t in thresholds)
                        {
                            if (t.Id == value.Id)
                            {
                                threshold = t.Threshold;
                            }
                        }

                        string description = "Unknown";
                        float? physical = null;
                        foreach (SmartAttribute a in SmartAttributes)
                        {
                            if (a.Id == value.Id)
                            {
                                description = a.Name;
                                if (a.HasRawValueConversion | a.SensorType.HasValue)
                                    physical = a.ConvertValue(value, null);
                                else
                                    physical = null;
                            }
                        }

                        string raw = BitConverter.ToString(value.RawValue);
                        r.AppendFormat(CultureInfo.InvariantCulture,
                                       " {0}{1}{2}{3}{4}{5}{6}{7}",
                                       value.Id.ToString("X2").PadRight(3),
                                       description.PadRight(35),
                                       raw.Replace("-", string.Empty).PadRight(13),
                                       value.WorstValue.ToString(CultureInfo.InvariantCulture).PadRight(6),
                                       value.CurrentValue.ToString(CultureInfo.InvariantCulture).PadRight(6),
                                       (threshold.HasValue
                                         ? threshold.Value.ToString(CultureInfo.InvariantCulture)
                                         : "-").PadRight(6),
                                       (physical.HasValue ? physical.Value.ToString(CultureInfo.InvariantCulture) : "-").PadRight(8),
                                       Environment.NewLine);
                    }

                    r.AppendLine();
                }
            }
        }

        protected static float RawToInt(byte[] raw, byte value, IReadOnlyList<IParameter> parameters)
        {
            return (raw[3] << 24) | (raw[2] << 16) | (raw[1] << 8) | raw[0];
        }

        public override void Close()
        {
            Smart.Close();
            base.Close();
        }
    }
}
