// Copyright (C) 2026, The Duplicati Team
// https://duplicati.com, hello@duplicati.com
//
// Permission is hereby granted, free of charge, to any person obtaining a
// copy of this software and associated documentation files (the "Software"),
// to deal in the Software without restriction, including without limitation
// the rights to use, copy, modify, merge, publish, distribute, sublicense,
// and/or sell copies of the Software, and to permit persons to whom the
// Software is furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
// OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
// DEALINGS IN THE SOFTWARE.

using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.IO;
using System.Linq;
using System.Text.RegularExpressions;
using System.Threading;
using Duplicati.Proprietary.DiskImage;
using NUnit.Framework;
using Duplicati.Proprietary.DiskImage.General;

#nullable enable

namespace Duplicati.UnitTest.DiskImage.Helpers;

/// <summary>
/// Linux implementation of <see cref="IDiskImageHelper"/> using loop devices and standard Linux tools.
/// </summary>
internal class LinuxDiskImageHelper : IDiskImageHelper
{
    /// <summary>
    /// Tracks loop devices created by this helper for cleanup.
    /// </summary>
    private readonly List<string> _trackedLoopDevices = new();

    /// <summary>
    /// Maps image paths to their associated loop devices.
    /// </summary>
    private readonly Dictionary<string, string> _imageToLoopDevice = new();

    /// <summary>
    /// Tracks any mount directories that were automatically created for cleanup.
    /// </summary>
    private readonly List<string> _autoCreatedMntDirectories = new();

    /// <inheritdoc />
    public string CreateDisk(string imagePath, long sizeB)
    {
        // Ensure the directory exists
        var directory = Path.GetDirectoryName(imagePath);
        if (!string.IsNullOrEmpty(directory) && !Directory.Exists(directory))
            Directory.CreateDirectory(directory);

        // Delete existing image if it exists
        if (File.Exists(imagePath))
        {
            try
            {
                CleanupDisk(imagePath);
            }
            catch
            {
                // Ignore errors during cleanup
            }
        }

        // Create a sparse disk image file using truncate
        RunProcess("truncate", $"-s {sizeB} \"{imagePath}\"");

        // Attach it as a loop device
        var loopDevice = RunProcess("losetup", $"--find --show --partscan \"{imagePath}\"").Trim();

        if (string.IsNullOrEmpty(loopDevice) || !loopDevice.StartsWith("/dev/loop"))
            throw new InvalidOperationException($"Failed to attach loop device for {imagePath}");

        // Track the loop device
        _trackedLoopDevices.Add(loopDevice);
        _imageToLoopDevice[imagePath] = loopDevice;

        return loopDevice;
    }

