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 ReachPushSpeed = 0.8f;
    private const float ReachDecaySpeed = 0.4f;
    private const float ReachDecayFactor = 0.92f;
    private const float WristOffset = 0.06f;
    private const float SoftIKFactor = 0.98f;
    private const float PoleHandWeight = 0.25f;
    private const float PoleSmoothRate = 10f;

    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;

    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 Vector3[] _lastPole = new Vector3[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"));
        }
    }

    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;

        Vector3 push = SolveArm(0, -_torsoBasis.X, HandLeft, dt, out bool decayLeft);
        push += SolveArm(1, _torsoBasis.X, HandRight, dt, out bool decayRight);
        UpdateReachBudget(push, decayLeft || decayRight, 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));
        }
        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
        };
        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);

    private void UpdateReachBudget(Vector3 push, bool decay, float dt)
    {
        if (push != Vector3.Zero)
            _torsoOffset += push * (dt * ReachPushSpeed);
        else if (decay)
            _torsoOffset = _torsoOffset.MoveToward(Vector3.Zero, ReachDecaySpeed * dt);

        Vector3 horizontal = new Vector3(_torsoOffset.X, 0f, _torsoOffset.Z).LimitLength(ReachBudgetMax);
        _torsoOffset = new Vector3(horizontal.X, Mathf.Clamp(_torsoOffset.Y, ReachVerticalMin, ReachVerticalMax), horizontal.Z);
        Torso.ReachOffset = _torsoOffset;
    }

    // Analytic 2-bone IK (moved from BodyRig): soft-IK clamp, torso-frame elbow pole with the
    // shoulder→hand axis projected out, hand-forward push, exp-smoothed pole.
    private Vector3 SolveArm(int index, Vector3 side, RigidBody3D hand, float delta, out bool decay)
    {
        Vector3 shoulder = _shoulder[index];
        Vector3 rawTarget = hand != null
            ? hand.GlobalPosition + hand.GlobalTransform.Basis * new Vector3(0f, 0f, WristOffset)
            : 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 clampedTarget = shoulder + dir * d;

        Vector3 push = Vector3.Zero;
        decay = false;
        if (rawDist > reach)
            push = (rawTarget - clampedTarget).Normalized();
        else if (rawDist < reach * ReachDecayFactor)
            decay = true;

        float cosA = (UpperLen * UpperLen + d * d - ForeLen * ForeLen) / (2f * UpperLen * d);
        float angle = Mathf.Acos(Mathf.Clamp(cosA, -1f, 1f));

        Vector3 handForward = hand != null ? hand.GlobalTransform.Basis * Vector3.Forward : Vector3.Zero;
        Vector3 pole = Vector3.Down * 2f + _torsoBasis * Vector3.Back * 0.8f + side * 0.3f;
        pole -= dir * pole.Dot(dir);
        pole -= handForward * PoleHandWeight;
        Vector3 poleDir;
        if (pole.LengthSquared() < 0.0001f)
            poleDir = _lastPole[index] == Vector3.Zero ? Vector3.Down : _lastPole[index];
        else
        {
            poleDir = _lastPole[index].Lerp(pole, 1f - Mathf.Exp(-delta * PoleSmoothRate));
            if (poleDir.LengthSquared() < 0.0001f)
                poleDir = _lastPole[index];
            poleDir = poleDir.Normalized();
            _lastPole[index] = poleDir;
        }

        Vector3 axis = dir.Cross(poleDir);
        if (axis.LengthSquared() < 0.0001f)
            axis = dir.Cross(Vector3.Up);
        if (axis.LengthSquared() < 0.0001f)
            axis = Vector3.Right;
        axis = axis.Normalized();

        Vector3 elbow = shoulder + dir.Rotated(axis, angle) * UpperLen;
        _elbowTarget[index] = elbow;
        _wristTarget[index] = elbow + (clampedTarget - elbow).Normalized() * ForeLen;
        return push;
    }
}