// Copyright (c) StellaOps. All rights reserved.
// Licensed under AGPL-3.0-or-later. See LICENSE in the project root.

using System.Diagnostics;
using System.Globalization;
using System.Runtime.InteropServices;
using System.Text;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Options;

namespace StellaOps.BinaryIndex.Ghidra;

/// <summary>
/// Manages Ghidra Headless process lifecycle.
/// Provides methods to run analysis with proper process isolation and cleanup.
/// </summary>
public sealed class GhidraHeadlessManager : IAsyncDisposable
{
    private readonly GhidraOptions _options;
    private readonly ILogger<GhidraHeadlessManager> _logger;
    private readonly SemaphoreSlim _semaphore;
    private bool _disposed;

    /// <summary>
    /// Creates a new GhidraHeadlessManager.
    /// </summary>
    /// <param name="options">Ghidra configuration options.</param>
    /// <param name="logger">Logger instance.</param>
    public GhidraHeadlessManager(
        IOptions<GhidraOptions> options,
        ILogger<GhidraHeadlessManager> logger)
    {
        _options = options.Value;
        _logger = logger;
        _semaphore = new SemaphoreSlim(_options.MaxConcurrentInstances, _options.MaxConcurrentInstances);

        EnsureWorkDirectoryExists();
    }

    /// <summary>
    /// Runs Ghidra analysis on a binary.
    /// </summary>
    /// <param name="binaryPath">Absolute path to the binary file.</param>
    /// <param name="scriptName">Name of the post-analysis script to run.</param>
    /// <param name="scriptArgs">Arguments to pass to the script.</param>
    /// <param name="runAnalysis">Whether to run full auto-analysis.</param>
    /// <param name="timeoutSeconds">Timeout in seconds (0 = use default).</param>
    /// <param name="ct">Cancellation token.</param>
    /// <returns>Standard output from Ghidra.</returns>
    public async Task<GhidraProcessResult> RunAnalysisAsync(
        string binaryPath,
        string? scriptName = null,
        string[]? scriptArgs = null,
        bool runAnalysis = true,
        int timeoutSeconds = 0,
        CancellationToken ct = default)
    {
        ObjectDisposedException.ThrowIf(_disposed, this);

        if (!File.Exists(binaryPath))
        {
            throw new FileNotFoundException("Binary file not found", binaryPath);
        }

        var effectiveTimeout = timeoutSeconds > 0 ? timeoutSeconds : _options.DefaultTimeoutSeconds;

        await _semaphore.WaitAsync(ct);
        try
        {
            var projectDir = CreateTempProjectDirectory();
            try
            {
                var args = BuildAnalyzeArgs(projectDir, binaryPath, scriptName, scriptArgs, runAnalysis);
                return await RunGhidraAsync(args, effectiveTimeout, ct);
            }
            finally
            {
                if (_options.CleanupTempProjects)
                {
                    CleanupProjectDirectory(projectDir);
                }
            }
        }
        finally
        {
            _semaphore.Release();
        }
    }

    /// <summary>
    /// Runs a Ghidra script on an existing project.
    /// </summary>
    /// <param name="projectDir">Path to the Ghidra project directory.</param>
    /// <param name="projectName">Name of the Ghidra project.</param>
    /// <param name="scriptName">Name of the script to run.</param>
    /// <param name="scriptArgs">Arguments to pass to the script.</param>
    /// <param name="timeoutSeconds">Timeout in seconds (0 = use default).</param>
    /// <param name="ct">Cancellation token.</param>
    /// <returns>Standard output from Ghidra.</returns>
    public async Task<GhidraProcessResult> RunScriptAsync(
        string projectDir,
        string projectName,
        string scriptName,
        string[]? scriptArgs = null,
        int timeoutSeconds = 0,
        CancellationToken ct = default)
    {
        ObjectDisposedException.ThrowIf(_disposed, this);

        if (!Directory.Exists(projectDir))
        {
            throw new DirectoryNotFoundException($"Project directory not found: {projectDir}");
        }

        var effectiveTimeout = timeoutSeconds > 0 ? timeoutSeconds : _options.DefaultTimeoutSeconds;

        await _semaphore.WaitAsync(ct);
        try
        {
            var args = BuildScriptArgs(projectDir, projectName, scriptName, scriptArgs);
            return await RunGhidraAsync(args, effectiveTimeout, ct);
        }
        finally
        {
            _semaphore.Release();
        }
    }

