using System.Collections.Immutable;
using System.Text.RegularExpressions;
using Injectio.Attributes;
using NLog;
using StabilityMatrix.Core.Extensions;
using StabilityMatrix.Core.Helper;
using StabilityMatrix.Core.Helper.Cache;
using StabilityMatrix.Core.Helper.HardwareInfo;
using StabilityMatrix.Core.Models.FileInterfaces;
using StabilityMatrix.Core.Models.Progress;
using StabilityMatrix.Core.Processes;
using StabilityMatrix.Core.Python;
using StabilityMatrix.Core.Services;
using LogLevel = Microsoft.Extensions.Logging.LogLevel;

namespace StabilityMatrix.Core.Models.Packages;

[RegisterSingleton<BasePackage, AiToolkit>(Duplicate = DuplicateStrategy.Append)]
public class AiToolkit(
    IGithubApiCache githubApi,
    ISettingsManager settingsManager,
    IDownloadService downloadService,
    IPrerequisiteHelper prerequisiteHelper,
    IPyInstallationManager pyInstallationManager,
    IPipWheelService pipWheelService
)
    : BaseGitPackage(
        githubApi,
        settingsManager,
        downloadService,
        prerequisiteHelper,
        pyInstallationManager,
        pipWheelService
    )
{
    private AnsiProcess? npmProcess;
    private static readonly Logger Logger = LogManager.GetCurrentClassLogger();

    public override string Name => "ai-toolkit";
    public override string DisplayName { get; set; } = "AI-Toolkit";
    public override string Author => "ostris";
    public override string Blurb => "AI Toolkit is an all in one training suite for diffusion models";
    public override string LicenseType => "MIT";
    public override string LicenseUrl => "https://github.com/ostris/ai-toolkit/blob/main/LICENSE";
    public override string LaunchCommand => string.Empty;

    public override Uri PreviewImageUri => new("https://cdn.lykos.ai/sm/packages/aitoolkit/preview.webp");

    public override string OutputFolderName => "output";
    public override IEnumerable<TorchIndex> AvailableTorchIndices => [TorchIndex.Cuda];
    public override PackageDifficulty InstallerSortOrder => PackageDifficulty.Advanced;
    public override SharedFolderMethod RecommendedSharedFolderMethod => SharedFolderMethod.None;
    public override List<LaunchOptionDefinition> LaunchOptions => [];
    public override Dictionary<SharedOutputType, IReadOnlyList<string>>? SharedOutputFolders => [];
    public override string MainBranch => "main";
    public override bool IsCompatible => HardwareHelper.HasNvidiaGpu();

    public override TorchIndex GetRecommendedTorchVersion() => TorchIndex.Cuda;

    public override PackageType PackageType => PackageType.SdTraining;
    public override bool OfferInOneClickInstaller => false;
    public override bool ShouldIgnoreReleases => true;
    public override PyVersion RecommendedPythonVersion => Python.PyInstallationManager.Python_3_11_13;

    public override IEnumerable<PackagePrerequisite> Prerequisites =>
        base.Prerequisites.Concat([PackagePrerequisite.Node]);

    public override async Task InstallPackage(
        string installLocation,
        InstalledPackage installedPackage,
        InstallPackageOptions options,
        IProgress<ProgressReport>? progress = null,
        Action<ProcessOutput>? onConsoleOutput = null,
        CancellationToken cancellationToken = default
    )
    {
        progress?.Report(new ProgressReport(-1, "Setting up venv", isIndeterminate: true));
        await using var venvRunner = await SetupVenvPure(
                installLocation,
                pythonVersion: options.PythonOptions.PythonVersion
            )
            .ConfigureAwait(false);
        venvRunner.UpdateEnvironmentVariables(GetEnvVars);

        var isBlackwell =
            SettingsManager.Settings.PreferredGpu?.IsBlackwellGpu() ?? HardwareHelper.HasBlackwellGpu();

        var config = new PipInstallConfig
        {
            RequirementsFilePaths = ["requirements.txt"],
            TorchVersion = "==2.7.0",
            TorchvisionVersion = "==0.22.0",
            TorchaudioVersion = "==2.7.0",
            CudaIndex = isBlackwell ? "cu128" : "cu126",
            ExtraPipArgs = [Compat.IsWindows ? "triton-windows" : "triton"],
            UpgradePackages = true,
        };

        await StandardPipInstallProcessAsync(
                venvRunner,
                options,
                installedPackage,
                config,
                onConsoleOutput,
                progress,
                cancellationToken
            )
            .ConfigureAwait(false);

        progress?.Report(new ProgressReport(-1f, "Installing AI Toolkit UI...", isIndeterminate: true));

        var uiDirectory = new DirectoryPath(installLocation, "ui");
        var envVars = GetEnvVars(venvRunner.EnvironmentVariables);
        await PrerequisiteHelper
            .RunNpm("install", uiDirectory, progress?.AsProcessOutputHandler(), envVars)
            .ConfigureAwait(false);
        await PrerequisiteHelper
            .RunNpm("run update_db", uiDirectory, progress?.AsProcessOutputHandler(), envVars)
            .ConfigureAwait(false);
        await PrerequisiteHelper
            .RunNpm("run build", uiDirectory, progress?.AsProcessOutputHandler(), envVars)
            .ConfigureAwait(false);
    }

    public override async Task RunPackage(
        string installLocation,
        InstalledPackage installedPackage,
        RunPackageOptions options,
        Action<ProcessOutput>? onConsoleOutput = null,
        CancellationToken cancellationToken = default
    )
    {
        await SetupVenv(installLocation, pythonVersion: PyVersion.Parse(installedPackage.PythonVersion))
            .ConfigureAwait(false);
        VenvRunner.UpdateEnvironmentVariables(GetEnvVars);

        var uiDirectory = new DirectoryPath(installLocation, "ui");
        var envVars = GetEnvVars(VenvRunner.EnvironmentVariables);
        npmProcess = PrerequisiteHelper.RunNpmDetached(
            "run start",
            uiDirectory,
            HandleConsoleOutput,
            envVars
        );
        npmProcess.EnableRaisingEvents = true;
        if (Compat.IsWindows)
        {
            ProcessTracker.AttachExitHandlerJobToProcess(npmProcess);
        }

        return;

        void HandleConsoleOutput(ProcessOutput s)
        {
            onConsoleOutput?.Invoke(s);

            if (!s.Text.Contains("Local:  ", StringComparison.OrdinalIgnoreCase))
                return;

            var regex = new Regex(@"(https?:\/\/)([^:\s]+):(\d+)");
            var match = regex.Match(s.Text);
            if (match.Success)
            {
                WebUrl = match.Value;
            }
            OnStartupComplete(WebUrl);
        }
    }

    public override async Task WaitForShutdown()
    {
        if (npmProcess is { HasExited: false })
        {
            npmProcess.Kill(true);
            try
            {
                await npmProcess
                    .WaitForExitAsync(new CancellationTokenSource(5000).Token)
                    .ConfigureAwait(false);
            }
            catch (OperationCanceledException e)
            {
                Logger.Warn(e, "Timed out waiting for npm to exit");
                npmProcess.CancelStreamReaders();
            }
        }
        npmProcess = null;
    }

    private ImmutableDictionary<string, string> GetEnvVars(ImmutableDictionary<string, string> env)
    {
        // set SETUPTOOLS_USE_DISTUTILS=setuptools to avoid job errors
        env = env.SetItem("SETUPTOOLS_USE_DISTUTILS", "setuptools");

        var pathBuilder = new EnvPathBuilder();

        if (env.TryGetValue("PATH", out var value))
        {
            pathBuilder.AddPath(value);
        }

        pathBuilder.AddPath(
            Compat.IsWindows
                ? Environment.GetFolderPath(Environment.SpecialFolder.System)
                : Path.Combine(SettingsManager.LibraryDir, "Assets", "nodejs", "bin")
        );

        pathBuilder.AddPath(Path.Combine(SettingsManager.LibraryDir, "Assets", "nodejs"));

        return env.SetItem("PATH", pathBuilder.ToString());
    }
}
