using Godot;

// Body-physics Stage 2: physical arm segments. Upper arm + forearm per side are chase-driven
// RigidBody3Ds (no joints — Box3D has no usable 3-DOF shoulder/wrist joint, and AGENTS.md
// forbids joints between tracked bodies). Each segment PD-chases its analytic 2-bone IK pose,
// solved every physics tick from the PHYSICAL chest shoulders to the PHYSICAL hand bodies.
// Hierarchical budgets (Nimso's muscle model): hand 90 N < forearm < upper arm < chest 2400 N.
// The arm IK + reach budget (torso lean on overreach) live here; BodyRig only draws.
public partial class PhysicsArms : Node
{
    private const float BodyMassStat = 60f;
    private const float ForearmMass = 1.2f;
    private const float UpperArmMass = 2.0f;
    private const float ForearmForcePerKg = 3f;   // 180 N at 60 kg
    private const float UpperArmForcePerKg = 6f;  // 360 N at 60 kg
    private const float SegmentMaxSpin = 40f;
    private const float ForearmMaxAlpha = 1500f;
    private const float UpperArmMaxAlpha = 1000f;
    private const float SegmentInset = 0.04f;
    public const float SnapDistance = 1f;

    private const float ShoulderRise = 0.25f;
    private const float ReachBudgetMax = 0.25f;
    private const float ReachVerticalMin = -0.05f;
    private const float ReachVerticalMax = 0.12f;
    private const float LeanReachFactor = 0.95f;
    private const float LeanSmoothRate = 6f;
    private const float WristOffset = 0.06f;
    private const float SoftIKFactor = 0.98f;
    private const float PoleHandWeight = 0.85f;
    private const float PoleSmoothRate = 12f;
    private const float PoleMaxAngle = 1.92f;  // 110 deg either way from down/back/out
    private const float PoleMaxSpeed = 6f;     // rad/s
    public const float HandArmSlack = 0.03f;
    private const float ElbowRiseAllowance = 0.05f;

    public static float ForearmMaxForce => BodyMassStat * ForearmForcePerKg;
    public static float UpperArmMaxForce => BodyMassStat * UpperArmForcePerKg;

    // Injected by headless spikes; resolved from the main scene otherwise.
    public Node3D HeadSource;
    public PhysicsTorso Torso;
    public RigidBody3D HandLeft;
    public RigidBody3D HandRight;
    // Hand ORIENTATION source for the elbow: the tracked controllers, not the physics palms
    // (a palm fighting contact must not swing the elbow). Falls back to the hand body.
    public Node3D OrientLeft;
    public Node3D OrientRight;

    public bool Calibrated { get; private set; }
    public float UpperLen { get; private set; }
    public float ForeLen { get; private set; }
    public float ShoulderHalfWidth { get; private set; }

    private readonly RigidBody3D[] _upper = new RigidBody3D[2];
    private readonly RigidBody3D[] _fore = new RigidBody3D[2];
    private readonly Vector3[] _shoulder = new Vector3[2];
    private readonly Vector3[] _elbowTarget = new Vector3[2];
    private readonly Vector3[] _wristTarget = new Vector3[2];
    private readonly float[] _poleAngle = new float[2];
    private readonly Vector3[] _lastRef = new Vector3[2];
    private readonly RigidBody3D[] _held = new RigidBody3D[2];
    private Vector3 _torsoOffset;
    private Basis _torsoBasis = Basis.Identity;

    public RigidBody3D UpperArm(int side) => _upper[side];
    public RigidBody3D Forearm(int side) => _fore[side];
    public Vector3 Shoulder(int side) => _shoulder[side];
    public Vector3 ElbowTarget(int side) => _elbowTarget[side];
    public Vector3 WristTarget(int side) => _wristTarget[side];

    // Physical elbow: where the two segment bodies meet (mean of their joint-side ends).
    public Vector3 Elbow(int side) =>
        (SegmentEnd(_upper[side], UpperLen, 1f) + SegmentEnd(_fore[side], ForeLen, -1f)) * 0.5f;

