using System;
using System.Collections.Generic;
using System.Runtime.InteropServices;
using System.Text;
using System.Threading;
using MultiWheelC.TrajectoryPlanning.Utils;

namespace MultiWheelC.TrajectoryPlanning.EMPlanner;

public sealed class OsqpNativeSolver : IQpSolver
{
    public QpSolveResult Solve(
        QuadraticProgram problem,
        QpSolverSettings settings,
        IReadOnlyList<double> warmStart,
        CancellationToken cancellationToken)
    {
        if (cancellationToken.IsCancellationRequested)
            return CreateFailure(QpSolveStatus.Cancelled, "OSQP solve was cancelled before native setup.");
        if (problem == null)
            return CreateFailure(QpSolveStatus.InvalidProblem, "A quadratic program is required.");
        if (settings == null)
            return CreateFailure(QpSolveStatus.InvalidProblem, "Solver settings are required.");

        double[] copiedWarmStart;
        string warmStartDiagnostic;
        if (!TryCopyWarmStart(warmStart, problem.VariableCount, out copiedWarmStart, out warmStartDiagnostic))
            return CreateFailure(QpSolveStatus.InvalidProblem, warmStartDiagnostic);

        OsqpNativeApi api;
        OsqpNativeLoadResult loadResult;
        if (!OsqpNativeLoader.TryGetLoadedApi(out api, out loadResult))
            return CreateFailure(loadResult.Status, loadResult.Diagnostic, loadResult.Version);

        var pinnedArrays = new OsqpPinnedArrays();
        OsqpNativeSetupMemory setupMemory = null;
        IntPtr solver = IntPtr.Zero;
        try
        {
            IntPtr pValues = pinnedArrays.Pin(Copy(problem.UpperTriangularP.Values));
            IntPtr pRows = pinnedArrays.Pin(Copy(problem.UpperTriangularP.RowIndices));
            IntPtr pColumns = pinnedArrays.Pin(Copy(problem.UpperTriangularP.ColumnPointers));
            IntPtr q = pinnedArrays.Pin(Copy(problem.LinearCost));
            IntPtr aValues = pinnedArrays.Pin(Copy(problem.ConstraintMatrix.Values));
            IntPtr aRows = pinnedArrays.Pin(Copy(problem.ConstraintMatrix.RowIndices));
            IntPtr aColumns = pinnedArrays.Pin(Copy(problem.ConstraintMatrix.ColumnPointers));
            IntPtr lowerBounds = pinnedArrays.Pin(Copy(problem.LowerBounds));
            IntPtr upperBounds = pinnedArrays.Pin(Copy(problem.UpperBounds));
            IntPtr warmStartPointer = pinnedArrays.Pin(copiedWarmStart);

            setupMemory = OsqpNativeStructures.AllocateSetupMemory(api.SetDefaultSettings);
            Marshal.StructureToPtr(
                CreateCsc(problem.UpperTriangularP, pColumns, pRows, pValues),
                setupMemory.UpperTriangularP,
                false);
            Marshal.StructureToPtr(
                CreateCsc(problem.ConstraintMatrix, aColumns, aRows, aValues),
                setupMemory.Constraints,
                false);

            var nativeSettings = (OsqpSettings)Marshal.PtrToStructure(setupMemory.Settings, typeof(OsqpSettings));
            nativeSettings.Verbose = 0;
            nativeSettings.WarmStarting = settings.EnableWarmStart ? 1 : 0;
            nativeSettings.Polishing = settings.EnablePolishing ? 1 : 0;
            nativeSettings.MaximumIterations = settings.MaximumIterations;
            nativeSettings.AbsoluteTolerance = settings.AbsoluteTolerance;
            nativeSettings.RelativeTolerance = settings.RelativeTolerance;
            nativeSettings.TimeLimit = settings.TimeLimit.TotalSeconds;
            Marshal.StructureToPtr(nativeSettings, setupMemory.Settings, false);

            int setupStatus = api.Setup(
                out solver,
                setupMemory.UpperTriangularP,
                q,
                setupMemory.Constraints,
                lowerBounds,
                upperBounds,
                problem.ConstraintCount,
                problem.VariableCount,
                setupMemory.Settings);
            if (setupStatus != 0 || solver == IntPtr.Zero)
                return CreateFailure(QpSolveStatus.NativeError, "OSQP setup failed with code " + setupStatus + ".", "setup=" + setupStatus);

            if (settings.EnableWarmStart && copiedWarmStart != null)
            {
                int warmStartStatus = api.WarmStart(solver, warmStartPointer, IntPtr.Zero);
                if (warmStartStatus != 0)
                    return CreateFailure(QpSolveStatus.NativeError, "OSQP warm start failed with code " + warmStartStatus + ".", "warm_start=" + warmStartStatus);
            }

            int solveStatus = api.Solve(solver);
            if (solveStatus != 0)
                return CreateFailure(QpSolveStatus.NativeError, "OSQP solve failed with code " + solveStatus + ".", "solve=" + solveStatus);
            if (cancellationToken.IsCancellationRequested)
                return CreateFailure(QpSolveStatus.Cancelled, "OSQP solve was cancelled after native completion.");

            return ReadSolveResult(solver, problem.VariableCount);
        }
        catch (Exception exception)
        {
            return CreateFailure(QpSolveStatus.NativeError, "OSQP native solve failed: " + exception.Message);
        }
        finally
        {
            if (solver != IntPtr.Zero)
            {
                try
                {
                    api.Cleanup(solver);
                }
                catch
                {
                }
            }

            if (setupMemory != null)
                setupMemory.Dispose();
            pinnedArrays.Dispose();
        }
    }