    /// <inheritdoc />
    public string[] InitializeDisk(string diskIdentifier, PartitionTableType tableType, (FileSystemType, long)[] partitions)
    {
        if (tableType == PartitionTableType.Unknown)
            return []; // No partition table, so nothing to initialize

        // Create partition table using parted
        var label = tableType == PartitionTableType.GPT ? "gpt" : "msdos";
        RunProcess("parted", $"-s {diskIdentifier} mklabel {label}");

        // Get disk size for calculating partition sizes
        var diskSize = GetDiskSize(diskIdentifier);
        var sectorSize = GetSectorSize(diskIdentifier);
        var mbrSize = 1 * sectorSize; // MBR requires 1 sector for the partition table
        var gptSize = mbrSize + 33 * sectorSize; // GPT requires 34 sectors reserved at the beginning (1 for protective MBR + 33 for GPT header and partition entries)
        long nextFreeSector = tableType == PartitionTableType.GPT ? gptSize : mbrSize;

        // Create partitions
        for (int i = 0; i < partitions.Length; i++)
        {
            var (fsType, size) = partitions[i];
            var partitionNumber = i + 1;

            // Calculate partition size
            long endOffset;
            if (size <= 0)
            {
                if (i != partitions.Length - 1)
                    throw new InvalidOperationException("Only the last partition can be sized to fill the remaining space");
                endOffset = diskSize - sectorSize;
            }
            else
                endOffset = nextFreeSector + size - sectorSize;

            // Check if we should clamp
            if (tableType == PartitionTableType.GPT && endOffset > diskSize - gptSize)
                if (size <= 0)
                    endOffset = diskSize - gptSize;
                else
                    throw new InvalidOperationException($"Partition {partitionNumber} exceeds available space. Max size for this partition is {diskSize - gptSize - nextFreeSector} bytes.");

            // Create partition
            RunProcess("parted", $"-s {diskIdentifier} mkpart primary {nextFreeSector}B {endOffset}B");

            // For MBR, set the boot flag on the first partition (some filesystems need this)
            if (tableType == PartitionTableType.MBR && i == 0)
            {
                try
                {
                    RunProcess("parted", $"-s {diskIdentifier} set {partitionNumber} boot on");
                }
                catch
                {
                    // Ignore errors setting boot flag
                }
            }

            nextFreeSector = endOffset + sectorSize;
        }

        // Re-read partition table
        RunProcess("partprobe", $"{diskIdentifier}");

        // Small delay to allow kernel to create partition devices
        System.Threading.Thread.Sleep(200);

        // Format each partition
        for (int i = 0; i < partitions.Length; i++)
        {
            var (fsType, _) = partitions[i];
            var partitionNumber = i + 1;
            var partitionDevice = GetPartitionDevice(diskIdentifier, partitionNumber);

            // Wait for partition device to exist
            var retryCount = 0;
            while (!File.Exists(partitionDevice) && retryCount < 10)
            {
                System.Threading.Thread.Sleep(100);
                retryCount++;
            }

            if (!File.Exists(partitionDevice))
                throw new InvalidOperationException($"Partition device {partitionDevice} did not appear after creation");

            // Format the partition
            FormatPartition(partitionDevice, fsType);
        }

        return partitions.Length == 0 ? [] : Mount(diskIdentifier, null, false, null);
    }

    /// <inheritdoc />
    public string[] Mount(string diskIdentifier, string? baseMountPath = null, bool readOnly = false, FileSystemType[]? fileSystemTypes = null)
    {
        var mountPoints = new List<string>();
        var partitions = GetPartitionNames(diskIdentifier);

        if (partitions.Count == 0)
            throw new InvalidOperationException($"No partitions found for disk {diskIdentifier}");
        if (fileSystemTypes is not null && fileSystemTypes.Length != partitions.Count)
            throw new ArgumentException($"Length of fileSystemTypes array ({fileSystemTypes.Length}) must match the number of partitions ({partitions.Count})");

        for (int i = 0; i < partitions.Count; i++)
        {
            var partitionDevice = partitions[i];

            // Get filesystem type to handle special cases
            FileSystemType? fsType = null;
            if (fileSystemTypes is not null)
                fsType = fileSystemTypes[i];

            // Create mount point
            string mountPoint;
            if (!string.IsNullOrEmpty(baseMountPath))
            {
                var partitionName = Path.GetFileName(partitionDevice);
                mountPoint = Path.Combine(baseMountPath, partitionName);
            }
            else
            {
                mountPoint = $"/mnt/duplicati_{Path.GetFileName(partitionDevice)}_{Guid.NewGuid():N}";
                _autoCreatedMntDirectories.Add(mountPoint);
            }

            Directory.CreateDirectory(mountPoint);

            // Mount the partition
            List<string> options = [];
            if (readOnly)
                options.Add("ro");
            if (fsType == FileSystemType.XFS)
                options.Add("nouuid"); // XFS does not allow mounting with duplicate UUIDs, so we ignore UUIDs for testing
            var optionsArg = options.Count > 0 ? $"-o {string.Join(",", options)}" : "";
            var mountArgs = $"{optionsArg} {partitionDevice} \"{mountPoint}\"".Trim();
            RunProcess("mount", mountArgs);

            mountPoints.Add(mountPoint);
        }

        return mountPoints.ToArray();
    }

