using System;
using System.Runtime.InteropServices;

namespace MultiWheelC.TrajectoryPlanning.EMPlanner;

[StructLayout(LayoutKind.Sequential, Pack = 8)]
internal struct OsqpCscMatrix
{
    public int RowCount;
    public int ColumnCount;
    public IntPtr ColumnPointers;
    public IntPtr RowIndices;
    public IntPtr Values;
    public int MaximumNonZeroCount;
    public int NonZeroCount;
    public int OwnsData;
}

[StructLayout(LayoutKind.Sequential, Pack = 8)]
internal struct OsqpSettings
{
    public int Device;
    public int LinearSystemSolver;
    public int AllocateSolution;
    public int Verbose;
    public int ProfilerLevel;
    public int WarmStarting;
    public int Scaling;
    public int Polishing;
    public double Rho;
    public int RhoIsVector;
    public double Sigma;
    public double Alpha;
    public int ConjugateGradientMaximumIterations;
    public int ConjugateGradientToleranceReduction;
    public double ConjugateGradientToleranceFraction;
    public int ConjugateGradientPreconditioner;
    public int AdaptiveRho;
    public int AdaptiveRhoInterval;
    public double AdaptiveRhoFraction;
    public double AdaptiveRhoTolerance;
    public int MaximumIterations;
    public double AbsoluteTolerance;
    public double RelativeTolerance;
    public double PrimalInfeasibilityTolerance;
    public double DualInfeasibilityTolerance;
    public int ScaledTermination;
    public int CheckTermination;
    public int CheckDualityGap;
    public double TimeLimit;
    public double PolishingDelta;
    public int PolishingRefinementIterations;
}

[StructLayout(LayoutKind.Sequential, Pack = 8)]
internal struct OsqpInfo
{
    [MarshalAs(UnmanagedType.ByValArray, SizeConst = 32, ArraySubType = UnmanagedType.I1)]
    public byte[] Status;
    public int StatusValue;
    public int PolishingStatus;
    public double ObjectiveValue;
    public double DualObjectiveValue;
    public double PrimalResidual;
    public double DualResidual;
    public double DualityGap;
    public int Iterations;
    public int RhoUpdates;
    public double RhoEstimate;
    public double SetupTime;
    public double SolveTime;
    public double UpdateTime;
    public double PolishingTime;
    public double RunTime;
    public double PrimalDualIntegral;
    public double RelativeKktError;
}

[StructLayout(LayoutKind.Sequential, Pack = 8)]
internal struct OsqpSolution
{
    public IntPtr Primal;
    public IntPtr Dual;
    public IntPtr PrimalInfeasibilityCertificate;
    public IntPtr DualInfeasibilityCertificate;
}

[StructLayout(LayoutKind.Sequential, Pack = 8)]
internal struct OsqpSolverPrefix
{
    public IntPtr Settings;
    public IntPtr Solution;
    public IntPtr Info;
    public IntPtr Workspace;
}