    public Vector3 Wrist(int side) => SegmentEnd(_fore[side], ForeLen, 1f);

    public override void _Ready()
    {
        if (HeadSource == null)
        {
            HeadSource = GetNode<XRCamera3D>("../XROrigin3D/XRCamera3D");
            Torso = GetNode<PhysicsTorso>("../PhysicsTorso");
            HandLeft = PhysicsHands.GetBody(GetNode<XRController3D>("../XROrigin3D/LeftHand"));
            HandRight = PhysicsHands.GetBody(GetNode<XRController3D>("../XROrigin3D/RightHand"));
            OrientLeft = GetNode<XRController3D>("../XROrigin3D/LeftHand");
            OrientRight = GetNode<XRController3D>("../XROrigin3D/RightHand");
        }
    }

    public override void _PhysicsProcess(double delta)
    {
        float dt = (float)delta;
        if (!Calibrated)
        {
            float headHeight = HeadSource.GlobalPosition.Y;
            if (headHeight <= 0.5f || Torso?.Body == null)
                return;
            Calibrate(headHeight);
        }

        _torsoBasis = Torso.Body.GlobalTransform.Basis;
        Vector3 topCenter = Torso.Body.GlobalPosition + Vector3.Up * ShoulderRise;
        _shoulder[0] = topCenter - _torsoBasis.X * ShoulderHalfWidth;
        _shoulder[1] = topCenter + _torsoBasis.X * ShoulderHalfWidth;

        SolveArm(0, -_torsoBasis.X, HandLeft, OrientLeft, dt);
        SolveArm(1, _torsoBasis.X, HandRight, OrientRight, dt);
        UpdateLean(dt);

        for (int side = 0; side < 2; side++)
        {
            Drive(_upper[side], SegmentTarget(_shoulder[side], _elbowTarget[side], _upper[side]), UpperArmMaxForce, UpperArmMaxAlpha, dt);
            Drive(_fore[side], SegmentTarget(_elbowTarget[side], _wristTarget[side], _fore[side]), ForearmMaxForce, ForearmMaxAlpha, dt);
        }
    }

    private void Calibrate(float headHeight)
    {
        Calibrated = true;
        float armTotal = Mathf.Clamp(headHeight * 0.46f, 0.5f, 0.72f);
        UpperLen = 0.5f * armTotal;
        ForeLen = 0.5f * armTotal;
        ShoulderHalfWidth = 0.16f + headHeight * 0.02f;

        for (int side = 0; side < 2; side++)
        {
            string name = side == 0 ? "Left" : "Right";
            _upper[side] = MakeSegment($"{name}UpperArmBody", UpperArmMass, new Vector3(0.07f, UpperLen - SegmentInset, 0.07f));
            _fore[side] = MakeSegment($"{name}ForearmBody", ForearmMass, new Vector3(0.06f, ForeLen - SegmentInset, 0.06f));
            SetHeldExceptions(side, _held[side], true);
        }
        GD.Print($"PhysicsArms: proportions calibrated from head height {headHeight:F2} m " +
                 $"(arm total {armTotal:F2} m, upper {UpperLen:F2} m, fore {ForeLen:F2} m, shoulder half-width {ShoulderHalfWidth:F2} m); " +
                 $"budgets fore {ForearmMaxForce:F0} N, upper {UpperArmMaxForce:F0} N");
    }

    private RigidBody3D MakeSegment(string name, float mass, Vector3 size)
    {
        RigidBody3D body = new()
        {
            Name = name,
            Mass = mass,
            GravityScale = 0f,
            ContinuousCd = true,
            CollisionLayer = 0b10000,  // layer 5 = player arms
            CollisionMask = 0b1,       // world + props only; own chest/hands/arms excluded
            ContactMonitor = true,
            MaxContactsReported = 4,
        };
        body.AddChild(new CollisionShape3D { Shape = new BoxShape3D { Size = size } });
        AddChild(body);
        // Spawn somewhere finite; the first Drive snaps it onto its target (> SnapDistance).
        body.GlobalPosition = Torso.Body.GlobalPosition + Vector3.Down * 100f;
        return body;
    }