    private static QpSolveResult ReadSolveResult(IntPtr solver, int variableCount)
    {
        var solverPrefix = (OsqpSolverPrefix)Marshal.PtrToStructure(solver, typeof(OsqpSolverPrefix));
        if (solverPrefix.Info == IntPtr.Zero)
            return CreateFailure(QpSolveStatus.NativeError, "OSQP solve returned no information block.");

        var info = (OsqpInfo)Marshal.PtrToStructure(solverPrefix.Info, typeof(OsqpInfo));
        string nativeStatus = ReadNativeStatus(info);
        QpSolveStatus status = OsqpStatusMapper.Map(info.StatusValue);
        double[] primal;
        string primalDiagnostic;
        if (!TryReadPrimal(solverPrefix.Solution, variableCount, out primal, out primalDiagnostic))
            return CreateFailure(QpSolveStatus.NativeError, primalDiagnostic, nativeStatus);

        string diagnostic = string.Empty;
        double objective = ToFiniteMetric(info.ObjectiveValue, "objective", ref diagnostic);
        double primalResidual = ToFiniteMetric(info.PrimalResidual, "primal residual", ref diagnostic);
        double dualResidual = ToFiniteMetric(info.DualResidual, "dual residual", ref diagnostic);
        return new QpSolveResult(
            status,
            primal,
            objective,
            primalResidual,
            dualResidual,
            Math.Max(0, info.Iterations),
            ToSolveTime(info.SolveTime),
            nativeStatus,
            diagnostic);
    }

    private static bool TryReadPrimal(IntPtr solutionPointer, int variableCount, out double[] primal, out string diagnostic)
    {
        primal = new double[0];
        diagnostic = string.Empty;
        if (solutionPointer == IntPtr.Zero)
        {
            diagnostic = "OSQP solve returned no solution block.";
            return false;
        }

        var solution = (OsqpSolution)Marshal.PtrToStructure(solutionPointer, typeof(OsqpSolution));
        if (variableCount == 0)
            return true;
        if (solution.Primal == IntPtr.Zero)
        {
            diagnostic = "OSQP solve returned no primal vector.";
            return false;
        }

        primal = new double[variableCount];
        Marshal.Copy(solution.Primal, primal, 0, primal.Length);
        for (int index = 0; index < primal.Length; index++)
        {
            if (!NumericGuard.IsFinite(primal[index]))
            {
                diagnostic = "OSQP solve returned a non-finite primal value.";
                return false;
            }
        }

        return true;
    }

