using UnityEngine;

namespace Unheimlich.Features.OctoArms
{

public class OctoIkArmsImproved : MonoBehaviour
{
    // Stores raycast hit information
    private class RaycastInfos
    {
        public bool HasHit;
        public float Distance;
        public float MaxDistance;
        public Vector3 Position;
        public Vector3 Normal;
        public Transform Transform;
    }

    [System.Serializable]
    private class RaycastRig
    {
        public Transform Root;
        public float Distance = 1f;
    }

    [System.Serializable]
    private class ArmRig
    {
        public bool Enabled = true;
        public OctoIkLimb Arm;
        public Vector3 RotationOffset;
        public RaycastRig[] Raycasts;
        
        [Header("Equipment System")]
        public ArmEquipment Equipment;

        [HideInInspector] public Vector3 TargetOriginPos;
        [HideInInspector] public Quaternion TargetOriginRot;
        [HideInInspector] public Vector3 PoleOriginPos;

        public void Init()
        {
            if (Arm)
            {
                if (Arm.Target)
                {
                    TargetOriginPos = Arm.Target.localPosition;
                    TargetOriginRot = Arm.Target.localRotation;
                }

                if (Arm.Pole)
                {
                    PoleOriginPos = Arm.Pole.localPosition;
                }
            }
            
            if (Equipment != null)
            {
                Equipment.Initialize();
            }
        }
        
        public bool IsEnvironmentalIkActive()
        {
            if (Equipment == null)
                return Enabled;
            
            return Enabled && Equipment.State == ArmState.Environmental;
        }
        
        public float GetEnvironmentalIkWeight()
        {
            if (Equipment == null)
                return 1f;
            
            return Equipment.IkWeight;
        }
    }

    [System.Serializable]
    private class Settings
    {
        [Header("Raycast")]
        public LayerMask LayerMask;

        [Header("Stickiness")]
        [Tooltip("Distance to start sticking (meters)")]
        [Range(0.1f, 2f)]
        public float StickStartDistance = 0.8f;

        [Tooltip("Distance to fully release (meters)")]
        [Range(0.5f, 3f)]
        public float StickReleaseDistance = 1.2f;

        [Header("Smoothing")]
        [Tooltip("Stick speed (units/sec)")]
        [Range(1f, 20f)]
        public float StickSpeed = 10f;

        [Tooltip("Release speed (units/sec)")]
        [Range(1f, 20f)]
        public float ReleaseSpeed = 6f;

        [Tooltip("Rotation speed (deg/sec)")]
        [Range(90f, 720f)]
        public float RotationSpeed = 360f;

        [Header("Arm Selection")]
        [Tooltip("Can both arms stick at the same time?")]
        public bool AllowBothArms = true;

        [Tooltip("Max surface angle to accept sticking (degrees)")]
        [Range(0f, 90f)]
        public float MaxSurfaceAngle = 60f;

        [Header("Temporal Smoothing")]
        [Tooltip("Frames before releasing (prevents flickering)")]
        [Range(0, 10)]
        public int StabilityFrames = 3;
    }

    [SerializeField] private ArmRig _leftArm;
    [SerializeField] private ArmRig _rightArm;
    [SerializeField] private Settings _settings;

    private ArmRig[] _arms;
    private RaycastInfos[][] _raycasts = new RaycastInfos[2][];
    private RaycastInfos[] _nearestHitRaycasts = new RaycastInfos[2];
    private int[] _noHitFrameCount = new int[2];
    private float[] _currentStickiness = new float[2];

    void Awake()
    {
        _arms = new ArmRig[2] { _leftArm, _rightArm };
        foreach (ArmRig armRig in _arms)
            armRig.Init();

        // Initialize raycast arrays for both arms
        _raycasts = new RaycastInfos[_arms.Length][];
        for (int i = 0; i < _arms.Length; i++)
        {
            _raycasts[i] = new RaycastInfos[_arms[i].Raycasts.Length];
            for (int u = 0; u < _raycasts[i].Length; u++)
                _raycasts[i][u] = new RaycastInfos() { MaxDistance = 1 };
        }

        _nearestHitRaycasts = new RaycastInfos[_arms.Length];
        for (int i = 0; i < _nearestHitRaycasts.Length; i++)
            _nearestHitRaycasts[i] = new RaycastInfos() { MaxDistance = 1 };

        _noHitFrameCount = new int[_arms.Length];
        _currentStickiness = new float[_arms.Length];
    }

