using System.Collections.Generic;
using System.IO;
using System.Text;
using UnityEngine;

namespace HexapodSim
{
    // Runs the real robot firmware (robot/, compiled into sim/build/hexapod_sim.dll)
    // against the existing Unity model:
    //   input (SimRemote) -> CmdPacket over the simulated radio -> firmware loop()
    //   -> 18 servo pulses -> joint rotations on the rig.
    // The body's motion (kinematic: it rests on its feet, grounded feet don't
    // slide) and the simulated MPU6050 readings are computed inside the DLL
    // (sim/SimBody.cpp), so the headless smoke test checks the same thing this
    // shows. Nothing here decides how the robot moves.
    [DefaultExecutionOrder(-100)]
    public class HexapodSimulation : MonoBehaviour
    {
        [Header("Firmware")]
        [Tooltip("Empty = <repo>/sim/build/hexapod_sim.dll")]
        public string dllPath = "";
        public bool radioPresent = true;
        public bool imuPresent = true;
        [Tooltip("Remote transmit rate. The real remote sends at 100 Hz.")]
        public float remoteTxHz = 100f;

        [Header("Rig (auto-detected from the scene if empty)")]
        public Transform chassis;
        public Transform ground;
        public ForwardAxis forwardAxis = ForwardAxis.NegativeX;
        [Tooltip("Unity units per millimetre. The model is built in cm.")]
        public float unitsPerMm = 0.1f;
        [Tooltip("Firmware tibia angle when the rig's Knie rotation is 0 (tibia straight down from a horizontal femur).")]
        public float kneeRestDeg = -90f;

        [Header("World")]
        [Tooltip("Tilts the ground (deg about robot X = roll, Y = pitch) to exercise the robot gyro.")]
        public Vector2 groundTiltDeg;
        public bool followCamera = true;

        // ---- Runtime state, read by SimTestPanel ----
        public HexapodSimNative Native { get; private set; }
        public HxGeometry Geometry { get; private set; }
        public HxStatus Status { get; private set; }
        public string LoadError { get; private set; }
        public readonly List<string> SerialLog = new List<string>();
        public readonly List<string> RigWarnings = new List<string>();
        public float[] JointsDeg { get; } = new float[18];
        public float[] FeetMm { get; } = new float[18];
        public bool[] Contact { get; } = new bool[6];
        public int[] ServoUs { get; } = new int[18];
        public int[] ServoAttached { get; } = new int[18];
        public uint[] ServoWrites { get; } = new uint[18];
        public Vector3 BodyPosMm => bodyPos_;
        public float BodyYawDeg => yaw_;

        SimRemote remote_;

        class LegRig
        {
            public Transform root, hip1, hip2, knie, footMarker;
            public float coxaOffsetDeg;
        }
        readonly LegRig[] legs_ = new LegRig[6];

        // Body pose from the DLL, robot-world frame (mm, Z up, ground z = 0 before tilt).
        float yaw_;
        Vector3 bodyPos_;
        Quaternion bodyRot_ = Quaternion.identity;   // robot-frame rotation used as a math object

        // Unity-side anchors captured at start.
        Vector3 originUnity_;
        Quaternion chassisRestRot_;
        Vector3 groundRestPos_;
        Quaternion groundRestRot_;
        Transform camera_;
        Vector3 cameraOffset_;
        StringBuilder serialLine_ = new StringBuilder();

        static readonly string[] LegacyMovementScripts =
            { "Walk", "RotationControl", "Leg", "HandleHoystick", "Serial2", "Serial3" };

        // Robot-frame axes, for readability (Unity's Vector3.up etc. are Unity axes).
        static readonly Vector3 AxisX = new Vector3(1, 0, 0);   // forward
        static readonly Vector3 AxisY = new Vector3(0, 1, 0);   // left

        void Awake()
        {
            DisableLegacyScripts();
            remote_ = GetComponent<SimRemote>();
            if (remote_ == null) remote_ = gameObject.AddComponent<SimRemote>();
            if (GetComponent<SimTestPanel>() == null) gameObject.AddComponent<SimTestPanel>();

            if (chassis == null) chassis = FindChassis();
            if (ground == null) { var g = GameObject.Find("Plane"); if (g) ground = g.transform; }
            if (chassis == null) { LoadError = "No leg rig found (expected 6 objects named 'Leg' with hip1/Hip2/Knie)."; return; }

            chassisRestRot_ = chassis.rotation;
            float groundY = ground ? ground.position.y : chassis.position.y - 150f * unitsPerMm;
            originUnity_ = new Vector3(chassis.position.x, groundY, chassis.position.z);
            if (ground) { groundRestPos_ = ground.position; groundRestRot_ = ground.rotation; }
            if (Camera.main) { camera_ = Camera.main.transform; cameraOffset_ = camera_.position - chassis.position; }

            ReloadFirmware();
        }

