using System.Collections;
using System.Collections.Generic;
using System.Linq;
using System.Threading.Tasks;
using UnityEngine;
using static RobotController;

public partial class Robot2Controller : MonoBehaviour
{
    // ===== Enums =====
    public enum StepType { Move, Pick, Drop, AttachA, AttachB, DetachA, DetachB, Detach, CustomAction, Delay, LinearMove, MoveJoints, Parallel }
    public enum GraspMode { Top, Right, Left, Font }
    public enum EffectorType { A, B }   // NEW

    // ===== Nested Classes =====
    [System.Serializable]
    public class Step
    {
        public StepType Type;
        public Frame[] Path;
        public System.Action Action;
        public float DelayTime;
        public float[] JointAngles;
        public float MoveDuration;
        public List<IEnumerator> ParallelCoroutines;
        public EffectorType Effector;   // NEW

        public Step(StepType type, EffectorType effector = EffectorType.A)
        {
            Type = type;
            Effector = effector;
        }
    }

    public class HomeData
    {
        public float[] Angles;
        public Frame FrameAtZero;
        public Frame FrameAtAngles;

        public HomeData(float[] angles, Frame frameAtZero, Frame frameAtAngles)
        {
            Angles = angles;
            FrameAtZero = frameAtZero;
            FrameAtAngles = frameAtAngles;
        }
    }

    // ===== Public Fields =====
    [Header("Joints Setup")]
    public RobotJoint[] Joints;

    [Header("Phoi từ ABB (đang xử lý)")]
    public Transform currentPhoiFromAbbA;   
    public Transform currentPhoiFromAbbB;

    [Header("Targets")]
    public Transform target3;
    public Transform target4;
    public Transform target42;
    public Transform linearTarget4;
    public Transform target51;
    public Transform target52;
    public Transform target7;
    public Transform target8;
    public float LinearSpeed = 1f;
    public Transform LinearBase;
    public Transform LinearHome;

    [Header("Motion Settings")]
    public int Steps = 80;
    public float StepDelay = 0.02f;
    public float LiftHeight = 0.2f;
    public int HoldFrames1 = 30;
    [Range(0f, 1f)] public float ApplyAlpha = 0.35f;

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

    [Header("References")]
    public RobotController robot1;

    [Header("Phoi Settings (clones)")]
    public Transform PhoiForBPrefab;     // prefab gốc
    public Transform PhoiForBInstance;  // instance của mỗi vòng


    [Header("Phoi từ ABB (đang xử lý)")]
    public Transform currentPhoiFromAbb;

    [Header("End Effectors")]
    public Transform EndEffectorA;
    public Transform EndEffectorB;

    // ===== Private Fields =====
    private SolverCCD solverA;
    private SolverCCD solverB;
    private MotionPlanner planner;
    private List<Step> steps;
    private int currentStepIndex;
    private float[] lastAppliedAngles;
    private PickDropHandler pickDropHandler;

    private readonly Queue<Transform> phoiQueue = new();
    private readonly Queue<Transform> abbPhoiQueue = new();

    public bool IsPart2Done { get; private set; } = true;

    // ===== Unity Lifecycle =====
    void Start()
    {
        if (Joints.Length == 0 || EndEffectorA == null || EndEffectorB == null || armAnimator == null)
        {
            Debug.LogError("Thiếu cấu hình Robot2Controller!");
            enabled = false;
            return;
        }

        foreach (var j in Joints) j.Init();

        solverA = new SolverCCD(Joints, EndEffectorA);
        solverB = new SolverCCD(Joints, EndEffectorB);
        planner = new MotionPlanner();

        lastAppliedAngles = ReadCurrentAngles();

        pickDropHandler = new PickDropHandler();
        pickDropHandler.Init(armAnimator);

        //InitPhoiStack();
    }