    void OnDrawGizmos()
    {
        if (Application.isPlaying || !enabled) return;

        foreach (ArmRig armRig in new ArmRig[2] { _leftArm, _rightArm })
        {
            if (armRig.Raycasts == null) continue;

            for (int i = 0; i < armRig.Raycasts.Length; i++)
            {
                RaycastRig raycastRig = armRig.Raycasts[i];
                if (raycastRig.Root == null) continue;

                Debug.DrawRay(raycastRig.Root.position, raycastRig.Root.TransformDirection(Vector3.forward) * raycastRig.Distance, Color.cyan);
            }
        }
    }

    void Update()
    {
        PerformRaycasts();
        FindNearestHits();
        UpdateArmPositions();
        UpdateEquippedPoses();
        UpdateHideShowAnimation();
        
        if (_keepHidden && !_isShowing)
        {
            _leftArm.Arm.Target.localPosition = _originalLeftTargetPos + _hideOffset;
            _rightArm.Arm.Target.localPosition = _originalRightTargetPos + _hideOffset;
        }
    }
    
    private void UpdateEquippedPoses()
    {
        for (int i = 0; i < _arms.Length; i++)
        {
            ArmRig armRig = _arms[i];
            if (armRig.Equipment != null)
            {
                armRig.Equipment.UpdateEquippedPose();
            }
        }
    }

    // Cast all configured raycasts and store results
    private void PerformRaycasts()
    {
        for (int i = 0; i < _arms.Length; i++)
        {
            ArmRig armRig = _arms[i];

            if (!armRig.IsEnvironmentalIkActive() || armRig.Raycasts == null)
                continue;

            for (int u = 0; u < armRig.Raycasts.Length; u++)
            {
                RaycastRig raycastRig = armRig.Raycasts[u];
                if (raycastRig.Root == null) continue;

                RaycastInfos raycastInfos = _raycasts[i][u];
                RaycastHit hit;

                Vector3 origin = raycastRig.Root.position;
                Vector3 direction = raycastRig.Root.TransformDirection(Vector3.forward);

                if (Physics.Raycast(origin, direction, out hit, raycastRig.Distance, _settings.LayerMask))
                {
                    Debug.DrawRay(origin, direction * hit.distance, Color.green);

                    raycastInfos.HasHit = true;
                    raycastInfos.Distance = hit.distance;
                    raycastInfos.MaxDistance = raycastRig.Distance;
                    raycastInfos.Position = hit.point;
                    raycastInfos.Normal = hit.normal;
                    raycastInfos.Transform = hit.transform;
                }
                else
                {
                    Debug.DrawRay(origin, direction * raycastRig.Distance, Color.red);

                    raycastInfos.HasHit = false;
                    raycastInfos.Distance = raycastRig.Distance;
                    raycastInfos.MaxDistance = raycastRig.Distance;
                    raycastInfos.Position = Vector3.zero;
                    raycastInfos.Normal = Vector3.zero;
                    raycastInfos.Transform = null;
                }
            }
        }
    }

    // Find the closest raycast hit for each arm
    private void FindNearestHits()
    {
        for (int i = 0; i < _arms.Length; i++)
        {
            ArmRig armRig = _arms[i];
            RaycastInfos nearest = _nearestHitRaycasts[i];

            nearest.HasHit = false;
            float closestDistance = float.MaxValue;

            if (armRig.Raycasts == null) continue;

            // Find the closest hit among all raycasts for this arm
            for (int u = 0; u < armRig.Raycasts.Length; u++)
            {
                RaycastInfos raycastInfos = _raycasts[i][u];

                if (raycastInfos.HasHit && raycastInfos.Distance < closestDistance)
                {
                    closestDistance = raycastInfos.Distance;

                    nearest.HasHit = true;
                    nearest.Distance = raycastInfos.Distance;
                    nearest.MaxDistance = raycastInfos.MaxDistance;
                    nearest.Position = raycastInfos.Position;
                    nearest.Normal = raycastInfos.Normal;
                    nearest.Transform = raycastInfos.Transform;
                }
            }

            // Track frames without hits for temporal smoothing
            if (!nearest.HasHit)
            {
                _noHitFrameCount[i]++;
            }
            else
            {
                _noHitFrameCount[i] = 0;
            }
        }
    }