    private static void Drive(RigidBody3D body, Transform3D target, float maxForce, float maxAlpha, float dt)
    {
        if (BodyMath.ResetIfNonFinite(body, target, "PhysicsArms"))
            return;
        if (body.GlobalPosition.DistanceTo(target.Origin) > SnapDistance)
        {
            body.GlobalTransform = target;
            body.LinearVelocity = Vector3.Zero;
            body.AngularVelocity = Vector3.Zero;
            return;
        }

        Vector3 desiredVelocity = (target.Origin - body.GlobalPosition) / dt;
        Vector3 force = ((desiredVelocity - body.LinearVelocity) / dt * body.Mass).LimitLength(maxForce);
        body.LinearVelocity += force / body.Mass * dt;

        Vector3 rotErr = BodyMath.RotationError(target.Basis.GetRotationQuaternion(),
            body.GlobalTransform.Basis.GetRotationQuaternion());
        Vector3 spinDesired = (rotErr / dt).LimitLength(SegmentMaxSpin);
        body.AngularVelocity += (spinDesired - body.AngularVelocity).LimitLength(maxAlpha * dt);
    }

    // Segment pose from joint A to joint B: local +Y runs A→B; twist follows the body's current
    // X axis so the long-axis roll never flips (a box arm has no meaningful roll).
    private static Transform3D SegmentTarget(Vector3 a, Vector3 b, RigidBody3D body)
    {
        Vector3 y = b - a;
        if (y.LengthSquared() < 1e-8f)
            y = Vector3.Down;
        y = y.Normalized();
        Vector3 x = body.GlobalTransform.Basis.X;
        x -= y * x.Dot(y);
        if (x.LengthSquared() < 1e-6f)
        {
            x = y.Cross(Vector3.Forward);
            if (x.LengthSquared() < 1e-6f)
                x = y.Cross(Vector3.Right);
        }
        x = x.Normalized();
        return new Transform3D(new Basis(x, y, x.Cross(y)), (a + b) * 0.5f);
    }

    private static Vector3 SegmentEnd(RigidBody3D body, float length, float sign) =>
        body.GlobalPosition + body.GlobalTransform.Basis.Y * (sign * length * 0.5f);

    // A prop held by this side's hand must not collide with this side's arm: the forearm
    // naturally lies along a held object, and 180 N of forearm shoving a prop welded to a 90 N
    // hand is a fight the hand loses. Called by PropGrabber on grab/release.
    public void SetHeld(int side, RigidBody3D prop)
    {
        if (_held[side] == prop)
            return;
        SetHeldExceptions(side, _held[side], false);
        _held[side] = prop;
        SetHeldExceptions(side, prop, true);
    }

    private void SetHeldExceptions(int side, RigidBody3D prop, bool add)
    {
        if (prop == null || !IsInstanceValid(prop) || _upper[side] == null)
            return;
        foreach (RigidBody3D segment in new[] { _upper[side], _fore[side] })
        {
            if (add)
                segment.AddCollisionExceptionWith(prop);
            else
                segment.RemoveCollisionExceptionWith(prop);
        }
    }

    // Stage 2b coupling: while an arm segment is in contact, the hand may lead the forearm's
    // end by at most HandArmSlack — a blocked forearm holds the hand back (the hand ghost
    // shows the controller). Free arms never limit the hand. PhysicsHands calls this with
    // the controller target before chasing; returns true when the limit is active.
    public bool ConstrainHand(int side, ref Transform3D target)
    {
        if (!Calibrated || _fore[side] == null)
            return false;
        if (_fore[side].GetContactCount() == 0 && _upper[side].GetContactCount() == 0)
            return false;
        Vector3 wristOffset = target.Basis * new Vector3(0f, 0f, WristOffset);
        Vector3 forearmEnd = Wrist(side);
        Vector3 lead = target.Origin + wristOffset - forearmEnd;
        if (lead.Length() <= HandArmSlack)
            return true;
        target.Origin = forearmEnd + lead.Normalized() * HandArmSlack - wristOffset;
        return true;
    }

