using CartographerSharp.Mapping;
using CartographerSharp.Models.Mapping;
using CartographerSharp.Models.Transform;
using CartographerSharp.Transform;
using RobotNet10.RobotApp.SLAM.Cartographer;
using RobotNet10.Shared.Geometry;
using System.Reflection;

namespace RobotNet10.RobotApp.SLAM.Cartographer.Helpers;

/// <summary>
/// Extension methods for IMapBuilder
/// Provides trajectory management, loading, and serialization operations
/// </summary>
public static class MapBuilderExtensions
{
    /// <summary>
    /// Add trajectory builder to MapBuilder with sensor IDs and optional callback.
    /// </summary>
    /// <param name="gridTypeOverride">Override grid type for submap creation. Use ProbabilityGrid for ScanMapping, Tsdf for Localizing.</param>
    public static (int trajectoryId, ITrajectoryBuilder? trajectoryBuilder) AddTrajectoryBuilder(
        this IMapBuilder mapBuilder,
        CartographerConfiguration config,
        TrajectoryBuilderOptions? trajectoryOptionsOverride = null,
        GridOptions2D.GridType? gridTypeOverride = null)
    {
        ArgumentNullException.ThrowIfNull(mapBuilder);
        ArgumentNullException.ThrowIfNull(config);

        var sensorIds = TrajectoryHelper.BuildSensorIds(config);
        var trajectoryOptions = trajectoryOptionsOverride ?? CreateTrajectoryBuilderOptions(config, gridTypeOverride);

        var trajectoryId = mapBuilder.AddTrajectoryBuilder(sensorIds, trajectoryOptions);
        var trajectoryBuilder = mapBuilder.GetTrajectoryBuilder(trajectoryId);

        return (trajectoryId, trajectoryBuilder);
    }

    /// <summary>
    /// Wait for trajectory to finish with async support
    /// </summary>
    public static async Task WaitForTrajectoryFinishedAsync(
        this IMapBuilder mapBuilder,
        int trajectoryId,
        CancellationToken cancellationToken)
    {
        ArgumentNullException.ThrowIfNull(mapBuilder);

        var checkInterval = TimeSpan.FromMilliseconds(200);

        while (true)
        {
            cancellationToken.ThrowIfCancellationRequested();

            if (mapBuilder.PoseGraph.IsTrajectoryFinished(trajectoryId))
                return;

            await Task.Delay(checkInterval, cancellationToken);
        }
    }

    /// <summary>
    /// Finish all active trajectories before final optimization
    /// </summary>
    public static async Task FinishAllActiveTrajectoriesAsync(
        this IMapBuilder mapBuilder,
        CancellationToken cancellationToken)
    {
        ArgumentNullException.ThrowIfNull(mapBuilder);

        var trajectoryStates = mapBuilder.PoseGraph.GetTrajectoryStates();
        var activeTrajectories = trajectoryStates
            .Where(kvp => kvp.Value == IPoseGraph.TrajectoryState.Active)
            .Select(kvp => kvp.Key)
            .ToList();

        if (activeTrajectories.Count == 0)
            return;

        foreach (var trajectoryId in activeTrajectories)
        {
            cancellationToken.ThrowIfCancellationRequested();
            mapBuilder.FinishTrajectory(trajectoryId);
            await mapBuilder.WaitForTrajectoryFinishedAsync(trajectoryId, cancellationToken);
        }
    }

    /// <summary>
    /// Find first frozen trajectory in a loaded map
    /// </summary>
    public static int? FindFrozenTrajectory(this IMapBuilder mapBuilder)
    {
        ArgumentNullException.ThrowIfNull(mapBuilder);

        var trajectoryStates = mapBuilder.PoseGraph.GetTrajectoryStates();
        if (trajectoryStates == null || trajectoryStates.Count == 0)
            return null;

        foreach (var kvp in trajectoryStates)
        {
            if (kvp.Value == IPoseGraph.TrajectoryState.Frozen)
            {
                var trajectoryId = kvp.Key;

                // Validate frozen trajectory has nodes
                var trajectoryNodePoses = mapBuilder.PoseGraph.GetTrajectoryNodePoses();
                if (trajectoryNodePoses != null && !trajectoryNodePoses.IsEmpty)
                {
                    var trajectoryNodes = trajectoryNodePoses.BeginOfTrajectory(trajectoryId).ToList();
                    if (trajectoryNodes.Count > 0)
                        return trajectoryId;
                }
            }
        }

        return null;
    }