    // Update arm positions based on raycast hits
    private void UpdateArmPositions()
    {
        for (int i = 0; i < _arms.Length; i++)
        {
            ArmRig armRig = _arms[i];
            if (armRig.Arm == null || armRig.Arm.Target == null) continue;
            
            if (!armRig.IsEnvironmentalIkActive())
                continue;

            RaycastInfos raycastInfos = _nearestHitRaycasts[i];

            Vector3 targetPos = armRig.TargetOriginPos;
            Quaternion targetRot = armRig.TargetOriginRot;

            bool shouldStick = false;

            if (raycastInfos.HasHit && _noHitFrameCount[i] < _settings.StabilityFrames)
            {
                if (IsValidSurface(raycastInfos, i))
                {
                    targetPos = armRig.Arm.Target.parent.InverseTransformPoint(raycastInfos.Position);
                    
                    targetRot = Quaternion.LookRotation(
                        -armRig.Arm.Target.parent.InverseTransformDirection(raycastInfos.Normal),
                        transform.up
                    ) * Quaternion.Euler(armRig.RotationOffset);

                    shouldStick = true;
                }
            }

            float targetStickiness = shouldStick ? 1f : 0f;
            float stickinessSpeed = shouldStick ? _settings.StickSpeed : _settings.ReleaseSpeed;
            _currentStickiness[i] = Mathf.MoveTowards(_currentStickiness[i], targetStickiness, stickinessSpeed * Time.deltaTime);

            float distanceFactor = 0f;
            if (shouldStick && raycastInfos.HasHit)
            {
                distanceFactor = 1f - Mathf.InverseLerp(_settings.StickStartDistance, _settings.StickReleaseDistance, raycastInfos.Distance);
                distanceFactor = Mathf.Clamp01(distanceFactor);
            }

            float finalStickiness = _currentStickiness[i] * distanceFactor * armRig.GetEnvironmentalIkWeight();

            float posSpeed = Mathf.Lerp(_settings.ReleaseSpeed, _settings.StickSpeed, finalStickiness);
            float rotSpeed = _settings.RotationSpeed * Mathf.Deg2Rad;

            float tPos = 1f - Mathf.Exp(-posSpeed * Time.deltaTime);
            float tRot = 1f - Mathf.Exp(-rotSpeed * Time.deltaTime);

            armRig.Arm.Target.localPosition = Vector3.Lerp(armRig.Arm.Target.localPosition, targetPos, tPos);
            armRig.Arm.Target.localRotation = Quaternion.Slerp(armRig.Arm.Target.localRotation, targetRot, tRot);
        }
    }

    // Check if the surface is valid for sticking
    private bool IsValidSurface(RaycastInfos raycast, int armIndex)
    {
        if (raycast.Transform == null)
            return false;

        // Check surface angle (reject surfaces too angled)
        float surfaceAngle = Vector3.Angle(-raycast.Normal, transform.forward);
        if (surfaceAngle > _settings.MaxSurfaceAngle)
            return false;

        // If not allowing both arms, check if arm is on correct side
        if (!_settings.AllowBothArms)
        {
            Vector3 hitLocalPos = transform.InverseTransformPoint(raycast.Position);
            bool isLeftSide = hitLocalPos.x < 0;
            bool isLeftArm = armIndex == 0;

            if (isLeftSide != isLeftArm)
                return false;
        }

        return true;
    }
    
    public void EquipObjectOnArm(int armIndex, GameObject objectToEquip, System.Action onComplete = null)
    {
        if (armIndex < 0 || armIndex >= _arms.Length)
        {
            Debug.LogError($"[OctoIkArms] Invalid arm index: {armIndex}");
            return;
        }
        
        ArmRig armRig = _arms[armIndex];
        
        if (armRig.Equipment == null)
        {
            Debug.LogError($"[OctoIkArms] Arm {armIndex} has no Equipment system configured!");
            return;
        }
        
        armRig.Equipment.StartTransitionToEquipped(objectToEquip, onComplete);
    }
    