    /// <summary>
    /// Checks if Ghidra is available and properly configured.
    /// </summary>
    /// <param name="ct">Cancellation token.</param>
    /// <returns>True if Ghidra is available.</returns>
    public async Task<bool> IsAvailableAsync(CancellationToken ct = default)
    {
        try
        {
            var executablePath = GetAnalyzeHeadlessPath();
            if (!File.Exists(executablePath))
            {
                _logger.LogDebug("Ghidra analyzeHeadless not found at: {Path}", executablePath);
                return false;
            }

            // Quick version check to verify Java is working
            var result = await RunGhidraAsync(["--help"], timeoutSeconds: 30, ct);
            return result.ExitCode == 0 || result.StandardOutput.Contains("analyzeHeadless", StringComparison.OrdinalIgnoreCase);
        }
        catch (Exception ex)
        {
            _logger.LogDebug(ex, "Ghidra availability check failed");
            return false;
        }
    }

    /// <summary>
    /// Gets Ghidra version information.
    /// </summary>
    /// <param name="ct">Cancellation token.</param>
    /// <returns>Version string.</returns>
    public async Task<string> GetVersionAsync(CancellationToken ct = default)
    {
        var result = await RunGhidraAsync(["--help"], timeoutSeconds: 30, ct);

        // Parse version from output - typically starts with "Ghidra X.Y"
        var lines = result.StandardOutput.Split('\n', StringSplitOptions.RemoveEmptyEntries);
        foreach (var line in lines)
        {
            if (line.Contains("Ghidra", StringComparison.OrdinalIgnoreCase) &&
                char.IsDigit(line.FirstOrDefault(c => char.IsDigit(c))))
            {
                return line.Trim();
            }
        }

        return "Unknown";
    }

    private string CreateTempProjectDirectory()
    {
        var projectDir = Path.Combine(
            _options.WorkDir,
            $"project_{DateTime.UtcNow:yyyyMMddHHmmssfff}_{Guid.NewGuid():N}");

        Directory.CreateDirectory(projectDir);
        _logger.LogDebug("Created temp project directory: {Path}", projectDir);
        return projectDir;
    }

    private void CleanupProjectDirectory(string projectDir)
    {
        try
        {
            if (Directory.Exists(projectDir))
            {
                Directory.Delete(projectDir, recursive: true);
                _logger.LogDebug("Cleaned up project directory: {Path}", projectDir);
            }
        }
        catch (Exception ex)
        {
            _logger.LogWarning(ex, "Failed to cleanup project directory: {Path}", projectDir);
        }
    }

    private void EnsureWorkDirectoryExists()
    {
        if (!Directory.Exists(_options.WorkDir))
        {
            Directory.CreateDirectory(_options.WorkDir);
            _logger.LogInformation("Created Ghidra work directory: {Path}", _options.WorkDir);
        }
    }

    private string[] BuildAnalyzeArgs(
        string projectDir,
        string binaryPath,
        string? scriptName,
        string[]? scriptArgs,
        bool runAnalysis)
    {
        var args = new List<string>
        {
            projectDir,
            "TempProject",
            "-import", binaryPath
        };

        if (!runAnalysis)
        {
            args.Add("-noanalysis");
        }

        if (!string.IsNullOrEmpty(scriptName))
        {
            args.AddRange(["-postScript", scriptName]);

            if (scriptArgs is { Length: > 0 })
            {
                args.AddRange(scriptArgs);
            }
        }

        if (!string.IsNullOrEmpty(_options.ScriptsDir))
        {
            args.AddRange(["-scriptPath", _options.ScriptsDir]);
        }

        args.AddRange(["-max-cpu", _options.MaxCpu.ToString(CultureInfo.InvariantCulture)]);

        return [.. args];
    }

    private static string[] BuildScriptArgs(
        string projectDir,
        string projectName,
        string scriptName,
        string[]? scriptArgs)
    {
        var args = new List<string>
        {
            projectDir,
            projectName,
            "-postScript", scriptName
        };

        if (scriptArgs is { Length: > 0 })
        {
            args.AddRange(scriptArgs);
        }

        return [.. args];
    }