    /// <inheritdoc />
    public void Unmount(string diskIdentifier)
    {
        // Find all mounted partitions of this device
        var partitions = GetPartitionNames(diskIdentifier);

        foreach (var partitionDevice in partitions)
        {
            // Check if mounted
            var mountInfo = GetMountPoint(partitionDevice);
            if (!string.IsNullOrEmpty(mountInfo))
            {
                RunProcess("umount", partitionDevice);
            }
        }
    }

    /// <inheritdoc />
    public void CleanupDisk(string imagePath, string? diskIdentifier = null)
    {
        // Determine the loop device if not provided
        if (string.IsNullOrEmpty(diskIdentifier))
        {
            if (_imageToLoopDevice.TryGetValue(imagePath, out var trackedDevice))
            {
                diskIdentifier = trackedDevice;
            }
            else
            {
                // Try to find the loop device using losetup -j
                try
                {
                    var output = RunProcess("losetup", $"-j \"{imagePath}\"");
                    if (!string.IsNullOrEmpty(output))
                    {
                        // Output format: /dev/loop0: [0005]:123456 (/path/to/image)
                        var match = Regex.Match(output, @"^(\S+):");
                        if (match.Success)
                            diskIdentifier = match.Groups[1].Value;
                    }
                }
                catch
                {
                    // Ignore errors
                }
            }
        }

        // Unmount partitions
        if (!string.IsNullOrEmpty(diskIdentifier))
        {
            bool unmounted = false;
            for (int i = 0; i < 5; i++)
            {
                try
                {
                    Unmount(diskIdentifier);
                    unmounted = true;
                    break;
                }
                catch
                {
                    // Try to wait
                    Thread.Sleep(1000);
                }
            }
            if (!unmounted)
            {
                TestContext.Progress.WriteLine($"Warning: Failed to unmount partitions for {diskIdentifier} after multiple attempts.");
            }

            // Detach loop device
            try
            {
                RunProcess("losetup", $"-d {diskIdentifier}");
            }
            catch (Exception ex)
            {
                TestContext.Progress.WriteLine($"Warning: Failed to detach loop device {diskIdentifier}: {ex.Message}");
            }

            _trackedLoopDevices.Remove(diskIdentifier);
        }

        // Delete the image file
        if (File.Exists(imagePath))
        {
            try
            {
                File.Delete(imagePath);
            }
            catch (Exception ex)
            {
                TestContext.Progress.WriteLine($"Warning: Failed to delete image file {imagePath}: {ex.Message}");
            }
        }

        _imageToLoopDevice.Remove(imagePath);

        // Clean up any auto-created mount directories
        foreach (var mntDir in _autoCreatedMntDirectories)
        {
            try
            {
                if (Directory.Exists(mntDir))
                    Directory.Delete(mntDir);
            }
            catch (Exception ex)
            {
                TestContext.Progress.WriteLine($"Warning: Failed to delete mount directory {mntDir}: {ex.Message}");
            }
        }
        _autoCreatedMntDirectories.Clear();
    }

    /// <inheritdoc />
    public bool HasRequiredPrivileges()
    {
        // Check if running as root (uid 0)
        try
        {
            var uid = RunProcess("id", "-u").Trim();
            return uid == "0";
        }
        catch
        {
            return false;
        }
    }

    /// <inheritdoc />
    public PartitionTableGeometry GetPartitionTable(string diskIdentifier)
    {
        var deviceName = Path.GetFileName(diskIdentifier);
        var tableType = PartitionTableType.Unknown;

        // Use blkid to get partition table type
        try
        {
            var output = RunProcess("blkid", $"-p -o value -s PTTYPE {diskIdentifier}").Trim();
            tableType = output.ToLower() switch
            {
                "gpt" => PartitionTableType.GPT,
                "dos" => PartitionTableType.MBR,
                "msdos" => PartitionTableType.MBR,
                _ => PartitionTableType.Unknown
            };
        }
        catch
        {
            // Ignore errors
        }

        var size = GetDiskSize(diskIdentifier);
        var sectorSize = GetSectorSize(diskIdentifier);

        return new PartitionTableGeometry
        {
            Type = tableType,
            Size = size,
            SectorSize = sectorSize
        };
    }

