using CartographerSharp.IO;
using CartographerSharp.Mapping;
using CartographerSharp.Transform;
using RobotNet10.RobotApp.Shared;
using RobotNet10.Shared.Geometry;
using RobotNet10.Shared.Localization;
using System.Text.Json;
using RobotNet10.RobotApp.SLAM.Cartographer;

namespace RobotNet10.RobotApp.SLAM.Cartographer.Helpers;

/// <summary>
/// Handles map saving workflow including scan matching, optimization, and file I/O
/// </summary>
public class MapSaveProcessor
{
    #region Fields and Constructor

    private readonly CartographerConfiguration _config;
    private readonly ILogger<MapSaveProcessor> _logger;
    private readonly string _mapsDirectory;
    private static readonly JsonSerializerOptions SerializerOptions = new() { WriteIndented = true };

    public MapSaveProcessor(
        CartographerConfiguration config,
        ILogger<MapSaveProcessor> logger)
    {
        _config = config ?? throw new ArgumentNullException(nameof(config));
        _logger = logger ?? throw new ArgumentNullException(nameof(logger));

        // Get maps directory from configuration
        _mapsDirectory = Path.GetFullPath(_config.MapStorage.Directory);
        Directory.CreateDirectory(_mapsDirectory);
    }

    #endregion

    #region Save Workflow

