journal/Journal.Core/Services/Vault/VaultStorageService.cs

using System.Diagnostics;
using System.IO.Compression;
using System.Globalization;
using System.Security.Cryptography;
using System.Text;

namespace Journal.Core.Services.Vault;

public class VaultStorageService(IVaultCryptoService crypto) : IVaultStorageService
{
    private readonly IVaultCryptoService _crypto = crypto;
    private readonly Dictionary<string, string> _monthFingerprintCache = new(StringComparer.Ordinal);
    private readonly object _vaultIoLock = new();

    public string GetMonthlyVaultFileName(DateTime date) => date.ToString("yyyy-MM") + ".vault";

    public bool LoadAllVaults(string password, string vaultDirectory, string dataDirectory)
    {
        EnsureRequiredArguments(password, vaultDirectory, dataDirectory);

        lock (_vaultIoLock)
        {
            _monthFingerprintCache.Clear();
            PrepareDataDirectory(dataDirectory);

            if (!Directory.Exists(vaultDirectory))
                return true;

            var vaultFiles = Directory.GetFiles(vaultDirectory, "*.vault")
                .OrderBy(Path.GetFileName, StringComparer.Ordinal)
                .ToArray();
            if (vaultFiles.Length == 0)
                return true;

            // Restore database vault files first
            RestoreDatabaseVaults(password, vaultDirectory, dataDirectory);

            var anyDecrypted = false;
            var anyVaultFiles = false;
            foreach (var vaultFile in vaultFiles)
            {
                var fileName = Path.GetFileName(vaultFile);
                if (string.Equals(fileName, "_init_vault.vault", StringComparison.OrdinalIgnoreCase))
                {
                    try
                    {
                        File.Delete(vaultFile);
                    }
                    catch (Exception ex)
                    {
                        Debug.WriteLine($"[VaultStorageService] Failed to delete legacy vault file {fileName}: {ex.Message}");
                    }
                    continue;
                }

                if (IsReservedVaultFile(fileName))
                    continue;

                anyVaultFiles = true;
                try
                {
                    var encrypted = File.ReadAllBytes(vaultFile);
                    var decryptedZip = _crypto.DecryptData(encrypted, password);
                    ExtractZipContent(decryptedZip, dataDirectory);
                    anyDecrypted = true;
                }
                catch (CryptographicException)
                {
                    Debug.WriteLine($"[VaultStorageService] Decryption failed for {fileName} (likely wrong password)");
                }
                catch (Exception ex)
                {
                    Debug.WriteLine($"[VaultStorageService] Failed to load vault {fileName}: {ex.GetType().Name} - {ex.Message}");
                }
            }

            if (!anyDecrypted && anyVaultFiles)
                return false;

            return true;
        }
    }

    public bool SaveCurrentMonthVault(string password, string vaultDirectory, string dataDirectory, DateTime now)
    {
        EnsureRequiredArguments(password, vaultDirectory, dataDirectory);

        lock (_vaultIoLock)
        {
            Directory.CreateDirectory(vaultDirectory);
            if (!Directory.Exists(dataDirectory))
                return false;

            var monthKey = now.ToString("yyyy-MM", CultureInfo.InvariantCulture);
            var filesInMonth = Directory.GetFiles(dataDirectory, "*.md")
                .Where(path => Path.GetFileNameWithoutExtension(path).StartsWith(monthKey, StringComparison.Ordinal))
                .OrderBy(Path.GetFileName, StringComparer.Ordinal)
                .ToList();

            var savedMonth = false;
            if (filesInMonth.Count > 0)
            {
                var currentFingerprint = ComputeMonthFingerprint(filesInMonth);
                if (!_monthFingerprintCache.TryGetValue(monthKey, out var cachedFingerprint) ||
                    !string.Equals(cachedFingerprint, currentFingerprint, StringComparison.Ordinal))
                {
                    SaveMonth(password, monthKey, filesInMonth, vaultDirectory);
                    savedMonth = true;
                }
            }

            // Also persist custom-named entries alongside the current month vault
            SaveCustomEntries(password, vaultDirectory, dataDirectory);

            return savedMonth;
        }
    }

    public void RebuildAllVaults(string password, string vaultDirectory, string dataDirectory)
    {
        EnsureRequiredArguments(password, vaultDirectory, dataDirectory);

        lock (_vaultIoLock)
        {
            Directory.CreateDirectory(vaultDirectory);
            if (!Directory.Exists(dataDirectory))
                return;

            var monthlyFiles = new Dictionary<string, List<string>>(StringComparer.Ordinal);
            foreach (var filePath in Directory.GetFiles(dataDirectory, "*.md"))
            {
                var stem = Path.GetFileNameWithoutExtension(filePath);
                if (!DateTime.TryParseExact(stem, "yyyy-MM-dd", CultureInfo.InvariantCulture, DateTimeStyles.None, out var fileDate))
                    continue;

                var monthKey = fileDate.ToString("yyyy-MM", CultureInfo.InvariantCulture);
                if (!monthlyFiles.TryGetValue(monthKey, out var files))
                {
                    files = [];
                    monthlyFiles[monthKey] = files;
                }

                files.Add(filePath);
            }

            foreach (var (monthKey, filesInMonth) in monthlyFiles.OrderBy(static pair => pair.Key, StringComparer.Ordinal))
                SaveMonth(password, monthKey, filesInMonth, vaultDirectory);

            SaveCustomEntries(password, vaultDirectory, dataDirectory);
            SaveDatabaseVaults(password, vaultDirectory, dataDirectory);
        }
    }

