using System;
using System.Globalization;
using System.IO;
using UnityEngine;
using UnityEngine.XR;

public class XRSessionLogger : MonoBehaviour
{
    public static XRSessionLogger Instance { get; private set; }

    [Header("Session")]
    [SerializeField] private string sessionLabel = "QuestSession";

    [Header("Tracked Transforms")]
    [SerializeField] private Transform hmd;
    [SerializeField] private Transform leftController;
    [SerializeField] private Transform rightController;

    [Header("Sampling")]
    [SerializeField, Range(1f, 120f)] private float poseSampleRateHz = 30f;

    public string SessionDirectory { get; private set; }
    public double ElapsedSeconds => Time.realtimeSinceStartupAsDouble - startRealtime;

    public int PoseSamples { get; private set; }
    public int EventCount { get; private set; }
    public int HandRows { get; private set; }
    public int CameraFrames { get; private set; }

    private StreamWriter poseWriter;
    private StreamWriter eventWriter;
    private StreamWriter handWriter;
    private StreamWriter cameraWriter;

    private double startRealtime;
    private double nextPoseTime;
    private double nextFlushTime;

    private InputDevice headDevice;
    private InputDevice leftDevice;
    private InputDevice rightDevice;

    private static readonly CultureInfo Inv = CultureInfo.InvariantCulture;

    [Serializable]
    private class SessionMetadata
    {
        public string sessionLabel;
        public string startedUtc;
        public string unityVersion;
        public string deviceModel;
        public string operatingSystem;
        public string applicationVersion;
        public string packageName;
        public string coordinateFrame;
        public float poseSampleRateHz;
    }

    private void Awake()
    {
        if (Instance != null && Instance != this)
        {
            Destroy(gameObject);
            return;
        }

        Instance = this;
    }

    private void Start()
    {
        BeginSession();
    }

    private void Update()
    {
        if (string.IsNullOrEmpty(SessionDirectory))
            return;

        double now = Time.realtimeSinceStartupAsDouble;

        if (now >= nextPoseTime)
        {
            SamplePosesAndControllerInputs();
            nextPoseTime = now + (1.0 / poseSampleRateHz);
        }

        if (now >= nextFlushTime)
        {
            FlushAll();
            nextFlushTime = now + 1.0;
        }
    }

    private void BeginSession()
    {
        startRealtime = Time.realtimeSinceStartupAsDouble;
        nextPoseTime = startRealtime;
        nextFlushTime = startRealtime + 1.0;

        string safeLabel = MakeSafeFilename(sessionLabel);
        string folderName = $"{safeLabel}_{DateTime.Now:yyyyMMdd_HHmmss}";

        SessionDirectory = Path.Combine(
            Application.persistentDataPath,
            "XRDataSessions",
            folderName);

        Directory.CreateDirectory(SessionDirectory);
        Directory.CreateDirectory(Path.Combine(SessionDirectory, "camera"));

        poseWriter = CreateWriter("poses.csv",
            "time_utc,elapsed_s," +
            "hmd_tracked,hmd_px,hmd_py,hmd_pz,hmd_qx,hmd_qy,hmd_qz,hmd_qw," +
            "left_tracked,left_px,left_py,left_pz,left_qx,left_qy,left_qz,left_qw," +
            "left_trigger,left_grip,left_axis_x,left_axis_y,left_primary,left_secondary," +
            "right_tracked,right_px,right_py,right_pz,right_qx,right_qy,right_qz,right_qw," +
            "right_trigger,right_grip,right_axis_x,right_axis_y,right_primary,right_secondary");

        eventWriter = CreateWriter("events.csv",
            "time_utc,elapsed_s,event_type,target,interactor,details");

        handWriter = CreateWriter("hands.csv",
            "time_utc,elapsed_s,hand,joint,hand_tracked,joint_valid,px,py,pz,qx,qy,qz,qw");

        cameraWriter = CreateWriter("camera.csv",
            "time_utc,elapsed_s,camera_timestamp_utc,filename," +
            "px,py,pz,qx,qy,qz,qw,width,height," +
            "focal_x,focal_y,principal_x,principal_y");

        var metadata = new SessionMetadata
        {
            sessionLabel = sessionLabel,
            startedUtc = DateTime.UtcNow.ToString("O"),
            unityVersion = Application.unityVersion,
            deviceModel = SystemInfo.deviceModel,
            operatingSystem = SystemInfo.operatingSystem,
            applicationVersion = Application.version,
            packageName = Application.identifier,
            coordinateFrame = "Unity world coordinates; meters; left-handed; Y-up",
            poseSampleRateHz = poseSampleRateHz
        };

        File.WriteAllText(
            Path.Combine(SessionDirectory, "session.json"),
            JsonUtility.ToJson(metadata, true));

        Debug.Log($"XR data session started: {SessionDirectory}");
    }