    public void UnequipObjectFromArm(int armIndex, System.Action onComplete = null)
    {
        if (armIndex < 0 || armIndex >= _arms.Length)
        {
            Debug.LogError($"[OctoIkArms] Invalid arm index: {armIndex}");
            return;
        }
        
        ArmRig armRig = _arms[armIndex];
        
        if (armRig.Equipment == null)
        {
            Debug.LogError($"[OctoIkArms] Arm {armIndex} has no Equipment system configured!");
            return;
        }
        
        armRig.Equipment.StartTransitionToEnvironmental(onComplete);
    }
    
    public ArmState GetArmState(int armIndex)
    {
        if (armIndex < 0 || armIndex >= _arms.Length)
            return ArmState.Environmental;
        
        if (_arms[armIndex].Equipment == null)
            return ArmState.Environmental;
        
        return _arms[armIndex].Equipment.State;
    }
    
    public bool IsArmInEnvironmentalMode(int armIndex)
    {
        return GetArmState(armIndex) == ArmState.Environmental;
    }
    
    public ArmEquipment GetArmEquipment(int armIndex)
    {
        if (armIndex < 0 || armIndex >= _arms.Length)
            return null;
        
        return _arms[armIndex].Equipment;
    }
    
    private SkinnedMeshRenderer _armsRenderer;
    private bool _armsHidden = false;
    private Vector3 _originalLeftTargetPos;
    private Vector3 _originalRightTargetPos;
    private float _hideAnimationProgress = 0f;
    private bool _isHiding = false;
    private bool _isShowing = false;
    private bool _keepHidden = false;
    private float _hideShowSpeed = 2f;
    private Vector3 _hideOffset = new Vector3(0, -0.5f, 0);
    
    public void HideArms(float speed = 2f)
    {
        if (_armsHidden)
            return;
        
        if (_armsRenderer == null)
            _armsRenderer = GetComponentInChildren<SkinnedMeshRenderer>();
        
        _originalLeftTargetPos = _leftArm.Arm.Target.localPosition;
        _originalRightTargetPos = _rightArm.Arm.Target.localPosition;
        
        _hideShowSpeed = speed;
        _hideAnimationProgress = 0f;
        _isHiding = true;
        _isShowing = false;
        _armsHidden = true;
        _keepHidden = false;
    }
    
    public void ShowArms(float speed = 2f)
    {
        if (!_armsHidden && !_isHiding)
            return;
        
        if (_armsRenderer == null)
            _armsRenderer = GetComponentInChildren<SkinnedMeshRenderer>();
        
        if (_armsRenderer != null)
            _armsRenderer.enabled = true;
        
        _hideShowSpeed = speed;
        _hideAnimationProgress = 0f;
        _isShowing = true;
        _isHiding = false;
    }
    
    private void UpdateHideShowAnimation()
    {
        if (!_isHiding && !_isShowing)
            return;
        
        _hideAnimationProgress += Time.deltaTime * _hideShowSpeed;
        float t = Mathf.Clamp01(_hideAnimationProgress);
        
        if (_isHiding)
        {
            _leftArm.Arm.Target.localPosition = Vector3.Lerp(_originalLeftTargetPos, _originalLeftTargetPos + _hideOffset, t);
            _rightArm.Arm.Target.localPosition = Vector3.Lerp(_originalRightTargetPos, _originalRightTargetPos + _hideOffset, t);
            
            if (t >= 1f)
            {
                _isHiding = false;
                _keepHidden = true;
                if (_armsRenderer != null)
                    _armsRenderer.enabled = false;
            }
        }
        else if (_isShowing)
        {
            _leftArm.Arm.Target.localPosition = Vector3.Lerp(_originalLeftTargetPos + _hideOffset, _originalLeftTargetPos, t);
            _rightArm.Arm.Target.localPosition = Vector3.Lerp(_originalRightTargetPos + _hideOffset, _originalRightTargetPos, t);
            
            if (t >= 1f)
            {
                _isShowing = false;
                _armsHidden = false;
                _keepHidden = false;
                _leftArm.Arm.Target.localPosition = _originalLeftTargetPos;
                _rightArm.Arm.Target.localPosition = _originalRightTargetPos;
            }
        }
    }
}

}