Backyard-CTF/Backyard CTF/Assets/Scripts/Flag.cs

using System.Collections;
using UnityEngine;

public class Flag : MonoBehaviour
{
    public Unit.Team team;
    public Transform homePosition;
    public Unit carriedBy;

    float dropTimer;
    bool isDropped;
    Coroutine returnCoroutine;

    void Update()
    {
        if (carriedBy != null)
        {
            // Follow carrier
            transform.position = carriedBy.transform.position + new Vector3(0.3f, 0.3f, 0);
        }
    }

    public void Pickup(Unit unit)
    {
        if (carriedBy != null) return;

        // Cancel return timer if picking up dropped flag
        if (returnCoroutine != null)
        {
            StopCoroutine(returnCoroutine);
            returnCoroutine = null;
        }

        carriedBy = unit;
        unit.hasFlag = true;
        isDropped = false;

        Debug.Log($"{unit.team} picked up {team} flag!");
    }

    public void Drop()
    {
        if (carriedBy == null) return;

        carriedBy.hasFlag = false;
        carriedBy = null;
        isDropped = true;

        Debug.Log($"{team} flag dropped!");

        // Start return timer
        returnCoroutine = StartCoroutine(ReturnAfterDelay());
    }

    IEnumerator ReturnAfterDelay()
    {
        yield return new WaitForSeconds(Game.FlagReturnDelay);

        if (isDropped && carriedBy == null)
        {
            ReturnHome();
        }
    }

    public void ReturnHome()
    {
        if (returnCoroutine != null)
        {
            StopCoroutine(returnCoroutine);
            returnCoroutine = null;
        }

        if (carriedBy != null)
        {
            carriedBy.hasFlag = false;
            carriedBy = null;
        }

        transform.position = homePosition.position;
        isDropped = false;

        Debug.Log($"{team} flag returned home!");
    }

    void OnTriggerEnter2D(Collider2D other)
    {
        // If flag is dropped and a friendly unit touches it, return it home immediately
        if (isDropped && carriedBy == null)
        {
            var unit = other.GetComponent<Unit>();
            if (unit != null && unit.team == team && !unit.isTaggedOut)
            {
                ReturnHome();
            }
        }
    }
}