using System;
using System.Runtime.InteropServices;
using System.Text;
using CeresSharp.Native;
using CeresSharp.Native.SafeHandles;
using CeresSharp;

namespace CeresSharp.Advanced;

/// <summary>
/// Gradient checker for validating cost function gradients.
/// </summary>
public sealed class GradientChecker : IDisposable
{
    private readonly GradientCheckerHandle _handle;
    private bool _disposed;

    /// <summary>
    /// Creates a new GradientChecker instance.
    /// </summary>
    /// <param name="costFunction">The cost function to check.</param>
    /// <param name="manifolds">Array of manifolds (can be null if no manifolds).</param>
    /// <param name="options">The gradient checker options.</param>
    public GradientChecker(CostFunction costFunction, Manifold[]? manifolds, GradientCheckerOptions options)
    {
        if (costFunction == null)
            throw new ArgumentNullException(nameof(costFunction));
        if (options == null)
            throw new ArgumentNullException(nameof(options));

        unsafe
        {
            IntPtr manifoldsPtr = IntPtr.Zero;
            int numManifolds = 0;

            if (manifolds != null && manifolds.Length > 0)
            {
                var manifoldHandles = new IntPtr[manifolds.Length];
                for (int i = 0; i < manifolds.Length; i++)
                {
                    if (manifolds[i] == null)
                        throw new ArgumentException($"Manifold {i} is null", nameof(manifolds));
                    manifoldHandles[i] = manifolds[i].Handle;
                }

                fixed (IntPtr* ptr = manifoldHandles)
                {
                    manifoldsPtr = (IntPtr)ptr;
                    numManifolds = manifolds.Length;
                }
            }

            var handle = CeresNative.ceres_wrapper_create_gradient_checker(
                costFunction.Handle,
                manifoldsPtr,
                numManifolds,
                options.Handle.DangerousGetHandle());

            _handle = GradientCheckerHandle.Create(handle);
        }
    }

    /// <summary>
    /// Gets the native handle.
    /// </summary>
    internal GradientCheckerHandle Handle => _handle;

    /// <summary>
    /// Probes gradients at the given parameters.
    /// </summary>
    /// <param name="parameters">Array of parameter block arrays.</param>
    /// <param name="relativePrecision">Relative precision for comparison.</param>
    /// <returns>True if gradients match, false otherwise.</returns>
    /// <exception cref="Exceptions.CeresException">Thrown if the operation fails (other than gradient mismatch).</exception>
    public bool Probe(double[][] parameters, double relativePrecision)
    {
        if (parameters == null || parameters.Length == 0)
            throw new ArgumentException("Parameters cannot be null or empty", nameof(parameters));

        unsafe
        {
            var parameterBlockPtrs = new IntPtr[parameters.Length];
            var pinnedArrays = new GCHandle[parameters.Length];

            try
            {
                for (int i = 0; i < parameters.Length; i++)
                {
                    if (parameters[i] == null)
                        throw new ArgumentException($"Parameter block {i} is null", nameof(parameters));

                    var handle = GCHandle.Alloc(parameters[i], GCHandleType.Pinned);
                    pinnedArrays[i] = handle;
                    parameterBlockPtrs[i] = handle.AddrOfPinnedObject();
                }

                var sb = new StringBuilder(512);
                fixed (IntPtr* ptrs = parameterBlockPtrs)
                {
                    var errorCode = CeresNative.ceres_wrapper_gradient_checker_probe(
                        _handle.DangerousGetHandle(),
                        (IntPtr)ptrs,
                        relativePrecision,
                        sb,
                        sb.Capacity);

                    // InvalidParameter means gradients don't match (expected behavior)
                    if (errorCode == Exceptions.CeresErrorCode.InvalidParameter)
                    {
                        return false;
                    }
                    
                    // Other error codes indicate actual errors
                    if (errorCode != Exceptions.CeresErrorCode.Success)
                    {
                        var message = sb.Length > 0 ? sb.ToString() : null;
                        throw new Exceptions.CeresException(errorCode, message);
                    }

                    return true;
                }
            }
            finally
            {
                foreach (var handle in pinnedArrays)
                {
                    if (handle.IsAllocated)
                        handle.Free();
                }
            }
        }
    }

    /// <summary>
    /// Probes gradients at the given parameters and returns error message if they don't match.
    /// </summary>
    /// <param name="parameters">Array of parameter block arrays.</param>
    /// <param name="relativePrecision">Relative precision for comparison.</param>
    /// <param name="errorMessage">Output error message if gradients don't match.</param>
    /// <returns>True if gradients match, false otherwise.</returns>
    /// <exception cref="Exceptions.CeresException">Thrown if the operation fails (other than gradient mismatch).</exception>
    public bool Probe(double[][] parameters, double relativePrecision, out string? errorMessage)
    {
        if (parameters == null || parameters.Length == 0)
            throw new ArgumentException("Parameters cannot be null or empty", nameof(parameters));

        unsafe
        {
            var parameterBlockPtrs = new IntPtr[parameters.Length];
            var pinnedArrays = new GCHandle[parameters.Length];

            try
            {
                for (int i = 0; i < parameters.Length; i++)
                {
                    if (parameters[i] == null)
                        throw new ArgumentException($"Parameter block {i} is null", nameof(parameters));

                    var handle = GCHandle.Alloc(parameters[i], GCHandleType.Pinned);
                    pinnedArrays[i] = handle;
                    parameterBlockPtrs[i] = handle.AddrOfPinnedObject();
                }

                var sb = new StringBuilder(512);
                fixed (IntPtr* ptrs = parameterBlockPtrs)
                {
                    var errorCode = CeresNative.ceres_wrapper_gradient_checker_probe(
                        _handle.DangerousGetHandle(),
                        (IntPtr)ptrs,
                        relativePrecision,
                        sb,
                        sb.Capacity);

                    // InvalidParameter means gradients don't match (expected behavior)
                    if (errorCode == Exceptions.CeresErrorCode.InvalidParameter)
                    {
                        errorMessage = sb.Length > 0 ? sb.ToString() : "Gradient check failed";
                        return false;
                    }
                    
                    // Other error codes indicate actual errors
                    if (errorCode != Exceptions.CeresErrorCode.Success)
                    {
                        var message = sb.Length > 0 ? sb.ToString() : null;
                        throw new Exceptions.CeresException(errorCode, message);
                    }

                    errorMessage = null;
                    return true;
                }
            }
            finally
            {
                foreach (var handle in pinnedArrays)
                {
                    if (handle.IsAllocated)
                        handle.Free();
                }
            }
        }
    }

    public void Dispose()
    {
        if (!_disposed)
        {
            _handle?.Dispose();
            _disposed = true;
        }
    }
}