using CartographerSharp.IO;
using CartographerSharp.Transform;
using RobotNet10.Shared.Geometry;
using RobotNet10.Shared.Numbers;

namespace RobotNet10.RobotApp.SLAM.Cartographer.Helpers;

/// <summary>
/// Helper class for transforming pbstream files.
/// Centralizes the logic for updating TransformToMap in PoseGraph proto.
/// </summary>
public static class PbstreamTransformHelper
{
    /// <summary>
    /// Transform pbstream file by updating TransformToMap in PoseGraph proto.
    /// This applies a coordinate transform to shift the map origin.
    /// </summary>
    /// <param name="pbstreamPath">Path to the pbstream file</param>
    /// <param name="newOrigin">New origin pose to apply</param>
    /// <param name="logger">Optional logger for logging</param>
    /// <exception cref="InvalidOperationException">Thrown when pbstream read/write fails</exception>
    /// <exception cref="ArgumentException">Thrown when newOrigin contains invalid values (NaN/Infinity)</exception>
    public static void TransformPbstreamOrigin(string pbstreamPath, Pose newOrigin, ILogger? logger = null)
    {
        // Validate newOrigin for NaN/Infinity
        ValidatePose(newOrigin);

        var tempPbstreamPath = pbstreamPath + ".tmp";
        try
        {
            using (var reader = new ProtoStreamReader(pbstreamPath))
            using (var writer = new ProtoStreamWriter(tempPbstreamPath))
            {
                // Read and write header
                if (!reader.ReadProto<CartographerSharp.Models.Mapping.SerializedData>(out var headerData))
                {
                    throw new InvalidOperationException("Failed to read serialization header");
                }
                if (!headerData.SerializationHeader.HasValue)
                {
                    throw new InvalidOperationException("Invalid serialization header");
                }
                writer.WriteProto(headerData);

                // Read PoseGraph, update TransformToMap, and write
                if (!reader.ReadProto<CartographerSharp.Models.Mapping.SerializedData>(out var poseGraphData))
                {
                    throw new InvalidOperationException("Failed to read pose graph");
                }
                if (!poseGraphData.PoseGraph.HasValue)
                {
                    throw new InvalidOperationException("Invalid pose graph data");
                }

                var poseGraphProto = poseGraphData.PoseGraph.Value;

                // Convert Pose to Rigid3d and compose with current transform
                // Key logic from xloc.cc ChangeMapOrigin:
                // SetTransformToMap(GetTransformToMap() * Rigid3d(position, orientation));
                //
                // In Cartographer:
                // - TransformToMap: converts map frame -> internal frame
                // - TransformToMapInverse: converts internal frame -> map frame (used in OccupancyGridGenerator)
                //
                // We want: newOrigin.Position (T) in old map becomes (0,0) in new map
                // Formula: newMapPoint = R^-1 * (oldMapPoint - T) where T = newOrigin.Position, R = newOrigin.Orientation
                var newOriginRigid = new Rigid3d(
                    new Vector3(newOrigin.Position.X, newOrigin.Position.Y, newOrigin.Position.Z),
                    new Quaternion(newOrigin.Orientation.X, newOrigin.Orientation.Y, newOrigin.Orientation.Z, newOrigin.Orientation.W));

                var currentTransform = poseGraphProto.TransformToMap.HasValue
                    ? (Rigid3d)poseGraphProto.TransformToMap.Value
                    : Rigid3d.Identity;
                var newTransform = currentTransform * newOriginRigid;

                // Update TransformToMap in the proto
                poseGraphProto.TransformToMap = (CartographerSharp.Models.Transform.Rigid3dProto)newTransform;

                // Write updated pose graph
                var updatedPoseGraphData = new CartographerSharp.Models.Mapping.SerializedData { PoseGraph = poseGraphProto };
                writer.WriteProto(updatedPoseGraphData);

                // Copy all remaining data unchanged (AllTrajectoryBuilderOptions, submaps, nodes, etc.)
                while (!reader.Eof)
                {
                    if (!reader.ReadProto<CartographerSharp.Models.Mapping.SerializedData>(out var data))
                    {
                        break;
                    }
                    writer.WriteProto(data);
                }
            }

            // Replace original file with updated file
            File.Move(tempPbstreamPath, pbstreamPath, overwrite: true);
            logger?.LogInformation("PbstreamTransformHelper: Transformed pbstream origin successfully");
        }
        catch (Exception)
        {
            // Clean up temp file on error
            if (File.Exists(tempPbstreamPath))
            {
                try { File.Delete(tempPbstreamPath); } catch { }
            }
            throw;
        }
    }

    /// <summary>
    /// Check if pose is identity (no transform needed)
    /// </summary>
    public static bool IsIdentityPose(Pose pose)
    {
        const double tolerance = 1e-6;

        // Check position is (0, 0, 0)
        var posIsZero = Math.Abs(pose.Position.X) < tolerance &&
                        Math.Abs(pose.Position.Y) < tolerance &&
                        Math.Abs(pose.Position.Z) < tolerance;

        // Check orientation is identity quaternion (0, 0, 0, 1)
        var quatIsIdentity = Math.Abs(pose.Orientation.X) < tolerance &&
                             Math.Abs(pose.Orientation.Y) < tolerance &&
                             Math.Abs(pose.Orientation.Z) < tolerance &&
                             Math.Abs(pose.Orientation.W - 1.0) < tolerance;

        return posIsZero && quatIsIdentity;
    }

    /// <summary>
    /// Validate pose for NaN/Infinity values
    /// </summary>
    /// <exception cref="ArgumentException">Thrown when pose contains invalid values</exception>
    public static void ValidatePose(Pose pose)
    {
        // Check position
        if (double.IsNaN(pose.Position.X) || double.IsInfinity(pose.Position.X) ||
            double.IsNaN(pose.Position.Y) || double.IsInfinity(pose.Position.Y) ||
            double.IsNaN(pose.Position.Z) || double.IsInfinity(pose.Position.Z))
        {
            throw new ArgumentException($"Pose position contains invalid values (NaN/Infinity): ({pose.Position.X}, {pose.Position.Y}, {pose.Position.Z})", nameof(pose));
        }

        // Check orientation
        if (double.IsNaN(pose.Orientation.X) || double.IsInfinity(pose.Orientation.X) ||
            double.IsNaN(pose.Orientation.Y) || double.IsInfinity(pose.Orientation.Y) ||
            double.IsNaN(pose.Orientation.Z) || double.IsInfinity(pose.Orientation.Z) ||
            double.IsNaN(pose.Orientation.W) || double.IsInfinity(pose.Orientation.W))
        {
            throw new ArgumentException($"Pose orientation contains invalid values (NaN/Infinity): ({pose.Orientation.X}, {pose.Orientation.Y}, {pose.Orientation.Z}, {pose.Orientation.W})", nameof(pose));
        }
    }

    /// <summary>
    /// Convert Pose to Rigid3d
    /// </summary>
    public static Rigid3d PoseToRigid3d(Pose pose)
    {
        var translation = new Vector3(pose.Position.X, pose.Position.Y, pose.Position.Z);
        var rotation = new Quaternion(pose.Orientation.X, pose.Orientation.Y, pose.Orientation.Z, pose.Orientation.W);
        return new Rigid3d(translation, rotation);
    }
}