using System;
using Microsoft.WindowsAzure.Storage;
using Microsoft.WindowsAzure.Storage.Table; // same SDK as your sample
using Microsoft.Azure.Devices;              // Install-Package Microsoft.Azure.Devices
using Microsoft.Azure.Devices.Common.Exceptions;

namespace Tango.MachineService.Telemetry
{
    internal sealed class TelemetryDeviceRegistrationEntity : TableEntity
    {
        public TelemetryDeviceRegistrationEntity() { }

        public TelemetryDeviceRegistrationEntity(string machineGuid, string serialNumber)
        {
            PartitionKey = machineGuid;
            RowKey = SanitizeRowKey(serialNumber);
        }

        public String Environment { get; set; }
        public string SerialNumber { get; set; }
        public string DeviceId { get; set; }
        public string ConnectionString { get; set; }
        public DateTime CreatedUtc { get; set; }
        public DateTime UpdatedUtc { get; set; }

        internal static string SanitizeRowKey(string value) =>
            (value ?? string.Empty)
            .Replace("/", "_").Replace("\\", "_")
            .Replace("#", "_").Replace("?", "_");
    }

    /// <summary>
    /// Creates/gets IoT Hub device identities and caches the per-device connection string
    /// in Azure Table Storage ("TelemetryDeviceRegistrations").
    /// PartitionKey = MachineGuid, RowKey = SerialNumber (sanitized).
    /// </summary>
    public sealed class TelemetryDeviceRegistrationManager
    {
        private readonly string _iotHubServiceConnectionString;
        private readonly string _iotHubHostName;
        private readonly CloudTable _table;
        private readonly object _locker = new object();

        public TelemetryDeviceRegistrationManager(
            string storageAccountConnectionString,
            string iotHubServiceConnectionString,
            string tableName = "TelemetryDeviceRegistrations")
        {
            if (string.IsNullOrWhiteSpace(storageAccountConnectionString))
                throw new ArgumentNullException(nameof(storageAccountConnectionString));
            if (string.IsNullOrWhiteSpace(iotHubServiceConnectionString))
                throw new ArgumentNullException(nameof(iotHubServiceConnectionString));
            if (string.IsNullOrWhiteSpace(tableName))
                throw new ArgumentNullException(nameof(tableName));

            _iotHubServiceConnectionString = iotHubServiceConnectionString;
            _iotHubHostName = ExtractHostName(iotHubServiceConnectionString);

            var storageAccount = CloudStorageAccount.Parse(storageAccountConnectionString);
            var tableClient = storageAccount.CreateCloudTableClient();
            _table = tableClient.GetTableReference(tableName);
            _table.CreateIfNotExists(); // sync, as in your sample
        }

        /// <summary>
        /// Returns the per-device IoT Hub connection string for (machineGuid, serialNumber).
        /// Creates the IoT Hub device (DeviceId = serialNumber) if needed, persists the connection string,
        /// then returns the cached value next time.
        /// </summary>
        public string GetOrCreateDeviceConnectionString(string machineGuid, string serialNumber)
        {
            if (string.IsNullOrWhiteSpace(machineGuid)) throw new ArgumentNullException(nameof(machineGuid));
            if (string.IsNullOrWhiteSpace(serialNumber)) throw new ArgumentNullException(nameof(serialNumber));

            // 1) Try table cache
            var cached = TryGetFromTable(machineGuid, serialNumber);
            if (!string.IsNullOrEmpty(cached))
                return cached;

            // 2) Create/get in IoT Hub (guard against concurrent callers)
            lock (_locker)
            {
                // double-check after lock
                cached = TryGetFromTable(machineGuid, serialNumber);
                if (!string.IsNullOrEmpty(cached))
                    return cached;

                var registry = RegistryManager.CreateFromConnectionString(_iotHubServiceConnectionString);

                Device device;
                try
                {
                    device = registry.AddDeviceAsync(new Device(serialNumber)).GetAwaiter().GetResult();
                }
                catch (DeviceAlreadyExistsException)
                {
                    device = registry.GetDeviceAsync(serialNumber).GetAwaiter().GetResult();
                }

                if (device == null)
                    throw new InvalidOperationException("Failed to get or create IoT Hub device.");

                var primaryKey = device.Authentication?.SymmetricKey?.PrimaryKey;
                if (string.IsNullOrEmpty(primaryKey))
                    throw new InvalidOperationException("Device has no symmetric key.");

                var deviceConnectionString =
                    $"HostName={_iotHubHostName};DeviceId={serialNumber};SharedAccessKey={primaryKey}";

                UpsertToTable(machineGuid, serialNumber, deviceConnectionString);

                return deviceConnectionString;
            }
        }

        // ---- Table helpers (sync) ----

        private string TryGetFromTable(string machineGuid, string serialNumber)
        {
            var rowKey = TelemetryDeviceRegistrationEntity.SanitizeRowKey(serialNumber);
            var retrieve = TableOperation.Retrieve<TelemetryDeviceRegistrationEntity>(machineGuid, rowKey);
            var result = _table.Execute(retrieve);
            var entity = result.Result as TelemetryDeviceRegistrationEntity;
            return entity?.ConnectionString;
        }

        private void UpsertToTable(string machineGuid, string serialNumber, string connectionString)
        {
            var entity = new TelemetryDeviceRegistrationEntity(machineGuid, serialNumber)
            {
                SerialNumber = serialNumber,
                DeviceId = serialNumber,
                Environment = MachineServiceConfig.DEPLOYMENT_SLOT.ToString(),
                ConnectionString = connectionString,
                CreatedUtc = DateTime.UtcNow,
                UpdatedUtc = DateTime.UtcNow
            };

            var op = TableOperation.InsertOrReplace(entity);
            _table.Execute(op);
        }

        // ---- Utility ----

        private static string ExtractHostName(string iotHubServiceConnectionString)
        {
            var parts = iotHubServiceConnectionString.Split(';');
            foreach (var p in parts)
                if (p.StartsWith("HostName=", StringComparison.OrdinalIgnoreCase))
                    return p.Substring("HostName=".Length);
            throw new ArgumentException("IoT Hub service connection string is missing HostName=…");
        }
    }
}