internal static class OsqpNativeStructures
{
    public static void ValidatePinnedLayout()
    {
        AssertSize(typeof(OsqpCscMatrix), 48);
        AssertOffset(typeof(OsqpCscMatrix), nameof(OsqpCscMatrix.ColumnPointers), 8);
        AssertOffset(typeof(OsqpCscMatrix), nameof(OsqpCscMatrix.Values), 24);
        AssertOffset(typeof(OsqpCscMatrix), nameof(OsqpCscMatrix.MaximumNonZeroCount), 32);

        AssertSize(typeof(OsqpSettings), 192);
        AssertOffset(typeof(OsqpSettings), nameof(OsqpSettings.Rho), 32);
        AssertOffset(typeof(OsqpSettings), nameof(OsqpSettings.Sigma), 48);
        AssertOffset(typeof(OsqpSettings), nameof(OsqpSettings.ConjugateGradientToleranceFraction), 72);
        AssertOffset(typeof(OsqpSettings), nameof(OsqpSettings.AdaptiveRhoFraction), 96);
        AssertOffset(typeof(OsqpSettings), nameof(OsqpSettings.AbsoluteTolerance), 120);
        AssertOffset(typeof(OsqpSettings), nameof(OsqpSettings.TimeLimit), 168);
        AssertOffset(typeof(OsqpSettings), nameof(OsqpSettings.PolishingRefinementIterations), 184);

        AssertSize(typeof(OsqpInfo), 152);
        AssertOffset(typeof(OsqpInfo), nameof(OsqpInfo.StatusValue), 32);
        AssertOffset(typeof(OsqpInfo), nameof(OsqpInfo.ObjectiveValue), 40);
        AssertOffset(typeof(OsqpInfo), nameof(OsqpInfo.Iterations), 80);
        AssertOffset(typeof(OsqpInfo), nameof(OsqpInfo.SetupTime), 96);
        AssertOffset(typeof(OsqpInfo), nameof(OsqpInfo.RelativeKktError), 144);

        AssertSize(typeof(OsqpSolution), 32);
        AssertSize(typeof(OsqpSolverPrefix), 32);
        AssertOffset(typeof(OsqpSolverPrefix), nameof(OsqpSolverPrefix.Workspace), 24);
    }

    public static OsqpNativeSetupMemory AllocateSetupMemory(OsqpNativeMethods.SetDefaultSettingsDelegate setDefaultSettings)
    {
        if (setDefaultSettings == null)
            throw new ArgumentNullException(nameof(setDefaultSettings));

        return new OsqpNativeSetupMemory(setDefaultSettings);
    }

    private static void AssertSize(Type type, int expectedSize)
    {
        int actualSize = Marshal.SizeOf(type);
        if (actualSize != expectedSize)
            throw new InvalidOperationException(type.Name + " has native size " + actualSize + " instead of " + expectedSize + ".");
    }

    private static void AssertOffset(Type type, string fieldName, int expectedOffset)
    {
        int actualOffset = checked((int)Marshal.OffsetOf(type, fieldName).ToInt64());
        if (actualOffset != expectedOffset)
            throw new InvalidOperationException(type.Name + "." + fieldName + " has native offset " + actualOffset + " instead of " + expectedOffset + ".");
    }
}

internal sealed class OsqpNativeSetupMemory : IDisposable
{
    private IntPtr upperTriangularP;
    private IntPtr constraints;
    private IntPtr settings;

    public OsqpNativeSetupMemory(OsqpNativeMethods.SetDefaultSettingsDelegate setDefaultSettings)
    {
        try
        {
            upperTriangularP = Marshal.AllocHGlobal(Marshal.SizeOf(typeof(OsqpCscMatrix)));
            constraints = Marshal.AllocHGlobal(Marshal.SizeOf(typeof(OsqpCscMatrix)));
            settings = Marshal.AllocHGlobal(Marshal.SizeOf(typeof(OsqpSettings)));
            Zero(upperTriangularP, Marshal.SizeOf(typeof(OsqpCscMatrix)));
            Zero(constraints, Marshal.SizeOf(typeof(OsqpCscMatrix)));
            Zero(settings, Marshal.SizeOf(typeof(OsqpSettings)));
            setDefaultSettings(settings);
        }
        catch
        {
            Dispose();
            throw;
        }
    }

    public IntPtr UpperTriangularP { get { return upperTriangularP; } }

    public IntPtr Constraints { get { return constraints; } }

    public IntPtr Settings { get { return settings; } }

    public void Dispose()
    {
        if (settings != IntPtr.Zero)
        {
            Marshal.FreeHGlobal(settings);
            settings = IntPtr.Zero;
        }
        if (constraints != IntPtr.Zero)
        {
            Marshal.FreeHGlobal(constraints);
            constraints = IntPtr.Zero;
        }
        if (upperTriangularP != IntPtr.Zero)
        {
            Marshal.FreeHGlobal(upperTriangularP);
            upperTriangularP = IntPtr.Zero;
        }
    }

    private static void Zero(IntPtr memory, int byteCount)
    {
        for (int index = 0; index < byteCount; index++)
            Marshal.WriteByte(memory, index, 0);
    }
}