    // ===== Public API =====
    public IEnumerator RunStepsExternal()
    {
        // spawn one phoi B at transfer location for this cycle (so AttachB has something)
        if (PhoiForBPrefab != null && target4 != null)
        {
            Quaternion zRot = Quaternion.Euler(0, 0, 90);
            currentPhoiFromAbbB = Instantiate(PhoiForBPrefab, target4.position, target4.rotation * zRot);
            currentPhoiFromAbbB.name = $"PhoiForB_{Time.frameCount}";
            Debug.Log($"[RunStepsExternal] Spawned {currentPhoiFromAbbB.name}");
        }

        BuildSteps();
        currentStepIndex = 0;

        while (currentStepIndex < steps.Count)
        {
            Step step = steps[currentStepIndex];

            switch (step.Type)
            {
                case StepType.Move:
                    yield return StartCoroutine(DoMove(step.Path, step.Effector));
                    break;
                case StepType.MoveJoints:
                    yield return StartCoroutine(DoMoveToJoints(step.JointAngles, step.MoveDuration));
                    break;
                case StepType.Pick:
                    yield return StartCoroutine(pickDropHandler.PlayPick(this));
                    break;
                case StepType.Drop:
                    yield return StartCoroutine(pickDropHandler.PlayDrop(this));
                    break;
                case StepType.AttachA:
                    HandleAttachA();
                    break;
                case StepType.DetachA:
                    HandleDetachA();
                    break;
                case StepType.AttachB:
                    HandleAttachB();
                    break;
                case StepType.DetachB:
                    HandleDetachB();
                    break;
                case StepType.Detach:
                    DA();
                    DB();
                    break;
                case StepType.CustomAction:
                    step.Action?.Invoke();
                    break;
                case StepType.Delay:
                    yield return new WaitForSeconds(step.DelayTime);
                    break;
                case StepType.LinearMove:
                    if (step.Path != null && step.Path.Length > 0)
                        yield return StartCoroutine(DoLinearMove(step.Path[0], step.DelayTime));
                    break;
                case StepType.Parallel:
                    if (step.ParallelCoroutines != null && step.ParallelCoroutines.Count > 0)
                        yield return StartCoroutine(RunParallel(step.ParallelCoroutines));
                    break;
            }

            currentStepIndex++;
        }

        // ✅ Sau khi kết thúc 1 vòng → xóa phôi của A
        if (currentPhoiFromAbbA != null) { Destroy(currentPhoiFromAbbA.gameObject); currentPhoiFromAbbA = null; }
    }

    public bool HasPhoiAtTarget3() => abbPhoiQueue.Count > 0;

    public void EnqueuePhoi(Transform phoi)
    {
        if (phoi == null)
        {
            Debug.LogWarning("EnqueuePhoi nhận null");
            return;
        }

        // Snap về target3
        phoi.SetPositionAndRotation(target3.position, target3.rotation);
        phoi.SetParent(null);

        abbPhoiQueue.Enqueue(phoi);
    }

    // ===== Attach / Detach =====
    void HandleAttachA()
    {
        if (abbPhoiQueue.Count == 0) return;

        currentPhoiFromAbbA = abbPhoiQueue.Dequeue();
        if (currentPhoiFromAbbA == null) return;

        pickDropHandler.AttachObjToAmr(EndEffectorA, currentPhoiFromAbbA);
        Debug.Log($"[AttachA] {currentPhoiFromAbbA.name} gắn vào EndEffectorA");
    }

    void HandleDetachA()
    {
        if (currentPhoiFromAbbA != null)
        {
            pickDropHandler.DetachObj(currentPhoiFromAbbA);
            Debug.Log($"[DetachA] Tháo {currentPhoiFromAbbA.name} khỏi EndEffectorA");
        }
    }
    void DA()
    {
        Destroy(currentPhoiFromAbbA.gameObject);   // ❌ Xóa luôn phôi
        currentPhoiFromAbbA = null;
    }

