using UnityEngine;
using UnityEngine.UI;
using System.Collections.Generic;

[ExecuteAlways]
[DisallowMultipleComponent]
[RequireComponent(typeof(RectTransform))]
public sealed class RectChildrenHeightFitter : MonoBehaviour
{
    public bool includeInactiveChildren;
    public float extraTopPadding;
    public float extraBottomPadding;
    public bool syncLayoutElement = true;
    public bool preserveDirectChildrenTopLeft = true;

    private RectTransform rectTransform;
    private LayoutElement layoutElement;
    private float lastAppliedHeight = -1f;
    private readonly Dictionary<int, ChildOffsets> childOffsets = new Dictionary<int, ChildOffsets>();

    private struct ChildOffsets
    {
        public float LeftOffset;
        public float TopOffset;
    }

    private void Awake()
    {
        CacheComponents();
        Refresh();
    }

    private void OnEnable()
    {
        CacheComponents();
        Refresh();
    }

    private void OnValidate()
    {
        CacheComponents();
        Refresh();
    }

    private void LateUpdate()
    {
        Refresh();
    }

    private void CacheComponents()
    {
        if (rectTransform == null)
        {
            rectTransform = GetComponent<RectTransform>();
        }

        if (syncLayoutElement && layoutElement == null)
        {
            layoutElement = GetComponent<LayoutElement>();
        }
    }

    private void Refresh()
    {
        if (rectTransform == null)
        {
            return;
        }

        CacheChildOffsetsIfNeeded();

        var hasBounds = false;
        var minY = float.PositiveInfinity;
        var maxY = float.NegativeInfinity;
        var corners = new Vector3[4];

        for (var i = 0; i < rectTransform.childCount; i++)
        {
            if (!(rectTransform.GetChild(i) is RectTransform child))
            {
                continue;
            }

            if (!includeInactiveChildren && !child.gameObject.activeSelf)
            {
                continue;
            }

            child.GetWorldCorners(corners);
            for (var c = 0; c < 4; c++)
            {
                var localPoint = rectTransform.InverseTransformPoint(corners[c]);
                minY = Mathf.Min(minY, localPoint.y);
                maxY = Mathf.Max(maxY, localPoint.y);
                hasBounds = true;
            }
        }

        var targetHeight = hasBounds
            ? Mathf.Max(0f, (maxY - minY) + extraTopPadding + extraBottomPadding)
            : 0f;

        if (!Mathf.Approximately(lastAppliedHeight, targetHeight))
        {
            rectTransform.SetSizeWithCurrentAnchors(RectTransform.Axis.Vertical, targetHeight);
            RestoreDirectChildPositions();

            if (syncLayoutElement)
            {
                if (layoutElement == null)
                {
                    layoutElement = GetComponent<LayoutElement>();
                }

                if (layoutElement != null)
                {
                    layoutElement.minHeight = -1f;
                    layoutElement.preferredHeight = targetHeight;
                    layoutElement.flexibleHeight = -1f;
                }
            }

            LayoutRebuilder.MarkLayoutForRebuild(rectTransform);
            if (rectTransform.parent is RectTransform parentRect)
            {
                LayoutRebuilder.MarkLayoutForRebuild(parentRect);
            }

            lastAppliedHeight = targetHeight;
        }
    }

    private void CacheChildOffsetsIfNeeded()
    {
        if (!preserveDirectChildrenTopLeft || rectTransform == null)
        {
            return;
        }

        if (childOffsets.Count == rectTransform.childCount)
        {
            var allPresent = true;
            for (var i = 0; i < rectTransform.childCount; i++)
            {
                if (!childOffsets.ContainsKey(rectTransform.GetChild(i).GetInstanceID()))
                {
                    allPresent = false;
                    break;
                }
            }

            if (allPresent)
            {
                return;
            }
        }

        childOffsets.Clear();
        var parentRect = rectTransform.rect;
        for (var i = 0; i < rectTransform.childCount; i++)
        {
            if (!(rectTransform.GetChild(i) is RectTransform child))
            {
                continue;
            }

            var localPivotPosition = (Vector2)child.localPosition;
            childOffsets[child.GetInstanceID()] = new ChildOffsets
            {
                LeftOffset = localPivotPosition.x - parentRect.xMin,
                TopOffset = parentRect.yMax - localPivotPosition.y
            };
        }
    }

    private void RestoreDirectChildPositions()
    {
        if (!preserveDirectChildrenTopLeft || rectTransform == null)
        {
            return;
        }

        var parentRect = rectTransform.rect;
        for (var i = 0; i < rectTransform.childCount; i++)
        {
            if (!(rectTransform.GetChild(i) is RectTransform child))
            {
                continue;
            }

            if (!childOffsets.TryGetValue(child.GetInstanceID(), out var offsets))
            {
                continue;
            }

            // Only handle fixed-anchor direct children. Stretched children should keep their own layout behavior.
            if (!Mathf.Approximately(child.anchorMin.x, child.anchorMax.x)
                || !Mathf.Approximately(child.anchorMin.y, child.anchorMax.y))
            {
                continue;
            }

            var desiredLocalPivot = new Vector2(
                parentRect.xMin + offsets.LeftOffset,
                parentRect.yMax - offsets.TopOffset);

            var anchorReference = new Vector2(
                Mathf.Lerp(parentRect.xMin, parentRect.xMax, child.anchorMin.x),
                Mathf.Lerp(parentRect.yMin, parentRect.yMax, child.anchorMin.y));

            child.anchoredPosition = desiredLocalPivot - anchorReference;
        }
    }
}