    private async Task<GhidraProcessResult> RunGhidraAsync(
        string[] args,
        int timeoutSeconds,
        CancellationToken ct)
    {
        var executablePath = GetAnalyzeHeadlessPath();

        var startInfo = new ProcessStartInfo
        {
            FileName = executablePath,
            Arguments = string.Join(" ", args.Select(QuoteArg)),
            RedirectStandardOutput = true,
            RedirectStandardError = true,
            UseShellExecute = false,
            CreateNoWindow = true,
            StandardOutputEncoding = Encoding.UTF8,
            StandardErrorEncoding = Encoding.UTF8
        };

        ConfigureEnvironment(startInfo);

        _logger.LogDebug("Starting Ghidra: {Command} {Args}", executablePath, startInfo.Arguments);

        var stopwatch = Stopwatch.StartNew();
        using var process = new Process { StartInfo = startInfo };

        var stdoutBuilder = new StringBuilder();
        var stderrBuilder = new StringBuilder();

        process.OutputDataReceived += (_, e) =>
        {
            if (e.Data is not null)
            {
                stdoutBuilder.AppendLine(e.Data);
            }
        };

        process.ErrorDataReceived += (_, e) =>
        {
            if (e.Data is not null)
            {
                stderrBuilder.AppendLine(e.Data);
            }
        };

        if (!process.Start())
        {
            throw new GhidraException("Failed to start Ghidra process");
        }

        process.BeginOutputReadLine();
        process.BeginErrorReadLine();

        using var timeoutCts = new CancellationTokenSource(TimeSpan.FromSeconds(timeoutSeconds));
        using var linkedCts = CancellationTokenSource.CreateLinkedTokenSource(ct, timeoutCts.Token);

        try
        {
            await process.WaitForExitAsync(linkedCts.Token);
        }
        catch (OperationCanceledException) when (timeoutCts.IsCancellationRequested)
        {
            try
            {
                process.Kill(entireProcessTree: true);
            }
            catch
            {
                // Best effort kill
            }

            throw new GhidraTimeoutException(timeoutSeconds);
        }

        stopwatch.Stop();

        var stdout = stdoutBuilder.ToString();
        var stderr = stderrBuilder.ToString();

        _logger.LogDebug(
            "Ghidra completed with exit code {ExitCode} in {Duration}ms",
            process.ExitCode,
            stopwatch.ElapsedMilliseconds);

        if (process.ExitCode != 0)
        {
            _logger.LogWarning("Ghidra failed: {Error}", stderr);
        }

        return new GhidraProcessResult(
            process.ExitCode,
            stdout,
            stderr,
            stopwatch.Elapsed);
    }

    private string GetAnalyzeHeadlessPath()
    {
        var basePath = Path.Combine(_options.GhidraHome, "support");

        if (RuntimeInformation.IsOSPlatform(OSPlatform.Windows))
        {
            return Path.Combine(basePath, "analyzeHeadless.bat");
        }

        return Path.Combine(basePath, "analyzeHeadless");
    }

    private void ConfigureEnvironment(ProcessStartInfo startInfo)
    {
        if (!string.IsNullOrEmpty(_options.JavaHome))
        {
            startInfo.EnvironmentVariables["JAVA_HOME"] = _options.JavaHome;
        }

        startInfo.EnvironmentVariables["MAXMEM"] = _options.MaxMemory;
        startInfo.EnvironmentVariables["GHIDRA_HOME"] = _options.GhidraHome;
    }

    private static string QuoteArg(string arg)
    {
        if (arg.Contains(' ', StringComparison.Ordinal) || arg.Contains('"', StringComparison.Ordinal))
        {
            return $"\"{arg.Replace("\"", "\\\"")}\"";
        }

        return arg;
    }

    /// <inheritdoc />
    public async ValueTask DisposeAsync()
    {
        if (_disposed)
        {
            return;
        }

        _disposed = true;

        // Wait for any in-flight operations to complete
        for (var i = 0; i < _options.MaxConcurrentInstances; i++)
        {
            await _semaphore.WaitAsync();
        }

        _semaphore.Dispose();
    }
}

/// <summary>
/// Result of a Ghidra process execution.
/// </summary>
/// <param name="ExitCode">Process exit code.</param>
/// <param name="StandardOutput">Standard output content.</param>
/// <param name="StandardError">Standard error content.</param>
/// <param name="Duration">Execution duration.</param>
public sealed record GhidraProcessResult(
    int ExitCode,
    string StandardOutput,
    string StandardError,
    TimeSpan Duration)
{
    /// <summary>
    /// Whether the process completed successfully (exit code 0).
    /// </summary>
    public bool IsSuccess => ExitCode == 0;
}