    void HandleAttachB()
    {
        if (currentPhoiFromAbbB == null)
        {
            Debug.LogWarning("[AttachB] currentPhoiFromAbbB NULL");
            return;
        }

        pickDropHandler.AttachObjToAmr(EndEffectorB, currentPhoiFromAbbB);
        Debug.Log($"[AttachB] {currentPhoiFromAbbB.name} gắn vào EndEffectorB");
    }

    void HandleDetachB()
    {
        if (currentPhoiFromAbbB != null)
        {
            pickDropHandler.DetachObj(currentPhoiFromAbbB);
            Debug.Log($"[DetachB] Tháo và Destroy {currentPhoiFromAbbB.name}");
        }
    }
    void DB()
    {
        Destroy(currentPhoiFromAbbB.gameObject);   // ❌ Xóa luôn phôi
        currentPhoiFromAbbB = null;
    }
    IEnumerator RunParallel(List<IEnumerator> coroutines)
    {
        var running = coroutines.Select(StartCoroutine).ToList();
        foreach (var r in running) yield return r;
    }

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

        HomeData homeAtLinearHome = GetHomeAtLinear(LinearHome);
        HomeData homeAtTarget4 = GetHomeAtLinear(linearTarget4, homeAtLinearHome.Angles);

        Frame t1 = CreateFrameFromTransform(target3, GraspMode.Top);
        Frame k1 = CreateLiftFrame(t1, LiftHeight, GraspMode.Top);

        Quaternion yOffset = Quaternion.Euler(0, -90, 0);
        Quaternion yOffsetv = Quaternion.Euler(0, 90, 0);   

        Frame t2 = CreateFrameFromTransform(target4, GraspMode.Right, yOffset);
        Frame t2v = CreateFrameFromTransform(target42, GraspMode.Right, yOffsetv);

        Frame t3 = CreateFrameFromTransform(target7);
        Frame t4 = CreateFrameFromTransform(target8);
        Frame k4 = CreateLiftFrame(t4, LiftHeight, GraspMode.Top);

        Frame k2 = CreateLiftFrame(t2, LiftHeight, GraspMode.Right);
        Frame k2v = CreateLiftFrame(t2v, LiftHeight, GraspMode.Right);


        // Về home
        steps.Add(new Step(StepType.MoveJoints)
        {
            JointAngles = CloneAngles(homeAtLinearHome.Angles),
            MoveDuration = 1f
        });

        // Pick object
        steps.Add(new Step(StepType.Move) { Path = planner.GeneratePath(new Frame(EndEffectorA.position, EndEffectorA.rotation), k1, Steps).ToArray() });
        steps.Add(new Step(StepType.Move) { Path = planner.GeneratePath(k1, t1, Steps).ToArray() });
        steps.Add(new Step(StepType.Move) { Path = planner.RepeatFrame(t1, HoldFrames1) });
        steps.Add(new Step(StepType.AttachA));
        steps.Add(new Step(StepType.Move) { Path = planner.GeneratePath(t1, k1, Steps).ToArray() });
        steps.Add(new Step(StepType.MoveJoints){JointAngles = CloneAngles(homeAtLinearHome.Angles),MoveDuration = 1f});
        steps.Add(new Step(StepType.LinearMove){Path = new[] { CreateFrameFromTransform(linearTarget4) },DelayTime = 3f});
        steps.Add(new Step(StepType.Move, EffectorType.B) { Path = planner.GeneratePath(homeAtTarget4.FrameAtZero, k2, Steps).ToArray() });

