using System.Diagnostics;
using WxAgent.Core;

namespace WxAgent.Windows;

public sealed record DatabaseScanResult(IReadOnlyList<AccountKeySet> Accounts, int ProcessCount, int CandidateCount, int DatabaseCount,
    int SaltMatchedBindings = 0, int VerificationAttempts = 0);

public static class WechatDatabaseScanner
{
    public static DatabaseScanResult Scan(string? xwechatFilesRoot, CancellationToken cancellationToken,
        IReadOnlySet<string>? verifiedKeys = null)
    {
        var accounts = WechatDatabaseDiscovery.FindAccountRoots(xwechatFilesRoot, cancellationToken);
        if (accounts.Count == 0)
        {
            throw new WxAgentException(WxAgentErrorCode.DataRootNotFound, "No xwechat_files account db_storage directory was found.");
        }

        var processes = Process.GetProcessesByName("Weixin").OrderByDescending(process => process.WorkingSet64).ToArray();
        if (processes.Length == 0)
        {
            throw new WxAgentException(WxAgentErrorCode.WechatNotRunning, "Weixin.exe is not running.");
        }

        var evidence = accounts.ToDictionary(account => account.Fingerprint, _ => new Dictionary<string, DatabaseKeyEvidence>(StringComparer.OrdinalIgnoreCase));
        var candidateCount = 0;
        var scannedProcessCount = 0;
        var saltMatchedBindings = 0;
        var verificationAttempts = 0;
        string? version = null;
        foreach (var process in processes)
        {
            cancellationToken.ThrowIfCancellationRequested();
            IReadOnlyList<WcdbKeyCandidate> candidates;
            try
            {
                candidates = ProcessMemoryScanner.Scan(process.Id, cancellationToken);
                scannedProcessCount++;
                version ??= SafeVersion(process);
            }
            catch (WxAgentException exception) when (exception.Code == WxAgentErrorCode.ProcessAccessDenied)
            {
                continue;
            }

            candidateCount += candidates.Count;
            foreach (var candidate in candidates)
            {
                cancellationToken.ThrowIfCancellationRequested();
                foreach (var account in accounts)
                {
                    cancellationToken.ThrowIfCancellationRequested();
                    foreach (var database in account.Databases)
                    {
                        cancellationToken.ThrowIfCancellationRequested();
                        if (verifiedKeys?.Contains(KeyCachePlanner.Key(account.Fingerprint, database.RelativePath)) == true
                            || evidence[account.Fingerprint].ContainsKey(database.RelativePath) ||
                            candidate.SaltHint is not null && !string.Equals(candidate.SaltHint, database.Salt, StringComparison.OrdinalIgnoreCase))
                        {
                            continue;
                        }
                        if (candidate.SaltHint is not null) saltMatchedBindings++;
                        verificationAttempts++;
                        if (!SqlCipherPageVerifier.VerifyHexKey(database.FirstPage, candidate.EncKey)) continue;

                        evidence[account.Fingerprint][database.RelativePath] = new DatabaseKeyEvidence(
                            database.RelativePath,
                            database.Salt,
                            candidate.EncKey,
                            process.Id,
                            DateTimeOffset.UtcNow,
                            KeyBindingConfidence.PageHmacVerified);
                    }
                }
            }
        }

        if (scannedProcessCount == 0)
        {
            throw new WxAgentException(WxAgentErrorCode.ProcessAccessDenied, "No Weixin process could be opened for read-only memory access.");
        }

        var results = accounts.Select(account => new AccountKeySet(
            account.Fingerprint,
            account.AccountRootPath,
            version,
            DateTimeOffset.UtcNow,
            evidence[account.Fingerprint].Values.OrderBy(item => item.RelativePath, StringComparer.OrdinalIgnoreCase).ToArray())).ToArray();

        return new DatabaseScanResult(results, processes.Length, candidateCount, accounts.Sum(account => account.Databases.Count),
            saltMatchedBindings, verificationAttempts);
    }

    private static string? SafeVersion(Process process)
    {
        try { return process.MainModule?.FileVersionInfo.FileVersion; } catch { return null; }
    }
}