    /// <inheritdoc />
    public PartitionGeometry[] GetPartitions(string diskIdentifier)
    {
        var partitions = new List<PartitionGeometry>();
        var partitionNames = GetPartitionNames(diskIdentifier);

        for (int i = 0; i < partitionNames.Count; i++)
        {
            var partitionDevice = partitionNames[i];
            var partitionNumber = i + 1;

            try
            {
                // Get partition info using blockdev
                var start = 0L;
                var size = 0L;

                try
                {
                    var devName = Path.GetFileName(partitionDevice);
                    var startFile = $"/sys/class/block/{devName}/start";
                    if (File.Exists(startFile))
                    {
                        var startStr = File.ReadAllText(startFile).Trim();
                        if (long.TryParse(startStr, out var startSector))
                        {
                            start = startSector * 512;
                        }
                    }
                    else
                    {
                        var startStr = RunProcess("blockdev", $"--getstart {partitionDevice}").Trim();
                        if (long.TryParse(startStr, out var startSector))
                        {
                            // blockdev --getstart always returns the start sector in 512-byte units
                            start = startSector * 512;
                        }
                    }
                }
                catch (Exception ex)
                {
                    TestContext.Progress.WriteLine($"Warning: Failed to get start for {partitionDevice}: {ex.Message}");
                }

                try
                {
                    size = GetDiskSize(partitionDevice);
                }
                catch { }

                // Get filesystem type
                var fsType = FileSystemType.Unknown;
                try
                {
                    var fsStr = RunProcess("blkid", $"-o value -s TYPE {partitionDevice}").Trim();
                    fsType = ParseFilesystemType(fsStr);
                }
                catch { }

                partitions.Add(new PartitionGeometry
                {
                    Number = partitionNumber,
                    Type = PartitionType.Unknown,
                    StartOffset = start,
                    Size = size,
                    FilesystemType = fsType,
                    TableType = PartitionTableType.Unknown
                });
            }
            catch (Exception ex)
            {
                TestContext.Progress.WriteLine($"Warning: Failed to get info for partition {partitionDevice}: {ex.Message}");
            }
        }

        return partitions.ToArray();
    }

    /// <inheritdoc />
    public void FlushDisk(string diskIdentifier)
    {
        // Sync all filesystems
        RunProcess("sync", "");

        // Flush buffers for the specific device
        try
        {
            RunProcess("blockdev", $"--flushbufs {diskIdentifier}");
        }
        catch
        {
            // Ignore errors
        }
    }

    /// <inheritdoc />
    public string ReAttach(string imagePath, string diskIdentifier, PartitionTableType tableType, bool readOnly = false)
    {
        // Unmount and detach
        try
        {
            Unmount(diskIdentifier);
        }
        catch { }

        try
        {
            RunProcess("losetup", $"-d {diskIdentifier}");
        }
        catch { }

        // Re-attach the loop device
        var readOnlyArg = readOnly ? "--read-only" : "";
        var newLoopDevice = RunProcess("losetup", $"--find --show --partscan {readOnlyArg} \"{imagePath}\"").Trim();

        if (string.IsNullOrEmpty(newLoopDevice) || !newLoopDevice.StartsWith("/dev/loop"))
            throw new InvalidOperationException($"Failed to re-attach loop device for {imagePath}");

        // Update tracking
        _trackedLoopDevices.Remove(diskIdentifier);
        _trackedLoopDevices.Add(newLoopDevice);
        _imageToLoopDevice[imagePath] = newLoopDevice;

        // Allow time for partition discovery
        System.Threading.Thread.Sleep(200);

        return newLoopDevice;
    }