    // Torso lean toward overreaching hands, computed OPEN-LOOP from the head-derived nominal
    // chest pose — never from the body's own shoulders (that loop made the chest hunt:
    // lean moves shoulders → overreach changes → lean changes). Exact two-arm combination:
    // the bigger excess first, then whatever the other arm still lacks after that shift.
    private void UpdateLean(float dt)
    {
        Basis basis = Torso.TargetBasis;
        Vector3 top = Torso.NominalTarget + Vector3.Up * ShoulderRise;
        Vector3 shoulderL = top - basis.X * ShoulderHalfWidth;
        Vector3 shoulderR = top + basis.X * ShoulderHalfWidth;
        Vector3 excessL = Excess(shoulderL, HandLeft);
        Vector3 excessR = Excess(shoulderR, HandRight);
        Vector3 lean = excessL.LengthSquared() >= excessR.LengthSquared()
            ? excessL + Excess(shoulderR + excessL, HandRight)
            : excessR + Excess(shoulderL + excessR, HandLeft);

        Vector3 horizontal = new Vector3(lean.X, 0f, lean.Z).LimitLength(ReachBudgetMax);
        lean = new Vector3(horizontal.X, Mathf.Clamp(lean.Y, ReachVerticalMin, ReachVerticalMax), horizontal.Z);
        _torsoOffset = _torsoOffset.Lerp(lean, 1f - Mathf.Exp(-dt * LeanSmoothRate));
        Torso.ReachOffset = _torsoOffset;
    }

    private Vector3 Excess(Vector3 shoulder, RigidBody3D hand)
    {
        if (hand == null)
            return Vector3.Zero;
        Vector3 toWrist = WristPoint(hand) - shoulder;
        float dist = toWrist.Length();
        float limit = (UpperLen + ForeLen) * LeanReachFactor;
        return dist > limit ? toWrist * ((dist - limit) / dist) : Vector3.Zero;
    }

    private static Vector3 WristPoint(RigidBody3D hand) =>
        hand.GlobalPosition + hand.GlobalTransform.Basis * new Vector3(0f, 0f, WristOffset);