    /// <summary>
    /// Execute full save map workflow with optimization and file I/O
    /// </summary>
    /// <param name="mapName">Name of the map to save</param>
    /// <param name="mapBuilder">MapBuilder containing map data</param>
    /// <param name="trajectoryId">Active trajectory ID to finish (-1 if none)</param>
    /// <param name="progressCallback">Progress callback (total, current, percent)</param>
    /// <param name="newOrigin">Optional new origin to transform the map before saving (e.g., wall-aligned pose)</param>
    /// <param name="cancellationToken">Cancellation token</param>
    public async Task<string> SaveMapAsync(
        string mapName,
        IMapBuilder mapBuilder,
        int trajectoryId,
        Func<int, int, int, Task> progressCallback,
        Pose? newOrigin,
        CancellationToken cancellationToken)
    {
        ArgumentNullException.ThrowIfNull(mapBuilder);
        ArgumentNullException.ThrowIfNull(mapName);

        try
        {
            _logger.LogInformation("MapSaveProcessor: Starting map save process for: {MapName}", mapName);

            // Wait for active scan matching to complete
            await WaitForScanMatchingAsync(cancellationToken);

            // Track total progress across all phases
            // Phase 1: Initial work queue drain (0-20%)
            // Phase 2: Finish trajectory (20-40%)
            // Phase 3: RunFinalOptimization - work queue + constraint builder (40-90%)
            // Phase 4: Final scan matching + file saving (90-100%)

            // Finish trajectory if active
            if (trajectoryId >= 0)
            {
                // Phase 1: Drain work queue before finishing trajectory (0-20%)
                _logger.LogDebug("MapSaveProcessor: Phase 1 - Draining work queue before finishing trajectory");
                await DrainWorkQueueWithProgressAsync(mapBuilder, progressCallback, 0, 20, cancellationToken);

                // Phase 2: Finish trajectory (20-40%)
                _logger.LogDebug("MapSaveProcessor: Phase 2 - Finishing trajectory ID: {TrajectoryId}", trajectoryId);
                await progressCallback(100, 20, 20);
                await FinishTrajectoryAsync(mapBuilder, trajectoryId, cancellationToken);
                await progressCallback(100, 40, 40);
            }
            else
            {
                // No trajectory to finish, skip to phase 3
                await progressCallback(100, 40, 40);
            }

            _logger.LogDebug("MapSaveProcessor: Finishing all active trajectories");
            // Finish all other active trajectories
            await mapBuilder.FinishAllActiveTrajectoriesAsync(cancellationToken);

            // Phase 3: Run final optimization (40-90%)
            // RunFinalOptimization internally calls WaitForAllComputations which:
            // - Drains work queue
            // - Waits for constraint builder to finish
            // - Runs optimization
            cancellationToken.ThrowIfCancellationRequested();
            _logger.LogDebug("MapSaveProcessor: Phase 3 - Running final pose graph optimization");

            // Start optimization in background task and track progress
            var optimizationTask = Task.Run(() => mapBuilder.PoseGraph.RunFinalOptimization(), cancellationToken);

            // Track progress while optimization is running
            await TrackOptimizationProgressAsync(mapBuilder, progressCallback, 40, 90, optimizationTask, cancellationToken);

            // Phase 4: Final scan matching + file saving (90-100%)
            _logger.LogDebug("MapSaveProcessor: Phase 4 - Final scan matching and file saving");
            await progressCallback(100, 90, 90);

            // Final wait for scan matching
            await WaitForScanMatchingAsync(cancellationToken);

            // Get map path
            var mapPath = GetMapPath(mapName);

            // Validate map before saving
            var (isValid, errorMessage) = ValidateMapBeforeSave(mapBuilder, mapName, mapPath);
            if (!isValid)
            {
                throw new InvalidOperationException($"Map validation failed: {errorMessage}");
            }

            _logger.LogDebug("MapSaveProcessor: Saving pbstream and metadata to: {MapPath}", mapPath);
            // Save pbstream file
            SavePbstream(mapBuilder, mapPath);

            // Transform origin if newOrigin is provided (e.g., wall alignment)
            IMapBuilder mapBuilderForMetadata = mapBuilder;
            if (newOrigin.HasValue && !PbstreamTransformHelper.IsIdentityPose(newOrigin.Value))
            {
                _logger.LogDebug("MapSaveProcessor: Transforming map origin with wall-aligned pose");
                var pbstreamPath = Path.Combine(mapPath, "map.pbstream");

                // For wall alignment, we need to apply the INVERSE of the compensation pose.
                // The compensation angle represents "how much the robot should rotate to align the wall",
                // but TransformToMap uses the inverse when rendering (via TransformToMapInverse).
                // So we invert the pose here to get the correct rotation direction in the final map.
                var poseRigid = PbstreamTransformHelper.PoseToRigid3d(newOrigin.Value);
                var invertedRigid = poseRigid.Inverse();
                var invertedPose = new Pose
                {
                    Position = new RobotNet10.Shared.Numbers.Vector3(
                        invertedRigid.Translation.X,
                        invertedRigid.Translation.Y,
                        invertedRigid.Translation.Z),
                    Orientation = new RobotNet10.Shared.Numbers.Quaternion(
                        invertedRigid.Rotation.X,
                        invertedRigid.Rotation.Y,
                        invertedRigid.Rotation.Z,
                        invertedRigid.Rotation.W)
                };

                PbstreamTransformHelper.TransformPbstreamOrigin(pbstreamPath, invertedPose, _logger);

                // Reload MapBuilder from transformed pbstream
                var loadResult = new MapCartographerLoadResult
                {
                    PbstreamPath = pbstreamPath,
                    MapPath = mapPath,
                    SavedConfig = null // Will use current config
                };
                mapBuilderForMetadata = MapBuilderHelper.CreateMapBuilderFromLoadResult(
                    loadResult, _config, _logger, loadFrozenState: false);
            }

            _logger.LogDebug("MapSaveProcessor: Generating occupancy grid and saving metadata");
            // Generate occupancy grid and save metadata
            var saveResult = await SaveMapMetadataAsync(mapName, mapBuilderForMetadata, mapPath, cancellationToken);

            // Dispose reloaded MapBuilder if it was created
            if (mapBuilderForMetadata != mapBuilder && mapBuilderForMetadata is IDisposable disposable)
            {
                disposable.Dispose();
            }

            // Final progress update (100%)
            await progressCallback(100, 100, 100);

            _logger.LogInformation("MapSaveProcessor: Map saved successfully: {MapName}", mapName);
            return saveResult;
        }
        catch (OperationCanceledException)
        {
            _logger.LogWarning("MapSaveProcessor: Map save cancelled for: {MapName}", mapName);
            throw;
        }
        catch (Exception ex)
        {
            _logger.LogError(ex, "MapSaveProcessor: Failed to save map: {MapName}", mapName);
            throw;
        }
    }