    /// <inheritdoc />
    public void Dispose()
    {
        // Clean up any remaining tracked loop devices
        foreach (var loopDevice in _trackedLoopDevices.ToList())
        {
            try
            {
                // Find associated image
                var imagePath = _imageToLoopDevice.FirstOrDefault(x => x.Value == loopDevice).Key;
                if (!string.IsNullOrEmpty(imagePath))
                {
                    CleanupDisk(imagePath, loopDevice);
                }
                else
                {
                    // Just try to detach the loop device
                    RunProcess("losetup", $"-d {loopDevice}");
                }
            }
            catch (Exception ex)
            {
                TestContext.Progress.WriteLine($"Warning: Failed to dispose loop device {loopDevice}: {ex.Message}");
            }
        }

        _trackedLoopDevices.Clear();
        _imageToLoopDevice.Clear();
        GC.SuppressFinalize(this);
    }

    #region Helper Methods

    /// <summary>
    /// Runs a process and returns the output.
    /// </summary>
    private static string RunProcess(string fileName, string arguments)
    {
        var psi = new ProcessStartInfo
        {
            FileName = fileName,
            Arguments = arguments,
            RedirectStandardOutput = true,
            RedirectStandardError = true,
            UseShellExecute = false,
            CreateNoWindow = true
        };

        using var process = new Process { StartInfo = psi };
        process.Start();
        string output = process.StandardOutput.ReadToEnd();
        string error = process.StandardError.ReadToEnd();
        process.WaitForExit();

        if (process.ExitCode != 0)
        {
            throw new InvalidOperationException($"Process {fileName} exited with code {process.ExitCode}: {error}");
        }

        return output;
    }

    /// <summary>
    /// Gets the partition device path for a given disk and partition number.
    /// </summary>
    private static string GetPartitionDevice(string diskIdentifier, int partitionNumber)
    {
        // For loop devices, partitions are named /dev/loop0p1, /dev/loop0p2, etc.
        if (diskIdentifier.StartsWith("/dev/loop"))
        {
            return $"{diskIdentifier}p{partitionNumber}";
        }

        // For NVMe devices, partitions are named /dev/nvme0n1p1, etc.
        if (diskIdentifier.Contains("nvme"))
        {
            return $"{diskIdentifier}p{partitionNumber}";
        }

        // For standard SCSI/SATA devices, partitions are named /dev/sda1, /dev/sda2, etc.
        return $"{diskIdentifier}{partitionNumber}";
    }

    /// <summary>
    /// Gets the list of partition device paths for a disk.
    /// </summary>
    private static List<string> GetPartitionNames(string diskIdentifier)
    {
        var partitions = new List<string>();
        var deviceName = Path.GetFileName(diskIdentifier);

        // Check for partition devices
        if (diskIdentifier.StartsWith("/dev/loop"))
        {
            // Loop device partitions: /dev/loop0p1, /dev/loop0p2, etc.
            for (int i = 1; i <= 16; i++)
            {
                var partitionDevice = $"{diskIdentifier}p{i}";
                if (File.Exists(partitionDevice))
                    partitions.Add(partitionDevice);
                else if (i > 1)
                    break; // Stop after first missing partition
            }
        }
        else if (diskIdentifier.Contains("nvme"))
        {
            // NVMe partitions: /dev/nvme0n1p1, etc.
            for (int i = 1; i <= 16; i++)
            {
                var partitionDevice = $"{diskIdentifier}p{i}";
                if (File.Exists(partitionDevice))
                    partitions.Add(partitionDevice);
                else if (i > 1)
                    break;
            }
        }
        else
        {
            // Standard partitions: /dev/sda1, /dev/sda2, etc.
            for (int i = 1; i <= 16; i++)
            {
                var partitionDevice = $"{diskIdentifier}{i}";
                if (File.Exists(partitionDevice))
                    partitions.Add(partitionDevice);
                else if (i > 1)
                    break;
            }
        }

        return partitions;
    }