    private StreamWriter CreateWriter(string filename, string header)
    {
        var writer = new StreamWriter(Path.Combine(SessionDirectory, filename), false);
        writer.WriteLine(header);
        writer.AutoFlush = false;
        return writer;
    }

    private void SamplePosesAndControllerInputs()
    {
        RefreshDevicesIfNeeded();

        bool headTracked = IsTracked(headDevice);
        bool leftTracked = IsTracked(leftDevice);
        bool rightTracked = IsTracked(rightDevice);

        Vector3 hmdPos = hmd != null ? hmd.position : Vector3.zero;
        Quaternion hmdRot = hmd != null ? hmd.rotation : Quaternion.identity;

        Vector3 leftPos = leftController != null ? leftController.position : Vector3.zero;
        Quaternion leftRot = leftController != null ? leftController.rotation : Quaternion.identity;

        Vector3 rightPos = rightController != null ? rightController.position : Vector3.zero;
        Quaternion rightRot = rightController != null ? rightController.rotation : Quaternion.identity;

        ReadController(leftDevice,
            out float lTrigger, out float lGrip, out Vector2 lAxis,
            out bool lPrimary, out bool lSecondary);

        ReadController(rightDevice,
            out float rTrigger, out float rGrip, out Vector2 rAxis,
            out bool rPrimary, out bool rSecondary);

        poseWriter.WriteLine(string.Join(",",
            UtcNow(), F(ElapsedSeconds),
            B(headTracked), F(hmdPos.x), F(hmdPos.y), F(hmdPos.z),
            F(hmdRot.x), F(hmdRot.y), F(hmdRot.z), F(hmdRot.w),
            B(leftTracked), F(leftPos.x), F(leftPos.y), F(leftPos.z),
            F(leftRot.x), F(leftRot.y), F(leftRot.z), F(leftRot.w),
            F(lTrigger), F(lGrip), F(lAxis.x), F(lAxis.y), B(lPrimary), B(lSecondary),
            B(rightTracked), F(rightPos.x), F(rightPos.y), F(rightPos.z),
            F(rightRot.x), F(rightRot.y), F(rightRot.z), F(rightRot.w),
            F(rTrigger), F(rGrip), F(rAxis.x), F(rAxis.y), B(rPrimary), B(rSecondary)));

        PoseSamples++;
    }

    private void RefreshDevicesIfNeeded()
    {
        if (!headDevice.isValid)
            headDevice = InputDevices.GetDeviceAtXRNode(XRNode.Head);
        if (!leftDevice.isValid)
            leftDevice = InputDevices.GetDeviceAtXRNode(XRNode.LeftHand);
        if (!rightDevice.isValid)
            rightDevice = InputDevices.GetDeviceAtXRNode(XRNode.RightHand);
    }

    private static bool IsTracked(InputDevice device)
    {
        if (!device.isValid)
            return false;
        if (device.TryGetFeatureValue(CommonUsages.isTracked, out bool tracked))
            return tracked;
        return true;
    }