        steps.Add(new Step(StepType.Move, EffectorType.B) { Path = planner.GeneratePath(k2, t2, Steps).ToArray() });
        steps.Add(new Step(StepType.Move, EffectorType.B) { Path = planner.RepeatFrame(t2, HoldFrames1) });
        steps.Add(new Step(StepType.AttachB));
        steps.Add(new Step(StepType.Move, EffectorType.B) { Path = planner.GeneratePath(t2, k2, Steps).ToArray() });
        steps.Add(new Step(StepType.Move, EffectorType.B) { Path = planner.GeneratePath(k2, homeAtTarget4.FrameAtAngles, Steps).ToArray() });
        steps.Add(new Step(StepType.Move) { Path = planner.GeneratePath(homeAtTarget4.FrameAtAngles, k2v, Steps).ToArray() });
        steps.Add(new Step(StepType.Move) { Path = planner.GeneratePath(k2v, t2v, Steps).ToArray() });
        steps.Add(new Step(StepType.Move) { Path = planner.RepeatFrame(t2v, HoldFrames1) });
        steps.Add(new Step(StepType.DetachA));
        steps.Add(new Step(StepType.Move) { Path = planner.GeneratePath(t2v,k2v, Steps).ToArray() });

        steps.Add(new Step(StepType.Move) { Path = planner.GeneratePath(k2v, homeAtTarget4.FrameAtAngles, Steps).ToArray() });
        steps.Add(new Step(StepType.Move) { Path = planner.GeneratePath(homeAtTarget4.FrameAtAngles, t3, Steps).ToArray() });
        steps.Add(new Step(StepType.Move) { Path = planner.GeneratePath(t3,k4, Steps).ToArray() });
        steps.Add(new Step(StepType.Move) { Path = planner.GeneratePath(k4,t4, Steps).ToArray() });
        steps.Add(new Step(StepType.Move) { Path = planner.RepeatFrame(t4, HoldFrames1) });
        steps.Add(new Step(StepType.DetachB));
        steps.Add(new Step(StepType.Move) { Path = planner.GeneratePath(t4,k4, Steps).ToArray() });
        steps.Add(new Step(StepType.Move) { Path = planner.GeneratePath(k4,t3, Steps).ToArray() });
        steps.Add(new Step(StepType.Move) { Path = planner.GeneratePath(t3,homeAtTarget4.FrameAtAngles, Steps).ToArray() });
        steps.Add(new Step(StepType.Detach));

        // Snap lại home
        steps.Add(new Step(StepType.MoveJoints){ JointAngles = CloneAngles(homeAtLinearHome.Angles), MoveDuration = 1f});

        // Parallel: Linear về + Part1 + Part2
        steps.Add(new Step(StepType.Parallel)
        {
            ParallelCoroutines = new List<IEnumerator> {
                DoLinearMove(new Frame(LinearHome.position, LinearHome.rotation), 3f),
                
                Part2Routine()
            }
        });

