using System.Collections;
using System.Collections.Generic;
using System.Linq;
using UnityEngine;

public partial class RobotController : MonoBehaviour
{
    public enum StepType { Move, MoveJoints, Pick, Drop, Attach, Detach, Action, Delay, CustomAction }

    [System.Serializable]
    public class Step
    {
        public StepType Type;
        public Frame[] Path;
        public System.Action Action;
        public float Delay;
        public float[] JointAngles;
        public float MoveDuration;

        public Step(StepType type, Frame[] path = null, System.Action action = null, float delay = 0f, float[] jointAngles = null, float moveDuration = 0)
        {
            Type = type;
            Path = path;
            Action = action;
            Delay = delay;
            JointAngles = jointAngles;
            MoveDuration = moveDuration;
        }
    }

    [Header("Joints Setup")]
    public RobotJoint[] joints;

    [Header("End Effector")]
    public Transform endEffector;

    [Header("Targets")]
    public Transform target1;
    public Transform target2;
    public Transform target3;
    public int HoldFrames = 30;

    [Header("Settings")]
    public int stepsPerMove = 80;
    public float stepDelay = 0.02f;
    public float liftHeight = 0.2f;
    [Range(0f, 1f)] public float applyAlpha = 0.35f;

    [Header("Animation")]
    public Animator armAnimator;

    [Header("Objects")]
    public Transform objectPrefab;
    public Transform currentObject;

    [Header("Stack Settings")]
    public Transform objectStacktran;
    public int stackCount = 12;
    public float stackSpacingY = 0.1f;
    public float shiftDelay = 3f;
    public float shiftDuration = 0.5f;

    SolverCCD solver;
    MotionPlanner planner;
    PickDropHandler pickDrop;
    List<Step> steps;
    float[] lastAngles;

    private readonly List<Transform> objectStack = new();
    private Transform nextTop;

    public Transform LastStackObject { get; set; }    // đối tượng ABB vừa đặt ở target3
    public System.Action OnCompleted;

    void Awake()
    {
        if (joints == null || joints.Length == 0 || endEffector == null)
        {
            Debug.LogError("RobotController: Thiếu joints hoặc endEffector!");
            enabled = false;
            return;
        }

        foreach (var j in joints) j.Init();
        solver = new SolverCCD(joints, endEffector);
        planner = new MotionPlanner();
        pickDrop = new PickDropHandler();
        pickDrop.Init(armAnimator);

        lastAngles = new float[joints.Length];
        for (int i = 0; i < joints.Length; i++) lastAngles[i] = joints[i].GetValue();

        InitObjectStack();
    }

    // ========== STACK ==========
    void InitObjectStack()
    {
        if (objectPrefab == null || objectStacktran == null) return;

        objectStack.Clear();
        objectStacktran.GetPositionAndRotation(out Vector3 basePos, out Quaternion baseRot);
        for (int i = 0; i < stackCount; i++)
        {
            Vector3 pos = basePos + Vector3.up * (stackSpacingY * i);
            Transform obj = Instantiate(objectPrefab, pos, baseRot, objectStacktran);
            obj.name = $"StackObject_{i}";
            objectStack.Add(obj);
        }
        nextTop = null;
    }

    Transform PeekTop() => objectStack.Count > 0 ? objectStack.Last() : null;
    Transform PopTop()
    {
        if (objectStack.Count == 0) return null;
        var t = objectStack.Last();
        objectStack.RemoveAt(objectStack.Count - 1);
        return t;
    }

    IEnumerator ShiftStackUpAfterDelay()
    {
        yield return new WaitForSeconds(shiftDelay);
        foreach (var obj in objectStack)
        {
            Vector3 targetPos = obj.position + Vector3.up * stackSpacingY;
            StartCoroutine(MoveToPosition(obj, targetPos, shiftDuration));
        }
        yield return new WaitForSeconds(shiftDuration);

        if (currentObject == null) currentObject = PeekTop();
        else nextTop = PeekTop();
    }

    IEnumerator MoveToPosition(Transform obj, Vector3 target, float duration)
    {
        if (obj == null) yield break;
        Vector3 start = obj.position;
        float elapsed = 0f;

        while (elapsed < duration)
        {
            elapsed += Time.deltaTime;
            float t = Mathf.Clamp01(elapsed / duration);
            obj.position = Vector3.Lerp(start, target, t);
            yield return null;
        }
        obj.position = target;
    }

    // ========== MAIN FLOW ==========
    public IEnumerator RunStepsExternal()
    {
        BuildSteps();

        foreach (var step in steps)
        {
            switch (step.Type)
            {
                case StepType.Move:
                    if (step.Path != null) yield return DoMove(step.Path);
                    break;

                case StepType.MoveJoints:
                    yield return StartCoroutine(DoMoveToJoints(step.JointAngles, step.MoveDuration));
                    break;

                case StepType.Pick:
                    yield return pickDrop.PlayPick(this);
                    break;

                case StepType.Drop:
                    yield return pickDrop.PlayDrop(this);
                    break;

                case StepType.Attach:
                    if (currentObject == null)
                    {
                        currentObject = PopTop();
                        if (currentObject == null) break;
                    }
                    pickDrop.AttachObjToAmr(endEffector, currentObject);
                    if (objectStack.Count > 0) StartCoroutine(ShiftStackUpAfterDelay());
                    break;

                case StepType.Detach:
                    pickDrop.DetachObj(currentObject);
                    if (currentObject != null) currentObject.rotation = Quaternion.identity;
                    break;

                case StepType.Action:
                case StepType.CustomAction:
                    step.Action?.Invoke();
                    break;

                case StepType.Delay:
                    yield return new WaitForSeconds(step.Delay);
                    break;
            }
        }
        OnCompleted?.Invoke();
    }