        void OnDestroy() => Native?.Dispose();

        public string ResolvedDllPath =>
            !string.IsNullOrEmpty(dllPath)
                ? dllPath
                : Path.GetFullPath(Path.Combine(Application.dataPath, "..", "..", "sim", "build", "hexapod_sim.dll"));

        // Fresh power-on: loads the current DLL build and resets the body.
        public void ReloadFirmware()
        {
            Native?.Dispose();
            Native = null;
            LoadError = null;
            try
            {
                Native = HexapodSimNative.Load(ResolvedDllPath);
                Native.SetRadioPresent(radioPresent);
                Native.SetImuPresent(imuPresent);
                Geometry = Native.Geometry;
                BindRig();
                Native.SetGroundTilt(groundTiltDeg.x, groundTiltDeg.y);
                Native.Start();
                remote_.ResetState();
                SerialLog.Add($"--- firmware loaded: {ResolvedDllPath}");
            }
            catch (System.Exception e)
            {
                LoadError = e.Message;
                Debug.LogError("[HexapodSim] " + e.Message);
                Native?.Dispose();
                Native = null;
            }
        }

        void FixedUpdate()
        {
            if (Native == null) return;
            float dt = Time.fixedDeltaTime;

            Native.SetGroundTilt(groundTiltDeg.x, groundTiltDeg.y);

            // Remote packets at the remote's rate, firmware time advanced in step.
            int sub = Mathf.Max(1, Mathf.RoundToInt(dt * remoteTxHz));
            float subDt = dt / sub;
            for (int n = 0; n < sub; n++)
            {
                remote_.Transmit(Native, subDt);
                Native.Step((uint)Mathf.RoundToInt(subDt * 1e6f));
            }

            Status = Native.Status;
            Native.GetJointsDeg(JointsDeg);
            Native.GetFeetMm(FeetMm);
            Native.GetServos(ServoUs, ServoAttached, ServoWrites);
            Native.GetBody(out bodyPos_, out bodyRot_, out yaw_, out int contactMask);
            for (int i = 0; i < 6; i++) Contact[i] = (contactMask >> i & 1) != 0;
            DrainSerial();

            ApplyBodyPose();
            ApplyJoints();
        }

        void LateUpdate()
        {
            if (followCamera && camera_ && chassis) camera_.position = chassis.position + cameraOffset_;
        }

        // ------------------------------------------------------------------
        // Rig
        // ------------------------------------------------------------------

        void DisableLegacyScripts()
        {
            foreach (var mb in FindObjectsByType<MonoBehaviour>(FindObjectsInactive.Include))
                if (System.Array.IndexOf(LegacyMovementScripts, mb.GetType().Name) >= 0)
                    mb.enabled = false;
        }

        static List<Transform> FindLegRoots()
        {
            var result = new List<Transform>();
            foreach (var t in FindObjectsByType<Transform>())
                if (t.name == "Leg" && FindDeep(t, "hip1") != null) result.Add(t);
            return result;
        }

        static Transform FindChassis()
        {
            var roots = FindLegRoots();
            return roots.Count > 0 ? roots[0].parent : null;
        }

        static Transform FindDeep(Transform parent, string name)
        {
            foreach (Transform c in parent)
            {
                if (string.Equals(c.name, name, System.StringComparison.OrdinalIgnoreCase)) return c;
                var d = FindDeep(c, name);
                if (d) return d;
            }
            return null;
        }

