/*
 * Copyright 2016 The Cartographer Authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

using CartographerSharp.Common;
using CartographerSharp.Mapping.Internal.D3D.ScanMatching;
using CartographerSharp.Models.Mapping;
using CartographerSharp.Transform;
using RobotNet10.Shared.Numbers;
using Submap3D = CartographerSharp.Mapping.D3D.Submap3D;

namespace CartographerSharp.Mapping.Internal.Constraints;

/// <summary>
/// Result of constraint building for 3D.
/// </summary>
public record struct ConstraintBuilder3DResult(List<IPoseGraph.Constraint> Constraints);

/// <summary>
/// Callback for constraint building completion.
/// </summary>
public delegate void ConstraintBuilder3DCallback(ConstraintBuilder3DResult result);

/// <summary>
/// Builds constraints for the 3D pose graph by matching nodes against submaps.
/// Match C++ ConstraintBuilder3D (constraint_builder_3d.h/cc)
/// </summary>
public class ConstraintBuilder3D : IDisposable
{
    private bool _disposed;
    private readonly ConstraintBuilderOptions _options;
    private readonly object _mutex = new();
    private readonly Dictionary<SubmapId, FixedRatioSampler> _perSubmapSampler = [];

    // Scan matchers
    private readonly CeresScanMatcher3D? _ceresScanMatcher;

    // Match C++: num_started_nodes_, num_finished_nodes_ (constraint_builder_3d.h line 157-159)
    private int _numStartedNodes;
    private int _numFinishedNodes;

    // Match C++: thread_pool_ (constraint_builder_3d.cc line 63)
    private readonly Common.Threading.ThreadPoolInterface _threadPool;

    // Match C++: finish_node_task_, when_done_task_ (constraint_builder_3d.h line 161-163)
    private Common.Threading.Task _finishNodeTask;
    private Common.Threading.Task _whenDoneTask;

    /// <summary>
    /// Submap scan matcher structure.
    /// Match C++ SubmapScanMatcher (constraint_builder_3d.h line 117-123)
    /// </summary>
    private class SubmapScanMatcher
    {
        public Mapping.D3D.HybridGrid? HighResolutionHybridGrid { get; set; }
        public Mapping.D3D.HybridGrid? LowResolutionHybridGrid { get; set; }
        public Mapping.D3D.IntensityHybridGrid? HighResolutionIntensityHybridGrid { get; set; }
        public RealTimeCorrelativeScanMatcher3D? FastCorrelativeScanMatcher { get; set; }
        public WeakReference<Common.Threading.Task>? CreationTaskHandle { get; set; }
    }

    // Match C++: submap_scan_matchers_ (constraint_builder_3d.h line 171-172)
    private readonly Dictionary<SubmapId, SubmapScanMatcher> _submapScanMatchers = [];

    // Match C++: constraints_ deque (constraint_builder_3d.h line 168)
    private readonly List<IPoseGraph.Constraint?> _pendingConstraints = [];

    // Match C++: when_done_ callback (constraint_builder_3d.h line 150-151)
    private ConstraintBuilder3DCallback? _whenDoneCallback;

    // Match C++: Constructor accepts thread_pool (constraint_builder_3d.cc line 61-68)
    public ConstraintBuilder3D(ConstraintBuilderOptions options, Common.Threading.ThreadPoolInterface threadPool)
    {
        _options = options;
        _threadPool = threadPool;

        // Initialize Ceres scan matcher if options are provided
        if (options.CeresScanMatcherOptions3D != null)
        {
            _ceresScanMatcher = new CeresScanMatcher3D(options.CeresScanMatcherOptions3D.Value);
        }

        // Match C++ (constraint_builder_3d.cc line 66-67): Initialize task objects
        _finishNodeTask = new Common.Threading.Task();
        _whenDoneTask = new Common.Threading.Task();
    }

    /// <summary>
    /// Schedules exploring a new constraint between 'submap' identified by
    /// 'submap_id', and the point cloud for 'node_id'.
    /// Match C++ MaybeAddConstraint (constraint_builder_3d.cc line 79-114)
    /// </summary>
    public void MaybeAddConstraint(
        SubmapId submapId,
        NodeId nodeId,
        Submap3D submap,
        TrajectoryNode node,
        Rigid3d globalNodePose,
        Rigid3d globalSubmapPose)
    {
        ArgumentNullException.ThrowIfNull(submap);
        ArgumentNullException.ThrowIfNull(node);
        if (node.ConstantData == null)
            return;

        // Check distance threshold
        var distance = (globalNodePose.Translation - globalSubmapPose.Translation).Length();
        if (distance > _options.MaxConstraintDistance)
            return;

        // Check sampling ratio
        if (!GetOrCreateSampler(submapId).Pulse())
            return;

        // Get point cloud from node
        var pointCloud = node.ConstantData.HighResolutionPointCloud;
        if (pointCloud == null || pointCloud.Count == 0)
            return;

        // Match C++ (constraint_builder_3d.cc line 95-113)
        lock (_mutex)
        {
            if (_whenDoneCallback != null)
            {
                // LOG(WARNING): MaybeAddConstraint was called while WhenDone was scheduled
            }

            // Add placeholder for constraint result
            var constraintIndex = _pendingConstraints.Count;
            _pendingConstraints.Add(null);

            // Get or create scan matcher (may schedule async construction)
            var scanMatcher = DispatchScanMatcherConstruction(submapId, submap);
            if (scanMatcher == null)
                return;

            // Schedule constraint computation task
            var constraintTask = new Common.Threading.Task();
            constraintTask.SetWorkItem(() =>
            {
                ComputeConstraint(submapId, nodeId, false, node.ConstantData,
                    globalNodePose, globalSubmapPose, scanMatcher, constraintIndex);
            });

            // Add dependency on scan matcher construction (match C++ line 110)
            constraintTask.AddDependency(scanMatcher.CreationTaskHandle);

            var constraintTaskHandle = _threadPool.Schedule(constraintTask);

            // Add dependency to finish_node_task (match C++ line 113)
            _finishNodeTask.AddDependency(constraintTaskHandle);
        }
    }

    /// <summary>
    /// Schedules exploring a new global constraint (full submap matching).
    /// Match C++ MaybeAddGlobalConstraint (constraint_builder_3d.cc line 116-142)
    /// </summary>
    public void MaybeAddGlobalConstraint(
        SubmapId submapId,
        NodeId nodeId,
        Submap3D submap,
        TrajectoryNode node,
        Quaternion globalNodeRotation,
        Quaternion globalSubmapRotation)
    {
        ArgumentNullException.ThrowIfNull(submap);
        ArgumentNullException.ThrowIfNull(node);
        if (node.ConstantData == null)
            return;

        // Get point cloud from node
        var pointCloud = node.ConstantData.HighResolutionPointCloud;
        if (pointCloud == null || pointCloud.Count == 0)
            return;

        // Match C++ (constraint_builder_3d.cc line 121-141)
        lock (_mutex)
        {
            if (_whenDoneCallback != null)
            {
                // LOG(WARNING): MaybeAddGlobalConstraint was called while WhenDone was scheduled
            }

            var constraintIndex = _pendingConstraints.Count;
            _pendingConstraints.Add(null);

            var scanMatcher = DispatchScanMatcherConstruction(submapId, submap);
            if (scanMatcher == null)
                return;

            // Create poses with only rotation (yaw is ignored for global matching)
            var globalNodePose = Rigid3d.FromRotation(globalNodeRotation);
            var globalSubmapPose = Rigid3d.FromRotation(globalSubmapRotation);

            var constraintTask = new Common.Threading.Task();
            constraintTask.SetWorkItem(() =>
            {
                ComputeConstraint(submapId, nodeId, true, node.ConstantData,
                    globalNodePose, globalSubmapPose, scanMatcher, constraintIndex);
            });

            constraintTask.AddDependency(scanMatcher.CreationTaskHandle);
            var constraintTaskHandle = _threadPool.Schedule(constraintTask);
            _finishNodeTask.AddDependency(constraintTaskHandle);
        }
    }

    /// <summary>
    /// Must be called after all computations related to one node have been added.
    /// Match C++ NotifyEndOfNode (constraint_builder_3d.cc line 144-156)
    /// </summary>
    public void NotifyEndOfNode()
    {
        lock (_mutex)
        {
            // Set work item for finish_node_task to increment num_finished_nodes
            _finishNodeTask.SetWorkItem(() =>
            {
                lock (_mutex)
                {
                    _numFinishedNodes++;
                }
            });

            // Schedule finish_node_task
            var finishNodeTaskHandle = _threadPool.Schedule(_finishNodeTask);

            // Create new finish_node_task for next node
            _finishNodeTask = new Common.Threading.Task();

            // Add dependency to when_done_task
            _whenDoneTask.AddDependency(finishNodeTaskHandle);

            _numStartedNodes++;
        }
    }

    /// <summary>
    /// Registers the callback to be called with the results, after all
    /// computations triggered by MaybeAdd*Constraint have finished.
    /// Match C++ WhenDone (constraint_builder_3d.cc line 158-168)
    /// </summary>
    public void WhenDone(ConstraintBuilder3DCallback callback)
    {
        lock (_mutex)
        {
            if (_whenDoneCallback != null)
            {
                throw new InvalidOperationException("WhenDone() called while another WhenDone() was pending");
            }

            _whenDoneCallback = callback;

            // Set work item for when_done_task to run callback
            _whenDoneTask.SetWorkItem(RunWhenDoneCallback);

            // Schedule when_done_task (it will wait for all dependencies)
            _threadPool.Schedule(_whenDoneTask);

            // Create new when_done_task for next cycle
            _whenDoneTask = new Common.Threading.Task();
        }
    }

    /// <summary>
    /// Match C++ RunWhenDoneCallback (constraint_builder_3d.cc line 307-333)
    /// </summary>
    private void RunWhenDoneCallback()
    {
        List<IPoseGraph.Constraint> result = [];
        ConstraintBuilder3DCallback? callback;

        lock (_mutex)
        {
            if (_whenDoneCallback == null)
            {
                throw new InvalidOperationException("RunWhenDoneCallback called without callback set");
            }

            // Collect all non-null constraints
            foreach (var constraint in _pendingConstraints)
            {
                if (constraint != null)
                {
                    result.Add(constraint.Value);
                }
            }

            // Clear pending constraints
            _pendingConstraints.Clear();

            // Take callback and clear
            callback = _whenDoneCallback;
            _whenDoneCallback = null;
        }

        // Invoke callback outside lock
        callback(new ConstraintBuilder3DResult(result));
    }

    /// <summary>
    /// Returns the number of consecutive finished nodes.
    /// </summary>
    public int GetNumFinishedNodes()
    {
        lock (_mutex) return _numFinishedNodes;
    }

    /// <summary>
    /// Returns the number of started nodes.
    /// </summary>
    public int GetNumStartedNodes()
    {
        lock (_mutex) return _numStartedNodes;
    }

    /// <summary>
    /// Delete data related to 'submap_id'.
    /// </summary>
    public void DeleteScanMatcher(SubmapId submapId)
    {
        lock (_mutex)
        {
            _submapScanMatchers.Remove(submapId);
            _perSubmapSampler.Remove(submapId);
        }
    }

    /// <summary>
    /// Returns the computed constraints.
    /// </summary>
    public List<IPoseGraph.Constraint> GetConstraints()
    {
        lock (_mutex)
        {
            return _pendingConstraints.Where(c => c != null).Select(c => c!.Value).ToList();
        }
    }

    /// <summary>
    /// Clears all constraints.
    /// </summary>
    public void Clear()
    {
        lock (_mutex)
        {
            _pendingConstraints.Clear();
        }
    }

    private FixedRatioSampler GetOrCreateSampler(SubmapId submapId)
    {
        lock (_mutex)
        {
            if (!_perSubmapSampler.TryGetValue(submapId, out var sampler))
            {
                sampler = new FixedRatioSampler(_options.SamplingRatio);
                _perSubmapSampler[submapId] = sampler;
            }
            return sampler;
        }
    }

    /// <summary>
    /// Dispatches scan matcher construction for a submap.
    /// Match C++ DispatchScanMatcherConstruction (constraint_builder_3d.cc line 170-198)
    /// MUST be called with _mutex held.
    /// </summary>
    private SubmapScanMatcher? DispatchScanMatcherConstruction(SubmapId submapId, Submap3D submap)
    {
        // Check if scan matcher already exists
        if (_submapScanMatchers.TryGetValue(submapId, out var existingMatcher))
        {
            return existingMatcher;
        }

        // Create new scan matcher entry
        var scanMatcher = new SubmapScanMatcher
        {
            HighResolutionHybridGrid = submap.HighResolutionHybridGrid,
            LowResolutionHybridGrid = submap.LowResolutionHybridGrid,
            HighResolutionIntensityHybridGrid = submap.HighResolutionIntensityHybridGrid
        };

        if (scanMatcher.HighResolutionHybridGrid == null)
        {
            return null;
        }

        _submapScanMatchers[submapId] = scanMatcher;

        var fastOptions = _options.FastCorrelativeScanMatcherOptions3D ??
            new FastCorrelativeScanMatcherOptions3D();

        // Get rotational scan matcher histogram from submap if available
        double[]? histogram = null;
        if (submap is Mapping.D3D.Submap3D submap3D)
        {
            var histogramList = submap3D.RotationalScanMatcherHistogram;
            if (histogramList != null && histogramList.Count > 0)
            {
                histogram = histogramList.ToArray();
            }
        }

        // Capture values for closure
        var highResGrid = scanMatcher.HighResolutionHybridGrid;
        var lowResGrid = scanMatcher.LowResolutionHybridGrid;

        // Schedule async construction of FastCorrelativeScanMatcher
        var scanMatcherTask = new Common.Threading.Task();
        scanMatcherTask.SetWorkItem(() =>
        {
            var matcher = new RealTimeCorrelativeScanMatcher3D(
                highResGrid,
                lowResGrid,
                histogram,
                fastOptions);

            lock (_mutex)
            {
                scanMatcher.FastCorrelativeScanMatcher = matcher;
            }
        });

        scanMatcher.CreationTaskHandle = _threadPool.Schedule(scanMatcherTask);

        return scanMatcher;
    }

    /// <summary>
    /// Computes a constraint between a node and submap.
    /// Match C++ ComputeConstraint (constraint_builder_3d.cc line 200-305)
    /// </summary>
    private void ComputeConstraint(
        SubmapId submapId,
        NodeId nodeId,
        bool matchFullSubmap,
        TrajectoryNode.Data constantData,
        Rigid3d globalNodePose,
        Rigid3d globalSubmapPose,
        SubmapScanMatcher scanMatcher,
        int constraintIndex)
    {
        // Get the scan matcher (should be ready by now due to task dependency)
        RealTimeCorrelativeScanMatcher3D? fastMatcher;
        lock (_mutex)
        {
            if (scanMatcher.FastCorrelativeScanMatcher == null)
            {
                return; // Scan matcher not ready (shouldn't happen with proper dependencies)
            }
            fastMatcher = scanMatcher.FastCorrelativeScanMatcher;
        }

        if (scanMatcher.HighResolutionHybridGrid == null)
            return;

        var pointCloud = constantData.HighResolutionPointCloud;
        if (pointCloud == null || pointCloud.Count == 0)
            return;

        // Step 1: Fast correlative scan matching for initial estimate
        FastCorrelativeScanMatcher3DResult? matchResult;
        if (matchFullSubmap)
        {
            matchResult = fastMatcher.MatchFullSubmap(
                globalNodePose.Rotation,
                globalSubmapPose.Rotation,
                constantData,
                _options.GlobalLocalizationMinScore);
        }
        else
        {
            matchResult = fastMatcher.Match(
                globalNodePose,
                globalSubmapPose,
                constantData,
                _options.MinScore);
        }

        if (matchResult == null)
            return; // Score too low

        var poseEstimate = matchResult.Value.PoseEstimate;

        // Step 2: Refine with Ceres scan matcher
        if (_ceresScanMatcher != null)
        {
            var pointCloudsAndGrids = new List<PointCloudAndHybridGridsPointers>();

            // Add high resolution point cloud and grid
            if (scanMatcher.HighResolutionHybridGrid != null)
            {
                pointCloudsAndGrids.Add(new PointCloudAndHybridGridsPointers
                {
                    PointCloud = pointCloud,
                    HybridGrid = scanMatcher.HighResolutionHybridGrid,
                    IntensityHybridGrid = scanMatcher.HighResolutionIntensityHybridGrid
                });
            }

            // Add low resolution point cloud and grid if available
            if (scanMatcher.LowResolutionHybridGrid != null && constantData.LowResolutionPointCloud != null)
            {
                pointCloudsAndGrids.Add(new PointCloudAndHybridGridsPointers
                {
                    PointCloud = constantData.LowResolutionPointCloud,
                    HybridGrid = scanMatcher.LowResolutionHybridGrid,
                    IntensityHybridGrid = null
                });
            }

            if (pointCloudsAndGrids.Count > 0)
            {
                CeresSharp.SolverSummary? summary = null;
                try
                {
                    _ceresScanMatcher.Match(
                        poseEstimate.Translation,
                        poseEstimate,
                        pointCloudsAndGrids,
                        out poseEstimate,
                        out summary
                    );
                }
                finally
                {
                    summary?.Dispose();
                }
            }
        }

        // Step 3: Create constraint
        // CRITICAL FIX: Match C++ (constraint_builder_3d.cc line 303-304)
        // constraint_transform = ComputeSubmapPose(*submap).inverse() * pose_estimate
        // poseEstimate is in global frame, constraint must be relative to submap's local frame
        var constraintTransform = globalSubmapPose.Inverse() * poseEstimate;
        var constraint = new IPoseGraph.Constraint(
            submapId,
            nodeId,
            new IPoseGraph.Constraint.Pose(
                constraintTransform,
                _options.LoopClosureTranslationWeight,
                _options.LoopClosureRotationWeight
            ),
            IPoseGraph.Constraint.Tag.InterSubmap,
            matchResult.Value.Score,
            IPoseGraph.Constraint.State.Enabled
        );

        // Store constraint at the pre-allocated index
        lock (_mutex)
        {
            _pendingConstraints[constraintIndex] = constraint;
        }
    }

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