        steps.Add(new Step(StepType.Delay) { DelayTime = 1f });
    }

    IEnumerator Part1Routine()
    {
        if (currentPhoiFromAbb == null) yield break;

        yield return MoveObject(currentPhoiFromAbb, target4, target4, 0f, true);
        currentPhoiFromAbb = null;
    }

    IEnumerator Part2Routine()
    {
        StartCoroutine(Part1Routine());
        Task.Delay(1000);
        IsPart2Done = false;
        if (phoiQueue.Count == 0) yield break;

        Transform currentPhoi = phoiQueue.Dequeue();
        currentPhoi.gameObject.SetActive(true);

        yield return MoveObject(currentPhoi, target51, target52, 3f, true);
        IsPart2Done = true;
    }

    // ===== Motion Helpers =====
    public IEnumerator MoveObject(Transform obj, Transform from, Transform to, float duration, bool destroyOnFinish = false)
    {
        if (obj == 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);
            obj.SetPositionAndRotation(
                Vector3.Lerp(startPos, to.position, t),
                Quaternion.Slerp(startRot, to.rotation, t)
            );
            yield return null;
        }

        if (destroyOnFinish)
        {
            Destroy(obj.gameObject);
        }
    }

    IEnumerator DoMove(Frame[] path, EffectorType effector)
    {
        SolverCCD solver = (effector == EffectorType.A) ? solverA : solverB;
        foreach (var frame in path)
        {
            float[] jointValues = solver.SolveIK(frame);
            for (int j = 0; j < Joints.Length && j < jointValues.Length; j++)
            {
                float applied = Mathf.Lerp(lastAppliedAngles[j], jointValues[j], ApplyAlpha);
                lastAppliedAngles[j] = applied;
                Joints[j].SetValue(applied);
            }
            yield return new WaitForSeconds(StepDelay);
        }
    }

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

        int frames = Steps;
        if (duration > 0f) StepDelay = duration / frames;

        float[] startAngles = ReadCurrentAngles();

        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);
                lastAppliedAngles[j] = angle;
            }
            yield return new WaitForSeconds(StepDelay);
        }
    }

    IEnumerator DoLinearMove(Frame target, float duration)
    {
        if (LinearBase == null) yield break;

        LinearBase.GetPositionAndRotation(out Vector3 startPos, out Quaternion startRot);
        float elapsed = 0f;
        while (elapsed < duration)
        {
            elapsed += Time.deltaTime;
            float t = Mathf.Clamp01(elapsed / duration);
            LinearBase.SetPositionAndRotation(
                Vector3.Lerp(startPos, target.Position, t),
                Quaternion.Slerp(startRot, target.Rotation, t)
            );
            yield return null;
        }
    }

    // ===== Utility Helpers =====
    HomeData GetHomeAtLinear(Transform linearPose, float[] referenceAngles = null)
    {
        if (linearPose == null) return null;

        LinearBase.GetPositionAndRotation(out Vector3 oldPos, out Quaternion oldRot);
        float[] backupAngles = ReadCurrentAngles();

        LinearBase.SetPositionAndRotation(linearPose.position, linearPose.rotation);

        float[] refAngles = referenceAngles != null ? CloneAngles(referenceAngles) : ReadCurrentAngles();

        for (int i = 0; i < Joints.Length; i++) Joints[i].SetValue(0f);
        Frame frameAtZero = new(EndEffectorA.position, EndEffectorA.rotation);

        for (int i = 0; i < Joints.Length; i++) Joints[i].SetValue(refAngles[i]);
        Frame frameAtAngles = new(EndEffectorA.position, EndEffectorA.rotation);

        for (int i = 0; i < Joints.Length; i++) Joints[i].SetValue(backupAngles[i]);
        LinearBase.SetPositionAndRotation(oldPos, oldRot);

        return new HomeData(refAngles, frameAtZero, frameAtAngles);
    }

    float[] ReadCurrentAngles()
    {
        float[] a = new float[Joints.Length];
        for (int i = 0; i < Joints.Length; i++) a[i] = Joints[i].GetValue();
        return a;
    }

    float[] CloneAngles(float[] src)
    {
        if (src == null) return null;
        var dst = new float[src.Length];
        System.Array.Copy(src, dst, src.Length);
        return dst;
    }

    private Frame CreateLiftFrame(Frame target, float lift, GraspMode _) =>
        new(target.Position + Vector3.up * lift, target.Rotation);

    Frame CreateFrameFromTransform(Transform t, GraspMode mode, Quaternion extraRotation) =>
        new(t.position, GetGraspRotation(mode) * extraRotation);

    Frame CreateFrameFromTransform(Transform t, GraspMode mode) =>
        new(t.position, GetGraspRotation(mode));


    Frame CreateFrameFromTransform(Transform t) =>
        new(t.position, t.rotation);


    Quaternion GetGraspRotation(GraspMode mode)
    {
        return mode switch
        {
            GraspMode.Right => Quaternion.LookRotation(Vector3.up, Vector3.left),
            GraspMode.Top => Quaternion.LookRotation(Vector3.right, Vector3.up),
            GraspMode.Left => Quaternion.LookRotation(Vector3.left),
            _ => Quaternion.identity,
        };
    }
}