    void BuildSteps()
    {
        steps = new List<Step>();

        float[] homeAngles = GetHomeJointAngles();

        Frame t1 = new(target1.position, target1.rotation);
        Frame t2 = new(target2.position, target2.rotation);
        Frame lift1 = LiftFrame(t1, liftHeight);
        Frame lift2 = LiftFrame(t2, liftHeight);

        // Về home = dùng MoveJoints
        steps.Add(new Step(StepType.MoveJoints, jointAngles: homeAngles));

        // Pick từ target1
        steps.Add(new Step(StepType.Move, planner.GeneratePath(new Frame(endEffector.position, endEffector.rotation), lift1, stepsPerMove).ToArray()));
        steps.Add(new Step(StepType.Move, planner.GeneratePath(lift1, t1, stepsPerMove).ToArray()));
        steps.Add(new Step(StepType.Attach));
        steps.Add(new Step(StepType.Move, planner.RepeatFrame(t1, HoldFrames)));
        steps.Add(new Step(StepType.Move, planner.GeneratePath(t1, lift1, stepsPerMove).ToArray()));
        steps.Add(new Step(StepType.MoveJoints, jointAngles: homeAngles) { MoveDuration = 1f });

        // Place tại target2
        steps.Add(new Step(StepType.Move, planner.GeneratePath(new Frame(endEffector.position, endEffector.rotation), lift2, stepsPerMove).ToArray()));
        steps.Add(new Step(StepType.Move, planner.GeneratePath(lift2, t2, stepsPerMove).ToArray()));
        steps.Add(new Step(StepType.Detach));
        steps.Add(new Step(StepType.Move, planner.RepeatFrame(t2, HoldFrames)));

        // Animate rơi tự do (rơi xuống target3)
        steps.Add(new Step(StepType.CustomAction, action: () => { StartCoroutine(MoveObject(target2, target3, 0.2f)); }) { MoveDuration = 0.75f });

        // Quay về home
        steps.Add(new Step(StepType.Move, planner.GeneratePath(t2, lift2, stepsPerMove).ToArray()));
        steps.Add(new Step(StepType.MoveJoints, jointAngles: homeAngles) { MoveDuration = 1f });
    }

    // ========== MOTION ==========
    IEnumerator DoMove(Frame[] path)
    {
        foreach (var f in path)
        {
            float[] solved = solver.SolveIK(f);
            for (int i = 0; i < joints.Length && i < solved.Length; i++)
            {
                float a = Mathf.Lerp(lastAngles[i], solved[i], applyAlpha);
                lastAngles[i] = a;
                joints[i].SetValue(a);
            }
            yield return new WaitForSeconds(stepDelay);
        }
    }

    IEnumerator DoMoveToJoints(float[] targetAngles, float duration = 0.75f)
    {
        if (targetAngles == null) yield break;

        int frames = stepsPerMove;
        if (duration > 0f) // nếu truyền duration, thì tính lại StepDelay
            stepDelay = duration / frames;

        float[] startAngles = new float[joints.Length];
        for (int i = 0; i < joints.Length; i++) startAngles[i] = joints[i].GetValue();

        for (int f = 0; f < frames; f++)
        {
            float t = (float)f / (frames - 1);
            for (int j = 0; j < joints.Length; j++)
            {
                float angle = Mathf.Lerp(startAngles[j], targetAngles[j], t);
                joints[j].SetValue(angle);
                lastAngles[j] = angle;
            }
            yield return new WaitForSeconds(stepDelay);
        }
    }
    float[] GetHomeJointAngles()
    {
        float[] home = new float[joints.Length];
        for (int i = 0; i < joints.Length; i++) home[i] = joints[i].GetValue();
        return home;
    }

    Frame LiftFrame(Frame f, float lift)
    {
        return new Frame(new Vector3(f.Position.x, f.Position.y + lift, f.Position.z), f.Rotation);
    }

    // ========== DROP ANIMATION ==========
    public IEnumerator MoveObject(Transform from, Transform to, float duration)
    {
        if (currentObject == null) yield break;

        from.GetPositionAndRotation(out Vector3 startPos, out Quaternion startRot);
        float elapsed = 0f;

        while (elapsed < duration)
        {
            elapsed += Time.deltaTime;
            float t = Mathf.Clamp01(elapsed / duration);
            currentObject.SetPositionAndRotation(
                Vector3.Lerp(startPos, to.position, t),
                Quaternion.Slerp(startRot, to.rotation, t)
            );
            yield return null;
        }

        // Sau khi thả xong tại target3
        LastStackObject = currentObject;
        // thông tin góc lastStackObject Rotation = Quaternion.identity;
        currentObject = null;
    }
}