    // Analytic 2-bone IK: soft-IK clamp, then pick the elbow on its circle. The elbow
    // direction is mostly HAND-DRIVEN — the forearm continues behind the wrist along the
    // hand's (aim pose) forward axis, which is where a real forearm is; that gives outward
    // elbows for wraps/hugs and keeps the forearm out of held objects. A torso-frame
    // "down/back/out" pole fills in where the hand axis says nothing (arm pointing straight
    // along the hand). The pole is smoothed by ROTATING it at a bounded rate around the arm
    // axis, so it can never pass through zero and flip.
    private void SolveArm(int index, Vector3 side, RigidBody3D hand, Node3D orient, float delta)
    {
        Vector3 shoulder = _shoulder[index];
        Vector3 rawTarget = hand != null ? WristPoint(hand) : shoulder + _torsoBasis * new Vector3(0f, -0.25f, -0.3f);

        float reach = UpperLen + ForeLen;
        Vector3 shoulderToTarget = rawTarget - shoulder;
        if (shoulderToTarget.LengthSquared() < 0.0001f)
            shoulderToTarget = _torsoBasis * new Vector3(0f, -0.25f, -0.3f);

        float rawDist = shoulderToTarget.Length();
        float d = Mathf.Clamp(rawDist, 0.01f, reach * SoftIKFactor);
        Vector3 dir = shoulderToTarget / rawDist;
        Vector3 wrist = shoulder + dir * d;

        // Elbow circle: center on the arm axis, radius from the law of cosines.
        float along = (UpperLen * UpperLen - ForeLen * ForeLen + d * d) / (2f * d);
        float radius = Mathf.Sqrt(Mathf.Max(UpperLen * UpperLen - along * along, 0f));
        Vector3 center = shoulder + dir * along;

        // Reference = torso-frame down/back/out projected off the arm axis. The elbow is
        // stored as an ANGLE around the arm axis relative to it, clamped to ±110° and
        // rate-limited — it can never swing through the inward/up side (the 180° flip).
        Vector3 reference = Vector3.Down * 2f + _torsoBasis * Vector3.Back * 0.8f + side * 0.3f;
        reference -= dir * reference.Dot(dir);
        if (reference.LengthSquared() > 0.06f)
            _lastRef[index] = reference.Normalized();
        else
        {
            // Arm pointing along the reference: keep the previous one, re-projected.
            reference = _lastRef[index] - dir * _lastRef[index].Dot(dir);
            reference = reference.LengthSquared() > 1e-8f ? reference.Normalized() : dir.Cross(_torsoBasis.X).Normalized();
        }
        Vector3 refN = reference.Normalized();
        Vector3 refB = dir.Cross(refN);

        float targetAngle = 0f;
        Basis handBasis = orient != null ? orient.GlobalTransform.Basis : hand?.GlobalTransform.Basis ?? Basis.Identity;
        if (orient != null || hand != null)
        {
            Vector3 handForward = (handBasis * Vector3.Forward).Normalized();
            Vector3 handElbow = wrist - handForward * ForeLen - center;
            handElbow -= dir * handElbow.Dot(dir);
            float confidence = Mathf.Clamp(handElbow.Length() / (0.5f * radius + 1e-4f), 0f, 1f);
            if (handElbow.LengthSquared() > 1e-8f)
                targetAngle = Mathf.Atan2(handElbow.Dot(refB), handElbow.Dot(refN)) * PoleHandWeight * confidence;
        }
        targetAngle = Mathf.Clamp(targetAngle, -PoleMaxAngle, PoleMaxAngle);
        float step = (targetAngle - _poleAngle[index]) * (1f - Mathf.Exp(-delta * PoleSmoothRate));
        _poleAngle[index] += Mathf.Clamp(step, -PoleMaxSpeed * delta, PoleMaxSpeed * delta);
        Vector3 pole = refN.Rotated(dir, _poleAngle[index]);

        Vector3 elbow = center + CapElbowHeight(pole, dir, center, radius, Mathf.Max(shoulder.Y, wrist.Y) + ElbowRiseAllowance) * radius;
        _elbowTarget[index] = elbow;
        _wristTarget[index] = elbow + (wrist - elbow).Normalized() * ForeLen;
    }

    // A real wrist flexes, so following the hand axis literally can throw the elbow above
    // the shoulder ("chicken wing"). Keep the elbow no higher than yMax by sliding it down
    // its circle, preserving its sideways heading.
    private static Vector3 CapElbowHeight(Vector3 pole, Vector3 dir, Vector3 center, float radius, float yMax)
    {
        if (center.Y + pole.Y * radius <= yMax)
            return pole;
        Vector3 vertical = Vector3.Up - dir * dir.Y;
        if (vertical.LengthSquared() < 1e-6f)
            return pole;  // arm vertical: the circle is horizontal, height can't change
        vertical = vertical.Normalized();
        Vector3 lateral = dir.Cross(vertical);
        float cosMax = Mathf.Clamp((yMax - center.Y) / (radius * vertical.Y), -1f, 1f);
        float sin = Mathf.Sqrt(1f - cosMax * cosMax);
        float sideSign = pole.Dot(lateral) >= 0f ? 1f : -1f;
        return (vertical * cosMax + lateral * (sin * sideSign)).Normalized();
    }
}