        // Matches each firmware leg (Config.h order) to the rig leg at its mount
        // position, and checks the rig against Config.h so a mismatch is loud.
        void BindRig()
        {
            RigWarnings.Clear();
            var roots = FindLegRoots();
            var g = Geometry;
            for (int i = 0; i < 6; i++)
            {
                Vector3 mountLocal = RobotFrame.ToUnity(new Vector3(g.mountXmm[i], g.mountYmm[i], 0), forwardAxis) * unitsPerMm;
                Transform best = null;
                float bestDist = float.MaxValue;
                foreach (var r in roots)
                {
                    float dist = Vector3.Distance(chassis.InverseTransformPoint(r.position), mountLocal);
                    if (dist < bestDist) { bestDist = dist; best = r; }
                }
                if (best == null) { legs_[i] = null; continue; }
                roots.Remove(best);

                var leg = new LegRig
                {
                    root = best,
                    hip1 = FindDeep(best, "hip1"),
                    hip2 = FindDeep(best, "Hip2"),
                    knie = FindDeep(best, "Knie"),
                    footMarker = FindDeep(best, "Point"),
                };
                if (bestDist / unitsPerMm > 5f)
                    RigWarnings.Add($"leg {i}: rig hip is {bestDist / unitsPerMm:F0} mm from Config.h mount ({g.mountXmm[i]}, {g.mountYmm[i]})");

                // Direction the rig's coxa points at rotation 0, in robot-frame degrees.
                Vector3 outward = RobotFrame.ToRobot(chassis.InverseTransformDirection(best.forward), forwardAxis);
                float rigMountDeg = Mathf.Atan2(outward.y, outward.x) * Mathf.Rad2Deg;
                leg.coxaOffsetDeg = RobotFrame.WrapDeg(g.mountAngleDeg[i] - rigMountDeg);
                if (Mathf.Abs(leg.coxaOffsetDeg) > 1f)
                    RigWarnings.Add($"leg {i}: Config mount angle {g.mountAngleDeg[i]:F0} vs rig {rigMountDeg:F0} deg (compensated)");

                if (leg.hip1 && leg.hip2 && leg.knie)
                {
                    float rigCoxa = Vector3.Distance(leg.hip1.position, leg.hip2.position) / unitsPerMm;
                    float rigFemur = Vector3.Distance(leg.hip2.position, leg.knie.position) / unitsPerMm;
                    if (Mathf.Abs(rigCoxa - g.coxaMm) > 2f || Mathf.Abs(rigFemur - g.femurMm) > 2f)
                        RigWarnings.Add($"leg {i}: rig coxa/femur {rigCoxa:F0}/{rigFemur:F0} mm vs Config {g.coxaMm:F0}/{g.femurMm:F0} mm");
                }
                else RigWarnings.Add($"leg {i}: missing hip1/Hip2/Knie under '{best.name}'");
                legs_[i] = leg;
            }
            foreach (var w in RigWarnings) Debug.LogWarning("[HexapodSim] " + w);
        }

        // Firmware joint convention -> rig rotations (see sim/HexapodSimApi.cpp FK):
        //   coxa + = CCW from above; Unity +Y rotation is clockwise from above.
        //   femur/tibia + = up; Unity +X rotation tips the leg's +Z down.
        void ApplyJoints()
        {
            for (int i = 0; i < 6; i++)
            {
                var leg = legs_[i];
                if (leg == null || leg.hip1 == null) continue;
                float coxa = JointsDeg[i * 3], femur = JointsDeg[i * 3 + 1], tibia = JointsDeg[i * 3 + 2];
                leg.hip1.localRotation = Quaternion.Euler(0, -(coxa + leg.coxaOffsetDeg), 0);
                if (leg.hip2) leg.hip2.localRotation = Quaternion.Euler(-femur, 0, 0);
                if (leg.knie) leg.knie.localRotation = Quaternion.Euler(-(tibia - kneeRestDeg), 0, 0);

                // The sphere marks where the firmware-side FK puts the foot. If it
                // doesn't sit on the tibia tip, the rig and Config.h disagree.
                if (leg.footMarker)
                    leg.footMarker.position = chassis.TransformPoint(FootLocal(i));
            }
        }

        Vector3 FootLocal(int i) =>
            RobotFrame.ToUnity(new Vector3(FeetMm[i * 3], FeetMm[i * 3 + 1], FeetMm[i * 3 + 2]), forwardAxis) * unitsPerMm;

        // ------------------------------------------------------------------
        // World
        // ------------------------------------------------------------------

        // Robot frame: X forward, Y left (Unity's Vector3.right/up are Unity axes).
        Quaternion GroundRotation() =>
            Quaternion.AngleAxis(groundTiltDeg.x, AxisX) * Quaternion.AngleAxis(groundTiltDeg.y, AxisY);

        // The robot-world frame is the chassis's starting orientation, so the
        // rig's local axes (which the Config.h mapping is defined against) and
        // the world stay consistent even if the chassis starts rotated.
        void ApplyBodyPose()
        {
            if (!chassis) return;
            chassis.position = originUnity_ + chassisRestRot_ * RobotFrame.ToUnity(bodyPos_, forwardAxis) * unitsPerMm;
            chassis.rotation = chassisRestRot_ * RobotFrame.ToUnity(bodyRot_, forwardAxis);
            if (ground)
            {
                Quaternion g = chassisRestRot_ * RobotFrame.ToUnity(GroundRotation(), forwardAxis) * Quaternion.Inverse(chassisRestRot_);
                ground.rotation = g * groundRestRot_;
                ground.position = originUnity_ + g * (groundRestPos_ - originUnity_);
            }
        }

        void DrainSerial()
        {
            string s = Native.SerialOutput();
            foreach (char ch in s)
            {
                if (ch == '\n')
                {
                    string line = serialLine_.ToString().TrimEnd('\r');
                    serialLine_.Clear();
                    SerialLog.Add(line);
                    Debug.Log("[robot] " + line);
                }
                else serialLine_.Append(ch);
            }
            if (SerialLog.Count > 200) SerialLog.RemoveRange(0, SerialLog.Count - 200);
        }
    }
}