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;

    public 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;
    public const float HandArmSlack = 0.03f;

    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 ArmSolver[] _solver = { new(), new() };
    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);

    private bool _injected;

    public override void _Ready()
    {
        _injected = HeadSource != null;
        BodyProfile.Changed += OnProfileChanged;
        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)
        {
            // Main scene: wait for the T-pose profile. Headless spikes (HeadSource injected) use
            // the anthropometric estimate from their head height.
            BodyProfile profile = BodyProfile.Current ?? (_injected ? BodyProfile.Estimate(HeadSource.GlobalPosition.Y) : null);
            if (profile == null || Torso?.Body == null)
                return;
            Calibrate(profile);
        }

        _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);
        }
    }

    public override void _ExitTree()
    {
        BodyProfile.Changed -= OnProfileChanged;
    }

    // Recalibration: rebuild the segments at the new lengths on the next tick.
    private void OnProfileChanged(BodyProfile profile)
    {
        if (!Calibrated)
            return;
        Calibrated = false;
        for (int side = 0; side < 2; side++)
        {
            _upper[side]?.QueueFree();
            _fore[side]?.QueueFree();
            _upper[side] = _fore[side] = null;
        }
    }

    private void Calibrate(BodyProfile profile)
    {
        Calibrated = true;
        UpperLen = profile.UpperArm;
        ForeLen = profile.Forearm;
        ShoulderHalfWidth = profile.ShoulderHalfWidth;

        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: body profile {profile}; 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);
        Basis? handBasis = orient != null ? orient.GlobalTransform.Basis : hand?.GlobalTransform.Basis;
        _solver[index].Solve(shoulder, rawTarget, _torsoBasis, side, handBasis, UpperLen, ForeLen, delta);
        _elbowTarget[index] = _solver[index].Elbow;
        _wristTarget[index] = _solver[index].Wrist;
    }
}