    /// <summary>
    /// Gets the mount point for a partition device.
    /// </summary>
    private static string? GetMountPoint(string partitionDevice)
    {
        try
        {
            var output = RunProcess("findmnt", $"--raw --noheadings -o TARGET {partitionDevice}").Trim();
            return string.IsNullOrEmpty(output) ? null : output;
        }
        catch
        {
            return null;
        }
    }

    /// <summary>
    /// Gets the filesystem type of a partition.
    /// </summary>
    private static string GetFilesystemType(string partitionDevice)
    {
        try
        {
            return RunProcess("blkid", $"-o value -s TYPE {partitionDevice}").Trim();
        }
        catch
        {
            return string.Empty;
        }
    }

    /// <summary>
    /// Formats a partition with the specified filesystem type.
    /// </summary>
    private static void FormatPartition(string partitionDevice, FileSystemType fsType)
    {
        var (command, args) = fsType switch
        {
            FileSystemType.FAT12 => ("mkfs.vfat", $"-F 12 {partitionDevice}"),
            FileSystemType.FAT16 => ("mkfs.vfat", $"-F 16 {partitionDevice}"),
            FileSystemType.FAT32 => ("mkfs.vfat", $"-F 32 {partitionDevice}"),
            FileSystemType.ExFAT => ("mkfs.exfat", $"{partitionDevice}"),
            FileSystemType.Ext2 => ("mkfs.ext2", $"-F {partitionDevice}"),
            FileSystemType.Ext3 => ("mkfs.ext3", $"-F {partitionDevice}"),
            FileSystemType.Ext4 => ("mkfs.ext4", $"-F {partitionDevice}"),
            FileSystemType.XFS => ("mkfs.xfs", $"-f {partitionDevice}"),
            FileSystemType.Btrfs => ("mkfs.btrfs", $"-f {partitionDevice}"),
            FileSystemType.NTFS => ("mkfs.ntfs", $"-F {partitionDevice}"),
            FileSystemType.ZFS => throw new NotSupportedException("ZFS is not supported for testing. ZFS requires pool creation which is incompatible with standard partition-based testing."),
            _ => throw new NotSupportedException($"Unsupported filesystem type on Linux: {fsType}")
        };

        RunProcess(command, args);
    }

    /// <summary>
    /// Gets the size of a disk or partition in bytes.
    /// </summary>
    private static long GetDiskSize(string devicePath)
    {
        var output = RunProcess("blockdev", $"--getsize64 {devicePath}").Trim();
        return long.TryParse(output, out var size) ? size : 0;
    }

    /// <summary>
    /// Gets the sector size of a disk.
    /// </summary>
    private static int GetSectorSize(string devicePath)
    {
        var output = RunProcess("blockdev", $"--getss {devicePath}").Trim();
        return int.TryParse(output, out var size) ? size : 512;
    }

    /// <summary>
    /// Parses a filesystem type string to the FileSystemType enum.
    /// </summary>
    private static FileSystemType ParseFilesystemType(string fsType)
    {
        return fsType.ToLower() switch
        {
            "ntfs" => FileSystemType.NTFS,
            "vfat" or "fat" => FileSystemType.FAT32, // Could be FAT12/16/32
            "exfat" => FileSystemType.ExFAT,
            "hfsplus" or "hfs+" => FileSystemType.HFSPlus,
            "apfs" => FileSystemType.APFS,
            "ext2" => FileSystemType.Ext2,
            "ext3" => FileSystemType.Ext3,
            "ext4" => FileSystemType.Ext4,
            "xfs" => FileSystemType.XFS,
            "btrfs" => FileSystemType.Btrfs,
            "zfs" => FileSystemType.ZFS,
            "refs" => FileSystemType.ReFS,
            _ => FileSystemType.Unknown
        };
    }

    #endregion
}