    private static void ReadController(InputDevice device, out float trigger, out float grip,
        out Vector2 axis, out bool primary, out bool secondary)
    {
        trigger = 0f;
        grip = 0f;
        axis = Vector2.zero;
        primary = false;
        secondary = false;

        if (!device.isValid)
            return;

        device.TryGetFeatureValue(CommonUsages.trigger, out trigger);
        device.TryGetFeatureValue(CommonUsages.grip, out grip);
        device.TryGetFeatureValue(CommonUsages.primary2DAxis, out axis);
        device.TryGetFeatureValue(CommonUsages.primaryButton, out primary);
        device.TryGetFeatureValue(CommonUsages.secondaryButton, out secondary);
    }

    public void LogEvent(string eventType, string target, string interactor = "", string details = "")
    {
        if (eventWriter == null)
            return;

        eventWriter.WriteLine(string.Join(",", UtcNow(), F(ElapsedSeconds),
            Csv(eventType), Csv(target), Csv(interactor), Csv(details)));
        EventCount++;
    }

    public void LogHandJoint(string hand, string joint, bool handTracked, bool jointValid,
        Vector3 position, Quaternion rotation)
    {
        if (handWriter == null)
            return;

        handWriter.WriteLine(string.Join(",", UtcNow(), F(ElapsedSeconds), Csv(hand), Csv(joint),
            B(handTracked), B(jointValid), F(position.x), F(position.y), F(position.z),
            F(rotation.x), F(rotation.y), F(rotation.z), F(rotation.w)));
        HandRows++;
    }

    public void LogCameraFrame(DateTime cameraTimestamp, string filename, Pose cameraPose,
        Vector2Int resolution, Vector2 focalLength, Vector2 principalPoint)
    {
        if (cameraWriter == null)
            return;

        cameraWriter.WriteLine(string.Join(",", UtcNow(), F(ElapsedSeconds),
            Csv(cameraTimestamp.ToUniversalTime().ToString("O")), Csv(filename),
            F(cameraPose.position.x), F(cameraPose.position.y), F(cameraPose.position.z),
            F(cameraPose.rotation.x), F(cameraPose.rotation.y), F(cameraPose.rotation.z),
            F(cameraPose.rotation.w), resolution.x.ToString(Inv), resolution.y.ToString(Inv),
            F(focalLength.x), F(focalLength.y), F(principalPoint.x), F(principalPoint.y)));
        CameraFrames++;
    }

    private void OnApplicationPause(bool pause)
    {
        if (pause) FlushAll();
    }

    private void OnApplicationQuit() => CloseAll();

    private void OnDestroy()
    {
        if (Instance == this)
        {
            CloseAll();
            Instance = null;
        }
    }

    private void FlushAll()
    {
        poseWriter?.Flush();
        eventWriter?.Flush();
        handWriter?.Flush();
        cameraWriter?.Flush();
    }

    private void CloseAll()
    {
        FlushAll();
        poseWriter?.Close();
        eventWriter?.Close();
        handWriter?.Close();
        cameraWriter?.Close();
        poseWriter = null;
        eventWriter = null;
        handWriter = null;
        cameraWriter = null;
    }

    private static string UtcNow() => DateTime.UtcNow.ToString("O");
    private static string F(double value) => value.ToString("0.######", Inv);
    private static string F(float value) => value.ToString("0.######", Inv);
    private static string B(bool value) => value ? "1" : "0";

    private static string Csv(string value)
    {
        if (string.IsNullOrEmpty(value)) return "";
        return "\"" + value.Replace("\"", "\"\"") + "\"";
    }

    private static string MakeSafeFilename(string value)
    {
        if (string.IsNullOrWhiteSpace(value)) return "Session";
        foreach (char c in Path.GetInvalidFileNameChars()) value = value.Replace(c, '_');
        return value.Replace(' ', '_');
    }
}