    /// <summary>
    /// Check if MapBuilder has trajectories and dispose if so
    /// </summary>
    public static bool DisposeIfHasTrajectories(this IMapBuilder mapBuilder)
    {
        ArgumentNullException.ThrowIfNull(mapBuilder);

        var trajectoryStates = mapBuilder.PoseGraph.GetTrajectoryStates();
        if (trajectoryStates != null && trajectoryStates.Count > 0)
        {
            (mapBuilder as IDisposable)?.Dispose();
            return true;
        }

        return false;
    }

    /// <summary>
    /// Finish trajectory if active, then dispose MapBuilder
    /// </summary>
    public static void FinishTrajectoryAndDispose(this IMapBuilder mapBuilder, int trajectoryId)
    {
        ArgumentNullException.ThrowIfNull(mapBuilder);

        if (trajectoryId >= 0)
        {
            var poseGraph = mapBuilder.PoseGraph;
            if (!poseGraph.IsTrajectoryFinished(trajectoryId))
            {
                mapBuilder.FinishTrajectory(trajectoryId);
            }
        }

        (mapBuilder as IDisposable)?.Dispose();
    }

    /// <summary>
    /// Add localization trajectory builder with optional initial pose.
    /// </summary>
    /// <param name="gridTypeOverride">Override grid type for submap creation. Use Tsdf for Localizing.</param>
    public static (int trajectoryId, ITrajectoryBuilder? trajectoryBuilder, Rigid3d? initialPoseInMapFrame) AddLocalizationTrajectoryBuilder(
            this IMapBuilder mapBuilder,
            CartographerConfiguration config,
            Pose? initialPose,
            GridOptions2D.GridType? gridTypeOverride = null)
    {
        ArgumentNullException.ThrowIfNull(mapBuilder);
        ArgumentNullException.ThrowIfNull(config);

        var frozenTrajectoryId = mapBuilder.FindFrozenTrajectory();
        var trajectoryOptions = CreateTrajectoryBuilderOptions(config, gridTypeOverride);
        CartographerSharp.Transform.Rigid3d? initialPoseInMapFrame = null;

        if (frozenTrajectoryId.HasValue)
        {
            CartographerSharp.Models.Transform.Rigid3dProto relativePoseProto;

            if (initialPose.HasValue)
            {
                var poseInMapFrame = mapBuilder.PoseGraph.GetTransformToMap() * PoseConverter.ToRigid3d(initialPose.Value);
                initialPoseInMapFrame = poseInMapFrame;
                relativePoseProto = (CartographerSharp.Models.Transform.Rigid3dProto)poseInMapFrame;
            }
            else
            {
                relativePoseProto = new CartographerSharp.Models.Transform.Rigid3dProto(
                    new CartographerSharp.Models.Transform.Vector3d(0, 0, 0),
                    new CartographerSharp.Models.Transform.Quaterniond(0, 0, 0, 1));
            }

            var initialTrajectoryPose = new CartographerSharp.Models.Mapping.InitialTrajectoryPose(
                relativePoseProto,
                frozenTrajectoryId.Value,
                DateTimeOffset.UtcNow.ToUnixTimeSeconds() * 1000000000L);

            trajectoryOptions = new TrajectoryBuilderOptions(
                trajectoryBuilder2DOptions: trajectoryOptions.TrajectoryBuilder2DOptions,
                trajectoryBuilder3DOptions: trajectoryOptions.TrajectoryBuilder3DOptions,
                initialTrajectoryPose: initialTrajectoryPose,
                pureLocalizationTrimmer: trajectoryOptions.PureLocalizationTrimmer,
                collateFixedFrame: trajectoryOptions.CollateFixedFrame,
                collateLandmarks: trajectoryOptions.CollateLandmarks,
                poseGraphOdometryMotionFilter: trajectoryOptions.PoseGraphOdometryMotionFilter);
        }

        var (trajectoryId, trajectoryBuilder) = mapBuilder.AddTrajectoryBuilder(config, trajectoryOptions);
        return (trajectoryId, trajectoryBuilder, initialPoseInMapFrame);
    }

    /// <summary>
    /// Create trajectory builder options from config.
    /// </summary>
    /// <param name="gridTypeOverride">Override grid type. ProbabilityGrid for ScanMapping, Tsdf for Localizing.</param>
    public static TrajectoryBuilderOptions CreateTrajectoryBuilderOptions(
        CartographerConfiguration config,
        GridOptions2D.GridType? gridTypeOverride = null)
    {
        ArgumentNullException.ThrowIfNull(config);

        if (config.TrajectoryBuilder.Use2D)
        {
            return TrajectoryBuilderOptionsFactory.Create2D(
                config.TrajectoryBuilder,
                config.Sensors.Imu.Enabled,
                gridTypeOverride);
        }
        else
        {
            return TrajectoryBuilderOptionsFactory.Create3D(
                config.TrajectoryBuilder,
                config.Sensors.Imu.Enabled);
        }
    }
}