RPG-RougeLiteBattler/Assets/Scripts/BattleSetup/BattleSetupMenu.cs

using UnityEngine;
using UnityEngine.UI;
using System.Collections.Generic;
using System.Threading.Tasks;
using UnityEngine.SceneManagement;

public class BattleSetupMenu : MonoBehaviour
{
    public Transform playerListParent;
    public Transform enemyListParent;
    public GameObject characterEntryPrefab; // Prefab dropdown for Weapon, remove button etc.
    public Button addPlayerButton;
    public Button addEnemyButton;
    public Button startBattleButton;
    public Button backToTeamSelectButton; // New back button

    private List<GameObject> playerEntries = new List<GameObject>();
    private List<GameObject> enemyEntries = new List<GameObject>();

    void Start()
    {
        addPlayerButton.onClick.AddListener(() => AddCharacterEntry(true));
        addEnemyButton.onClick.AddListener(() => AddCharacterEntry(false));
        startBattleButton.onClick.AddListener(StartBattle);

        // Add back button listener
        if (backToTeamSelectButton != null)
        {
            backToTeamSelectButton.onClick.AddListener(BackToTeamSelect);
        }

        UpdateStartButton();
    }

    void AddCharacterEntry(bool isPlayer)
    {
        var entry = Instantiate(characterEntryPrefab, isPlayer ? playerListParent : enemyListParent);
        var entryScript = entry.GetComponent<CharacterEntryUI>();
        entryScript.Setup(isPlayer, RemoveEntry, UpdateStartButton);
        if (isPlayer)
        {
            playerEntries.Add(entry);
        }
        else
        {
            enemyEntries.Add(entry);
        }
        UpdateStartButton();
    }

    void RemoveEntry(GameObject entry, bool isPlayer)
    {
        if (isPlayer)
        {
            playerEntries.Remove(entry);
        }
        else
        {
            enemyEntries.Remove(entry);
        }
        Destroy(entry);
        UpdateStartButton();
    }

    void UpdateStartButton()
    {
        startBattleButton.interactable = playerEntries.Count > 0 && enemyEntries.Count > 0;
    }

    void StartBattle()
    {
        BattleSetupData.playerSelections.Clear();
        BattleSetupData.enemySelections.Clear();
        foreach (var entry in playerEntries)
        {
            BattleSetupData.playerSelections.Add(entry.GetComponent<CharacterEntryUI>().GetSelection());
        }
        foreach (var entry in enemyEntries)
        {
            BattleSetupData.enemySelections.Add(entry.GetComponent<CharacterEntryUI>().GetSelection());
        }

        SceneManager.LoadScene("BattleScene");
    }

    void BackToTeamSelect()
    {
        SceneManager.LoadScene("MainTeamSelectScene");
    }

}