    #endregion

    #region Progress Tracking

    /// <summary>
    /// Wait for active scan matching operations to complete
    /// </summary>
    private static async Task WaitForScanMatchingAsync(CancellationToken cancellationToken, int checkIntervalMs = 100)
    {
        while (true)
        {
            cancellationToken.ThrowIfCancellationRequested();

            var activeCount = CartographerSharp.Mapping.Internal.D2D.ScanMatching.CeresScanMatcher2D.GetActiveScanMatchingCount();
            if (activeCount == 0)
                break;

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

    /// <summary>
    /// Finish active trajectory and wait for completion
    /// </summary>
    private static async Task FinishTrajectoryAsync(IMapBuilder mapBuilder, int trajectoryId, CancellationToken cancellationToken)
    {
        cancellationToken.ThrowIfCancellationRequested();
        
        // Wait for scan matching to complete before finishing trajectory
        // Wait for scan matching to complete
        await WaitForScanMatchingAsync(cancellationToken);

        // Finish trajectory
        cancellationToken.ThrowIfCancellationRequested();

        // FinishTrajectory is called synchronously
        mapBuilder.FinishTrajectory(trajectoryId);

        // Wait for trajectory to finish asynchronously
        // Wait for trajectory to finish
        await mapBuilder.WaitForTrajectoryFinishedAsync(trajectoryId, cancellationToken);

        // Trajectory finished successfully
    }

    /// <summary>
    /// Drain work queue with progress updates mapped to a percentage range.
    /// </summary>
    private static async Task DrainWorkQueueWithProgressAsync(
        IMapBuilder mapBuilder,
        Func<int, int, int, Task> progressCallback,
        int progressStart,
        int progressEnd,
        CancellationToken cancellationToken)
    {
        var poseGraph = mapBuilder.PoseGraph;
        var initialQueueCount = poseGraph.WorkQueueCount;

        if (initialQueueCount == 0)
        {
            await progressCallback(100, progressEnd, progressEnd);
            return;
        }

        var progressRange = progressEnd - progressStart;

        while (poseGraph.WorkQueueCount > 0)
        {
            cancellationToken.ThrowIfCancellationRequested();

            var currentQueueCount = poseGraph.WorkQueueCount;
            var processed = initialQueueCount - currentQueueCount;

            // Calculate progress within the range
            int percentComplete;
            if (initialQueueCount > 0)
            {
                var phaseProgress = (double)processed / initialQueueCount;
                percentComplete = progressStart + (int)(phaseProgress * progressRange);
            }
            else
            {
                percentComplete = progressEnd;
            }
            percentComplete = Math.Clamp(percentComplete, progressStart, progressEnd);

            await progressCallback(initialQueueCount, processed, percentComplete);

            await Task.Delay(500, cancellationToken);
        }

        await progressCallback(initialQueueCount, initialQueueCount, progressEnd);
    }

    /// <summary>
    /// Track optimization progress (work queue + constraint builder) while optimization task is running.
    /// </summary>
    private static async Task TrackOptimizationProgressAsync(
        IMapBuilder mapBuilder,
        Func<int, int, int, Task> progressCallback,
        int progressStart,
        int progressEnd,
        Task optimizationTask,
        CancellationToken cancellationToken)
    {
        var poseGraph = mapBuilder.PoseGraph;
        var progressRange = progressEnd - progressStart;

        // Capture the remaining work at the start of tracking.
        // Progress is calculated based on how much of this remaining work has been completed,
        // so the remaining portion at call time = 100%.
        var initialWorkQueueCount = poseGraph.WorkQueueCount;
        var initialConstraintTasksFinished = poseGraph.ConstraintTasksFinished;
        var initialConstraintTasksTotal = poseGraph.ConstraintTasksTotal;
        var remainingConstraintTasks = initialConstraintTasksTotal - initialConstraintTasksFinished;

        // Tracking optimization progress

        var lastLogTime = DateTime.UtcNow;

        while (!optimizationTask.IsCompleted)
        {
            cancellationToken.ThrowIfCancellationRequested();

            var workQueueCount = poseGraph.WorkQueueCount;
            var constraintTasksTotal = poseGraph.ConstraintTasksTotal;
            var constraintTasksFinished = poseGraph.ConstraintTasksFinished;

            // Recalculate remaining tasks as total may grow during processing
            var currentRemainingTotal = constraintTasksTotal - initialConstraintTasksFinished;
            var effectiveRemaining = Math.Max(remainingConstraintTasks, currentRemainingTotal);

            // Calculate progress based on remaining work from the start of tracking:
            // 1. Work queue drain progress (30% of phase)
            double workQueueProgress;
            if (initialWorkQueueCount > 0)
            {
                var workQueueProcessed = initialWorkQueueCount - workQueueCount;
                workQueueProgress = Math.Min(1.0, (double)workQueueProcessed / initialWorkQueueCount);
            }
            else
            {
                workQueueProgress = 1.0; // No work queue items at start
            }

            // 2. Constraint task progress (70% of phase)
            double constraintProgress;
            if (effectiveRemaining > 0)
            {
                var newlyFinished = constraintTasksFinished - initialConstraintTasksFinished;
                constraintProgress = Math.Min(1.0, (double)newlyFinished / effectiveRemaining);
            }
            else
            {
                constraintProgress = 1.0; // No constraint tasks to process
            }

            // Combined progress weighted by actual work amounts
            var totalWork = initialWorkQueueCount + effectiveRemaining;
            var combinedProgress = totalWork > 0
                ? (workQueueProgress * initialWorkQueueCount + constraintProgress * effectiveRemaining) / totalWork
                : 1.0;
            var percentComplete = progressStart + (int)(combinedProgress * progressRange);
            percentComplete = Math.Clamp(percentComplete, progressStart, progressEnd);

            await progressCallback(constraintTasksTotal, constraintTasksFinished, percentComplete);

            // Update last log time periodically (logging removed to reduce noise)
            if ((DateTime.UtcNow - lastLogTime).TotalSeconds >= 10)
            {
                lastLogTime = DateTime.UtcNow;
            }

            // Check more frequently for quicker updates
            await Task.Delay(1000, cancellationToken);
        }

        // Wait for optimization task to complete (may throw if cancelled)
        await optimizationTask;

        await progressCallback(poseGraph.ConstraintTasksTotal, poseGraph.ConstraintTasksFinished, progressEnd);
        // Optimization completed successfully
    }

    #endregion

    #region File I/O

    /// <summary>
    /// Save pbstream file from MapBuilder state
    /// </summary>
    private void SavePbstream(IMapBuilder mapBuilder, string mapPath)
    {
        Directory.CreateDirectory(mapPath);
        var pbstreamPath = Path.Combine(mapPath, "map.pbstream");
        var includeUnfinished = _config.MapStorage.IncludeUnfinishedSubmaps;

        try
        {
            using var writer = new ProtoStreamWriter(pbstreamPath);
            mapBuilder.SerializeState(includeUnfinishedSubmaps: includeUnfinished, writer);
            if (!writer.Close())
            {
                throw new InvalidOperationException($"ProtoStreamWriter.Close() failed for {pbstreamPath}");
            }

            _logger.LogInformation("MapSaveProcessor: Pbstream saved to: {Path}", pbstreamPath);
        }
        catch (Exception ex)
        {
            _logger.LogError(ex, "MapSaveProcessor: Failed to save pbstream to: {Path}", pbstreamPath);
            throw new InvalidOperationException($"Failed to save pbstream: {ex.Message}", ex);
        }
    }

    /// <summary>
    /// Generate occupancy grid and save metadata files (PGM, PNG, JSON).
    /// Can be used for both new map creation and map update (e.g., after transform).
    /// </summary>
    /// <param name="mapName">Map name for sanitization</param>
    /// <param name="mapBuilder">MapBuilder containing the map data</param>
    /// <param name="mapPath">Directory path to save files</param>
    /// <param name="cancellationToken">Cancellation token</param>
    /// <returns>The map path</returns>
    public async Task<string> SaveMapMetadataAsync(
        string mapName,
        IMapBuilder mapBuilder,
        string mapPath,
        CancellationToken cancellationToken = default)
    {
        cancellationToken.ThrowIfCancellationRequested();

        // Generate occupancy grid
        var (occupancyGrid, submapCount) = GenerateOccupancyGridForSave(mapBuilder, _config.MapStorage.OccupancyGridResolution);

        var poseGraph = mapBuilder.PoseGraph;
        var trajectoryNodePoses = poseGraph.GetTrajectoryNodePoses();
        var trajectoryNodeCount = trajectoryNodePoses?.Count ?? 0;

        var finalSubmapCount = submapCount;
        if (finalSubmapCount == 0)
        {
            var allSubmapData = poseGraph.GetAllSubmapData();
            finalSubmapCount = allSubmapData?.Count ?? 0;
        }

        // Save occupancy grid files if available
        if (occupancyGrid != null)
        {
            await OccupancyGridFileHelper.SaveAllFormatsAsync(occupancyGrid, mapPath, cancellationToken, _logger);
        }
        else
        {
            _logger.LogWarning("MapSaveProcessor: No occupancy grid generated, skipping PGM/PNG/JPG files");
        }

        // Calculate bounds and size (in meters)
        var bounds = new BoundingBox();
        var size = new MapSize();
        if (occupancyGrid != null)
        {
            var originX = occupancyGrid.Origin.Position.X;
            var originY = occupancyGrid.Origin.Position.Y;
            var resolution = occupancyGrid.Resolution;
            // Size in meters = pixels * resolution
            var widthMeters = occupancyGrid.Width * resolution;
            var heightMeters = occupancyGrid.Height * resolution;
            size = new MapSize(widthMeters, heightMeters);
            bounds = new BoundingBox
            {
                MinX = originX,
                MinY = originY,
                MaxX = originX + widthMeters,
                MaxY = originY + heightMeters
            };
        }

        // Create metadata
        var sanitizedMapName = MapNameHelper.Sanitize(mapName);
        var metadata = new MapCartographerInfo
        {
            Name = sanitizedMapName,
            FolderPath = mapPath,
            CreatedDate = DateTime.UtcNow,
            Resolution = occupancyGrid?.Resolution ?? _config.MapStorage.OccupancyGridResolution,
            Size = size,
            Origin = occupancyGrid?.Origin ?? new Pose
            {
                Position = new RobotNet10.Shared.Numbers.Vector3(0, 0, 0),
                Orientation = new RobotNet10.Shared.Numbers.Quaternion(0, 0, 0, 1)
            },
            Bounds = bounds,
            TrajectoryNodeCount = trajectoryNodeCount,
            MapConfig = MapBuilderHelper.CreateConfigSnapshot(_config)
        };

        // Save metadata JSON
        var metadataPath = Path.Combine(mapPath, "map.json");
        await File.WriteAllTextAsync(metadataPath, JsonSerializer.Serialize(metadata, SerializerOptions), cancellationToken);

        return mapPath;
    }

    #endregion

    #region Validation

    /// <summary>
    /// Get map path from map name
    /// </summary>
    private string GetMapPath(string mapName)
    {
        var sanitized = MapNameHelper.Sanitize(mapName);
        return Path.Combine(_mapsDirectory, sanitized);
    }

    /// <summary>
    /// Validate map before saving
    /// </summary>
    private (bool IsValid, string? ErrorMessage) ValidateMapBeforeSave(IMapBuilder mapBuilder, string mapName, string mapPath)
    {
        try
        {
            var poseGraph = mapBuilder.PoseGraph;
            var trajectoryStates = poseGraph.GetTrajectoryStates();
            var trajectoryNodePoses = poseGraph.GetTrajectoryNodePoses();
            var allSubmapData = poseGraph.GetAllSubmapData();

            // Check at least 1 trajectory
            if (trajectoryStates.Count == 0)
            {
                return (false, "No trajectories found. Cannot save empty map.");
            }

            // Check at least 1 node
            if (trajectoryNodePoses.Count == 0)
            {
                return (false, "No trajectory nodes found. Cannot save map without nodes.");
            }

            // Check at least 1 submap
            if (allSubmapData.Count == 0)
            {
                return (false, "No submaps found. Cannot save map without submaps.");
            }

            // Validate map name
            var sanitizedMapName = MapNameHelper.Sanitize(mapName);
            if (string.IsNullOrWhiteSpace(sanitizedMapName))
            {
                return (false, $"Invalid map name: '{mapName}' (sanitized to empty string)");
            }

            // Check poses for NaN/Infinity
            int invalidPoseCount = 0;
            foreach (var nodePose in trajectoryNodePoses)
            {
                var globalPose = nodePose.Data.GlobalPose;
                if (!globalPose.IsValid())
                {
                    invalidPoseCount++;
                }
            }

            if (invalidPoseCount > 0)
            {
                return (false, $"Found {invalidPoseCount} trajectory node(s) with invalid poses (NaN/Infinity or invalid quaternion).");
            }

            // Check submap poses for NaN/Infinity
            int invalidSubmapPoseCount = 0;
            foreach (var submapPose in poseGraph.GetAllSubmapPoses())
            {
                var pose = submapPose.Data.Pose;
                if (!pose.IsValid())
                {
                    invalidSubmapPoseCount++;
                }
            }

            if (invalidSubmapPoseCount > 0)
            {
                return (false, $"Found {invalidSubmapPoseCount} submap(s) with invalid poses (NaN/Infinity or invalid quaternion).");
            }

            // Validate occupancy grid resolution
            if (_config.MapStorage.OccupancyGridResolution <= 0)
            {
                return (false, $"Invalid occupancy grid resolution: {_config.MapStorage.OccupancyGridResolution}. Must be > 0.");
            }

            // All validations passed
            return (true, string.Empty);
        }
        catch (Exception ex)
        {
            return (false, $"Validation error: {ex.Message}");
        }
    }

    #endregion

    #region Grid Generation

    /// <summary>
    /// Generate occupancy grid from MapBuilder for save
    /// </summary>
    private (OccupancyGrid? Grid, int SubmapCount) GenerateOccupancyGridForSave(IMapBuilder mapBuilder, double targetResolution)
    {
        try
        {
            var padding = _config.MapStorage.MapPadding;

            // Use the same Generate() dispatch as OccupancyGridManager (UI display)
            // to ensure saved files match what the user sees on the UI.
            // Generate() respects OccupancyGridConfiguration.MergeStrategy (default: LogOddsSum).
            var occupancyGrid = OccupancyGridGenerator.Generate(
                mapBuilder, targetResolution, padding, _logger, _config.OccupancyGrid);

            if (occupancyGrid == null)
            {
                _logger.LogWarning("MapSaveProcessor: Occupancy grid generation returned null");
                return (null, 0);
            }

            // Get submap count for result
            var submapCount = mapBuilder.PoseGraph.GetAllSubmapData()?.Count ?? 0;

            return (occupancyGrid, submapCount);
        }
        catch (Exception ex)
        {
            _logger.LogError(ex, "MapSaveProcessor: Failed to generate occupancy grid from MapBuilder");
            return (null, 0);
        }
    }

    #endregion
}