    public void ClearDataDirectory(string dataDirectory)
    {
        if (string.IsNullOrWhiteSpace(dataDirectory))
            throw new ArgumentException("Data directory is required.", nameof(dataDirectory));

        lock (_vaultIoLock)
        {
            PrepareDataDirectory(dataDirectory);
            _monthFingerprintCache.Clear();
        }
    }

    private static void PrepareDataDirectory(string dataDirectory)
    {
        DeleteDirectoryWithRetries(dataDirectory);
        Directory.CreateDirectory(dataDirectory);
    }

    private static void DeleteDirectoryWithRetries(string dataDirectory, int retries = 5, int delayMs = 200)
    {
        if (!Directory.Exists(dataDirectory))
            return;

        for (var attempt = 0; attempt < retries; attempt++)
        {
            try
            {
                Directory.Delete(dataDirectory, recursive: true);
                return;
            }
            catch (IOException) when (attempt < retries - 1)
            {
                Thread.Sleep(delayMs);
            }
            catch (UnauthorizedAccessException) when (attempt < retries - 1)
            {
                Thread.Sleep(delayMs);
            }
        }

        Directory.Delete(dataDirectory, recursive: true);
    }

    // ── Custom entries vault helpers ──────────────────────────────

    private const string CustomEntriesVaultFileName = "_custom_entries.vault";

    private List<string> GetCustomEntryFiles(string dataDirectory)
    {
        return Directory.GetFiles(dataDirectory, "*.md")
            .Where(path =>
            {
                var stem = Path.GetFileNameWithoutExtension(path);
                return !DateTime.TryParseExact(stem, "yyyy-MM-dd", CultureInfo.InvariantCulture, DateTimeStyles.None, out _);
            })
            .OrderBy(Path.GetFileName, StringComparer.Ordinal)
            .ToList();
    }

    private void SaveCustomEntries(string password, string vaultDirectory, string dataDirectory)
    {
        var customFiles = GetCustomEntryFiles(dataDirectory);
        var vaultPath = Path.Combine(vaultDirectory, CustomEntriesVaultFileName);

        if (customFiles.Count == 0)
        {
            // Remove stale custom vault if no custom entries remain
            if (File.Exists(vaultPath))
                File.Delete(vaultPath);
            return;
        }

        var zipBytes = CreateMonthlyArchive(customFiles);
        var encrypted = _crypto.EncryptData(zipBytes, password);
        File.WriteAllBytes(vaultPath, encrypted);
    }

    private static bool IsCustomEntriesVaultFile(string fileName)
        => string.Equals(fileName, CustomEntriesVaultFileName, StringComparison.OrdinalIgnoreCase);

    // ── Database vault helpers ─────────────────────────────────────

    private const string DatabaseVaultPrefix = "_db_";
    private const string DatabaseVaultSuffix = ".vault";

    private static bool IsReservedVaultFile(string fileName)
        => fileName.StartsWith(DatabaseVaultPrefix, StringComparison.OrdinalIgnoreCase);

    private void SaveDatabaseVaults(string password, string vaultDirectory, string dataDirectory)
    {
        var dbFiles = Directory.GetFiles(dataDirectory, "*.db");
        foreach (var dbPath in dbFiles)
        {
            try
            {
                var dbFileName = Path.GetFileName(dbPath);
                var vaultFileName = $"{DatabaseVaultPrefix}{dbFileName}{DatabaseVaultSuffix}";
                var vaultPath = Path.Combine(vaultDirectory, vaultFileName);

                var dbBytes = File.ReadAllBytes(dbPath);
                var encrypted = _crypto.EncryptData(dbBytes, password);
                File.WriteAllBytes(vaultPath, encrypted);

                Debug.WriteLine($"[VaultStorageService] Saved database vault: {vaultFileName}");
            }
            catch (Exception ex)
            {
                Debug.WriteLine($"[VaultStorageService] Failed to save database vault for {Path.GetFileName(dbPath)}: {ex.Message}");
            }
        }
    }