    private static OsqpCscMatrix CreateCsc(SparseCscMatrix matrix, IntPtr columnPointers, IntPtr rowIndices, IntPtr values)
    {
        return new OsqpCscMatrix
        {
            RowCount = matrix.RowCount,
            ColumnCount = matrix.ColumnCount,
            ColumnPointers = columnPointers,
            RowIndices = rowIndices,
            Values = values,
            MaximumNonZeroCount = matrix.Values.Count,
            NonZeroCount = -1,
            OwnsData = 0,
        };
    }

    private static bool TryCopyWarmStart(
        IReadOnlyList<double> warmStart,
        int variableCount,
        out double[] copiedWarmStart,
        out string diagnostic)
    {
        copiedWarmStart = null;
        diagnostic = string.Empty;
        if (warmStart == null)
            return true;
        if (warmStart.Count != variableCount)
        {
            diagnostic = "OSQP warm start length must match the quadratic-program variable count.";
            return false;
        }

        copiedWarmStart = Copy(warmStart);
        for (int index = 0; index < copiedWarmStart.Length; index++)
        {
            if (!NumericGuard.IsFinite(copiedWarmStart[index]))
            {
                diagnostic = "OSQP warm start values must be finite.";
                return false;
            }
        }

        return true;
    }

    private static T[] Copy<T>(IReadOnlyList<T> values)
    {
        var copied = new T[values.Count];
        for (int index = 0; index < copied.Length; index++)
            copied[index] = values[index];
        return copied;
    }

    private static string ReadNativeStatus(OsqpInfo info)
    {
        byte[] bytes = info.Status ?? new byte[0];
        int length = 0;
        while (length < bytes.Length && bytes[length] != 0)
            length++;
        string nativeStatus = Encoding.ASCII.GetString(bytes, 0, length).Trim();
        return string.IsNullOrEmpty(nativeStatus) ? "status=" + info.StatusValue : nativeStatus;
    }

    private static double ToFiniteMetric(double value, string metricName, ref string diagnostic)
    {
        if (NumericGuard.IsFinite(value))
            return value;

        diagnostic = string.IsNullOrEmpty(diagnostic)
            ? "OSQP returned a non-finite " + metricName + "; the planner-neutral diagnostic was normalized to zero."
            : diagnostic;
        return 0d;
    }

    private static TimeSpan ToSolveTime(double seconds)
    {
        if (!NumericGuard.IsFinite(seconds) || seconds <= 0d)
            return TimeSpan.Zero;
        if (seconds >= TimeSpan.MaxValue.TotalSeconds)
            return TimeSpan.MaxValue;
        return TimeSpan.FromSeconds(seconds);
    }

    private static QpSolveResult CreateFailure(QpSolveStatus status, string diagnostic, string nativeStatus = "")
    {
        return new QpSolveResult(status, new double[0], 0d, 0d, 0d, 0, TimeSpan.Zero, nativeStatus, diagnostic);
    }

    private sealed class OsqpPinnedArrays : IDisposable
    {
        private readonly List<GCHandle> handles = new List<GCHandle>();

        public IntPtr Pin<T>(T[] values) where T : struct
        {
            if (values == null || values.Length == 0)
                return IntPtr.Zero;

            GCHandle handle = GCHandle.Alloc(values, GCHandleType.Pinned);
            handles.Add(handle);
            return handle.AddrOfPinnedObject();
        }

        public void Dispose()
        {
            for (int index = handles.Count - 1; index >= 0; index--)
            {
                if (handles[index].IsAllocated)
                    handles[index].Free();
            }
            handles.Clear();
        }
    }
}