    private void RestoreDatabaseVaults(string password, string vaultDirectory, string dataDirectory)
    {
        var dbVaultFiles = Directory.GetFiles(vaultDirectory, $"{DatabaseVaultPrefix}*{DatabaseVaultSuffix}");
        foreach (var vaultFile in dbVaultFiles)
        {
            try
            {
                var vaultFileName = Path.GetFileName(vaultFile);
                // Strip prefix and suffix to get original DB filename
                var dbFileName = vaultFileName[DatabaseVaultPrefix.Length..^DatabaseVaultSuffix.Length];
                if (string.IsNullOrWhiteSpace(dbFileName))
                    continue;

                var encrypted = File.ReadAllBytes(vaultFile);
                var dbBytes = _crypto.DecryptData(encrypted, password);
                var targetPath = Path.Combine(dataDirectory, dbFileName);
                File.WriteAllBytes(targetPath, dbBytes);

                Debug.WriteLine($"[VaultStorageService] Restored database from vault: {vaultFileName} → {dbFileName}");
            }
            catch (CryptographicException)
            {
                Debug.WriteLine($"[VaultStorageService] Database vault decryption failed for {Path.GetFileName(vaultFile)} (likely wrong password)");
            }
            catch (Exception ex)
            {
                Debug.WriteLine($"[VaultStorageService] Failed to restore database vault {Path.GetFileName(vaultFile)}: {ex.Message}");
            }
        }
    }

    private static void EnsureRequiredArguments(string password, string vaultDirectory, string dataDirectory)
    {
        if (string.IsNullOrWhiteSpace(password))
            throw new ArgumentException("Password cannot be empty.", nameof(password));
        if (string.IsNullOrWhiteSpace(vaultDirectory))
            throw new ArgumentException("Vault directory is required.", nameof(vaultDirectory));
        if (string.IsNullOrWhiteSpace(dataDirectory))
            throw new ArgumentException("Data directory is required.", nameof(dataDirectory));
    }

    private void SaveMonth(string password, string monthKey, List<string> filesInMonth, string vaultDirectory)
    {
        var monthDate = DateTime.ParseExact(monthKey, "yyyy-MM", CultureInfo.InvariantCulture);
        var monthlyVaultPath = Path.Combine(vaultDirectory, GetMonthlyVaultFileName(monthDate));

        var zipBytes = CreateMonthlyArchive(filesInMonth);
        var encryptedPayload = _crypto.EncryptData(zipBytes, password);
        File.WriteAllBytes(monthlyVaultPath, encryptedPayload);

        _monthFingerprintCache[monthKey] = ComputeMonthFingerprint(filesInMonth);
    }

    private static byte[] CreateMonthlyArchive(List<string> filesInMonth)
    {
        using var memoryStream = new MemoryStream();
        using (var archive = new ZipArchive(memoryStream, ZipArchiveMode.Create, leaveOpen: true))
        {
            foreach (var filePath in filesInMonth.OrderBy(Path.GetFileName, StringComparer.Ordinal))
            {
                var fileName = Path.GetFileName(filePath);
                if (string.IsNullOrWhiteSpace(fileName))
                    continue;

                var entry = archive.CreateEntry(fileName, CompressionLevel.Optimal);
                using var entryStream = entry.Open();
                using var sourceStream = File.OpenRead(filePath);
                sourceStream.CopyTo(entryStream);
            }
        }

        return memoryStream.ToArray();
    }

    private static string ComputeMonthFingerprint(List<string> files)
    {
        using var hash = IncrementalHash.CreateHash(HashAlgorithmName.SHA256);

        foreach (var filePath in files.OrderBy(Path.GetFileName, StringComparer.Ordinal))
        {
            var fileInfo = new FileInfo(filePath);
            if (!fileInfo.Exists)
                continue;

            AppendUtf8(hash, fileInfo.Name);
            AppendAscii(hash, fileInfo.LastWriteTimeUtc.Ticks.ToString(CultureInfo.InvariantCulture));
            AppendAscii(hash, fileInfo.Length.ToString(CultureInfo.InvariantCulture));
        }

        return Convert.ToHexString(hash.GetHashAndReset()).ToLowerInvariant();
    }

    private static void AppendUtf8(IncrementalHash hash, string value) => hash.AppendData(Encoding.UTF8.GetBytes(value));
    private static void AppendAscii(IncrementalHash hash, string value) => hash.AppendData(Encoding.ASCII.GetBytes(value));

    private static void ExtractZipContent(byte[] zipBytes, string dataDirectory)
    {
        using var stream = new MemoryStream(zipBytes);
        using var archive = new ZipArchive(stream, ZipArchiveMode.Read);

        var dataRoot = Path.GetFullPath(dataDirectory);
        if (!dataRoot.EndsWith(Path.DirectorySeparatorChar))
            dataRoot += Path.DirectorySeparatorChar;

        foreach (var entry in archive.Entries)
        {
            if (string.IsNullOrEmpty(entry.Name))
                continue;

            var destinationPath = Path.GetFullPath(Path.Combine(dataDirectory, entry.FullName));
            if (!destinationPath.StartsWith(dataRoot, StringComparison.OrdinalIgnoreCase))
                throw new InvalidDataException("Zip entry path escapes target data directory.");

            var destinationDir = Path.GetDirectoryName(destinationPath);
            if (!string.IsNullOrWhiteSpace(destinationDir))
                Directory.CreateDirectory(destinationDir);

            entry.ExtractToFile(destinationPath, overwrite: true);
        }
    }
}