utad-ui/Source/UTAD_UI/SkillTree/SkillTreeComponent.cpp

#include "SkillTreeComponent.h"
#include "Engine/DataTable.h"

USkillTreeComponent::USkillTreeComponent()
{
	PrimaryComponentTick.bCanEverTick = false;

	AvailableSkillPoints   = 5;
	TotalSkillPointsEarned = 5;
	CurrentSelectionCost   = 0;
}

void USkillTreeComponent::BeginPlay()
{
	Super::BeginPlay();
	InitializeSkillTree();
}

void USkillTreeComponent::InitializeSkillTree()
{
	AllSkills.Empty();
	PurchasedSkills.Empty();

	if (!SkillDataTable)
	{
		UE_LOG(LogTemp, Warning,
			   TEXT("SkillTreeComponent: No skill data table assigned!"));
		return;
	}

	// Load all skills from the data table using RowNames as IDs
	TArray<FName> RowNames = SkillDataTable->GetRowNames();
	for (const FName& RowName : RowNames)
	{
		FSkillNodeData* Row = SkillDataTable->FindRow<FSkillNodeData>(
			RowName, TEXT("SkillTreeComponent"));
		if (Row)
		{
			FSkillNodeRuntime RuntimeNode;
			RuntimeNode.NodeData	 = *Row;
			RuntimeNode.CurrentState = ESkillNodeState::Locked;

			// Use RowName as the skill ID
			AllSkills.Add(RowName, RuntimeNode);
		}
	}

	// Auto-purchase starting skills (these don't cost skill points)
	for (auto& Pair : AllSkills)
	{
		const FName& SkillID = Pair.Key;
		FSkillNodeRuntime& SkillNode = Pair.Value;

		if (SkillNode.NodeData.bStartingSkill)
		{
			// Mark as purchased
			SkillNode.CurrentState = ESkillNodeState::Purchased;
			PurchasedSkills.Add(SkillID);

			// Apply skill effects
			ApplySkillEffects(SkillNode.NodeData);

			// Broadcast purchase event
			OnSkillPurchased.Broadcast(SkillID);

			UE_LOG(LogTemp, Log, TEXT("SkillTreeComponent: Auto-purchased starting skill: %s"),
				   *SkillNode.NodeData.DisplayName.ToString());
		}
	}

	UpdateSkillStates();

	UE_LOG(LogTemp, Log, TEXT("SkillTreeComponent: Initialized with %d skills (%d starting)"),
		   AllSkills.Num(), PurchasedSkills.Num());
}

void USkillTreeComponent::AddSkillPoints(int32 Amount)
{
	if (Amount > 0)
	{
		AvailableSkillPoints += Amount;
		TotalSkillPointsEarned += Amount;
		OnSkillPointsChanged.Broadcast(AvailableSkillPoints);
		UpdateSkillStates();
	}
}

void USkillTreeComponent::RemoveSkillPoints(int32 Amount)
{
	if (Amount > 0)
	{
		AvailableSkillPoints = FMath::Max(0, AvailableSkillPoints - Amount);
		OnSkillPointsChanged.Broadcast(AvailableSkillPoints);
		UpdateSkillStates();
	}
}

bool USkillTreeComponent::SelectSkill(FName SkillID)
{
	FText ErrorMessage;

	if (!AllSkills.Contains(SkillID))
	{
		ErrorMessage = FText::FromString(TEXT("Skill does not exist"));
		OnSelectionError.Broadcast(SkillID, ErrorMessage);
		return false;
	}

	FSkillNodeRuntime* SkillNode = AllSkills.Find(SkillID);
	switch (SkillNode->CurrentState)
	{
	case ESkillNodeState::Purchased:
		ErrorMessage = FText::FromString(TEXT("Skill already purchased"));
		OnSelectionError.Broadcast(SkillID, ErrorMessage);
		return false;

	case ESkillNodeState::Selected:
		ErrorMessage = FText::FromString(TEXT("Skill already selected"));
		OnSelectionError.Broadcast(SkillID, ErrorMessage);
		return false;

	case ESkillNodeState::Locked:
		ErrorMessage = FText::FromString(TEXT("Prerequisites not met or not enough skill points"));
		OnSelectionError.Broadcast(SkillID, ErrorMessage);
		return false;

	case ESkillNodeState::Available:
		// Double-check affordability (should already be Available only if affordable)
		if ((AvailableSkillPoints - CurrentSelectionCost) < SkillNode->NodeData.Cost)
		{
			ErrorMessage = FText::FromString(TEXT("Not enough skill points"));
			OnSelectionError.Broadcast(SkillID, ErrorMessage);
			return false;
		}
		break;
	}

	SkillNode->CurrentState = ESkillNodeState::Selected;
	SelectedSkills.Add(SkillID);
	UpdateSelectionCost();

	UpdateSkillStates();

	OnSkillSelectionChanged.Broadcast(CurrentSelectionCost);
	OnSkillStateChanged.Broadcast(SkillID);

	return true;
}

bool USkillTreeComponent::DeselectSkill(FName SkillID)
{
	if (!AllSkills.Contains(SkillID))
		return false;

	FSkillNodeRuntime* SkillNode = AllSkills.Find(SkillID);

	if (SkillNode->CurrentState != ESkillNodeState::Selected)
		return false;

	CascadeDeselectDependents(SkillID);

	SelectedSkills.Remove(SkillID);
	UpdateSelectionCost();

	UpdateSkillStates();
	OnSkillSelectionChanged.Broadcast(CurrentSelectionCost);

	return true;
}

void USkillTreeComponent::ClearSelection()
{
	SelectedSkills.Empty();
	UpdateSelectionCost();

	UpdateSkillStates();

	OnSkillSelectionChanged.Broadcast(CurrentSelectionCost);
}

bool USkillTreeComponent::ConfirmPurchase()
{
	// Validate we can afford it
	if (!CanAffordSelection())
	{
		FText ErrorMessage = FText::FromString(TEXT("Not enough skill points"));
		OnSelectionError.Broadcast(NAME_None, ErrorMessage);
		return false;
	}

	// Sort selected skills by purchase order to ensure prerequisites are
	// purchased first
	TArray<FName> SortedSelection = SelectedSkills;
	SortedSelection.Sort(
		[this](const FName& A, const FName& B)
		{
			FSkillNodeRuntime* NodeA = AllSkills.Find(A);
			FSkillNodeRuntime* NodeB = AllSkills.Find(B);
			if (NodeA && NodeB)
			{
				return NodeA->NodeData.PurchaseOrder
					< NodeB->NodeData.PurchaseOrder;
			}
			return false;
		});

	// Validate all skills can be purchased in order
	for (const FName& SkillID : SortedSelection)
	{
		FSkillNodeRuntime* SkillNode = AllSkills.Find(SkillID);
		if (!SkillNode)
			continue;

		// Check prerequisites (considering already processed skills in this
		// batch)
		bool bPrereqsMet = true;
		for (const FName& PrereqID : SkillNode->NodeData.Prerequisites)
		{
			if (!IsSkillPurchased(PrereqID))
			{
				// Check if it's in our current batch and would be purchased
				// before this
				int32 PrereqIndex  = SortedSelection.IndexOfByKey(PrereqID);
				int32 CurrentIndex = SortedSelection.IndexOfByKey(SkillID);

				if (PrereqIndex == INDEX_NONE || PrereqIndex >= CurrentIndex)
				{
					bPrereqsMet = false;
					break;
				}
			}
		}

		if (!bPrereqsMet)
		{
			FText ErrorMessage = FText::Format(
				FText::FromString(
					TEXT("Cannot purchase {0}: prerequisites not met")),
				SkillNode->NodeData.DisplayName);
			OnSelectionError.Broadcast(SkillID, ErrorMessage);
			return false;
		}
	}

	// Purchase all selected skills
	for (const FName& SkillID : SortedSelection)
	{
		FSkillNodeRuntime* SkillNode = AllSkills.Find(SkillID);
		if (SkillNode)
		{
			// Mark as purchased
			SkillNode->CurrentState = ESkillNodeState::Purchased;
			PurchasedSkills.Add(SkillID);

			// Apply effects
			ApplySkillEffects(SkillNode->NodeData);

			// Broadcast purchase event
			OnSkillPurchased.Broadcast(SkillID);
			OnSkillStateChanged.Broadcast(SkillID);
		}
	}

	// Deduct skill points
	RemoveSkillPoints(CurrentSelectionCost);

	// Clear selection
	SelectedSkills.Empty();
	UpdateSelectionCost();
	UpdateSkillStates();

	OnSkillSelectionChanged.Broadcast(CurrentSelectionCost);

	return true;
}

bool USkillTreeComponent::IsSkillPurchased(FName SkillID) const
{
	return PurchasedSkills.Contains(SkillID);
}

bool USkillTreeComponent::IsSkillSelected(FName SkillID) const
{
	return SelectedSkills.Contains(SkillID);
}

bool USkillTreeComponent::IsSkillAvailable(FName SkillID) const
{
	const FSkillNodeRuntime* SkillNode = AllSkills.Find(SkillID);
	if (!SkillNode)
		return false;

	return SkillNode->CurrentState == ESkillNodeState::Available;
}

bool USkillTreeComponent::CanSelectSkill(FName	SkillID,
										 FText& OutErrorMessage) const
{
	const FSkillNodeRuntime* SkillNode = AllSkills.Find(SkillID);
	if (!SkillNode)
	{
		OutErrorMessage = FText::FromString(TEXT("Skill does not exist"));
		return false;
	}

	switch (SkillNode->CurrentState)
	{
	case ESkillNodeState::Purchased:
		OutErrorMessage = FText::FromString(TEXT("Skill already purchased"));
		return false;

	case ESkillNodeState::Selected:
		OutErrorMessage = FText::FromString(TEXT("Skill already selected"));
		return false;

	case ESkillNodeState::Locked:
		OutErrorMessage = FText::FromString(TEXT("Prerequisites not met or not enough skill points"));
		return false;

	case ESkillNodeState::Available:
		// Double-check affordability
		if ((AvailableSkillPoints - CurrentSelectionCost) < SkillNode->NodeData.Cost)
		{
			OutErrorMessage = FText::FromString(TEXT("Not enough skill points"));
			return false;
		}
		OutErrorMessage = FText::GetEmpty();
		return true;

	default:
		OutErrorMessage = FText::FromString(TEXT("Unknown skill state"));
		return false;
	}
}

ESkillNodeState USkillTreeComponent::GetSkillState(FName SkillID) const
{
	const FSkillNodeRuntime* SkillNode = AllSkills.Find(SkillID);
	if (SkillNode)
		return SkillNode->CurrentState;

	return ESkillNodeState::Locked;
}

bool USkillTreeComponent::GetSkillNodeData(FName			  SkillID,
										   FSkillNodeRuntime& OutNodeData) const
{
	const FSkillNodeRuntime* SkillNode = AllSkills.Find(SkillID);
	if (SkillNode)
	{
		OutNodeData = *SkillNode;
		return true;
	}

	return false;
}

TArray<FSkillNodeRuntime> USkillTreeComponent::GetAllSkillNodes() const
{
	TArray<FSkillNodeRuntime> Result;
	for (const auto& Pair : AllSkills)
		Result.Add(Pair.Value);

	return Result;
}

float USkillTreeComponent::GetTotalHealthBonus() const
{
	float Total = 0.0f;
	for (const FName& SkillID : PurchasedSkills)
	{
		if (const FSkillNodeRuntime* SkillNode = AllSkills.Find(SkillID))
			Total += SkillNode->NodeData.HealthBonus;
	}

	return Total;
}

float USkillTreeComponent::GetTotalDamageBonus() const
{
	float Total = 0.0f;
	for (const FName& SkillID : PurchasedSkills)
	{
		if (const FSkillNodeRuntime* SkillNode = AllSkills.Find(SkillID))
			Total += SkillNode->NodeData.DamageBonus;
	}

	return Total;
}

float USkillTreeComponent::GetTotalSpeedBonus() const
{
	float Total = 0.0f;
	for (const FName& SkillID : PurchasedSkills)
	{
		if (const FSkillNodeRuntime* SkillNode = AllSkills.Find(SkillID))
			Total += SkillNode->NodeData.SpeedBonus;
	}

	return Total;
}

float USkillTreeComponent::GetTotalHealthBonusPercent() const
{
	float Total = 0.0f;
	for (const FName& SkillID : PurchasedSkills)
	{
		if (const FSkillNodeRuntime* SkillNode = AllSkills.Find(SkillID))
			Total += SkillNode->NodeData.HealthBonusPercent;
	}

	return Total;
}

float USkillTreeComponent::GetTotalDamageBonusPercent() const
{
	float Total = 0.0f;
	for (const FName& SkillID : PurchasedSkills)
	{
		if (const FSkillNodeRuntime* SkillNode = AllSkills.Find(SkillID))
			Total += SkillNode->NodeData.DamageBonusPercent;
	}

	return Total;
}

float USkillTreeComponent::GetTotalSpeedBonusPercent() const
{
	float Total = 0.0f;
	for (const FName& SkillID : PurchasedSkills)
	{
		if (const FSkillNodeRuntime* SkillNode = AllSkills.Find(SkillID))
			Total += SkillNode->NodeData.SpeedBonusPercent;
	}

	return Total;
}

void USkillTreeComponent::ResetAllSkills()
{
	// Refund skill points (but not for starting skills)
	int32 RefundAmount = 0;
	for (const FName& SkillID : PurchasedSkills)
	{
		if (const FSkillNodeRuntime* SkillNode = AllSkills.Find(SkillID))
		{
			// Remove skill effects for all purchased skills
			RemoveSkillEffects(SkillNode->NodeData);

			// Only refund cost for non-starting skills
			if (!SkillNode->NodeData.bStartingSkill)
			{
				RefundAmount += SkillNode->NodeData.Cost;
			}
		}
	}

	// Clear purchases
	PurchasedSkills.Empty();

	// Reset all skill states
	for (auto& Pair : AllSkills)
		Pair.Value.CurrentState = ESkillNodeState::Locked;

	// Clear selection
	SelectedSkills.Empty();
	CurrentSelectionCost = 0;

	// Add refunded points
	if (RefundAmount > 0)
	{
		AddSkillPoints(RefundAmount);
	}

	// Re-purchase starting skills
	for (auto& Pair : AllSkills)
	{
		const FName& SkillID = Pair.Key;
		FSkillNodeRuntime& SkillNode = Pair.Value;

		if (SkillNode.NodeData.bStartingSkill)
		{
			// Mark as purchased
			SkillNode.CurrentState = ESkillNodeState::Purchased;
			PurchasedSkills.Add(SkillID);

			// Apply skill effects
			ApplySkillEffects(SkillNode.NodeData);

			// Broadcast purchase event
			OnSkillPurchased.Broadcast(SkillID);
		}
	}

	// Update states
	UpdateSkillStates();

	// Broadcast events
	OnSkillSelectionChanged.Broadcast(CurrentSelectionCost);

	UE_LOG(LogTemp, Log, TEXT("SkillTreeComponent: Reset all skills. Refunded %d points. Re-purchased %d starting skills."),
		   RefundAmount, PurchasedSkills.Num());
}

void USkillTreeComponent::UpdateSkillStates()
{
	for (auto& Pair : AllSkills)
	{
		const FName& SkillID = Pair.Key;
		FSkillNodeRuntime& SkillNode = Pair.Value;

		switch (SkillNode.CurrentState)
		{
		case ESkillNodeState::Purchased:
			// Never change purchased state
			continue;

		case ESkillNodeState::Selected:
			// If marked as Selected but NOT in SelectedSkills array, it was deselected
			if (!SelectedSkills.Contains(SkillID))
			{
				// Check if it can be Available (prerequisites + affordability)
				bool bPrereqsMet = ArePrerequisitesMet(SkillNode.NodeData, true);
				bool bCanAfford = (AvailableSkillPoints - CurrentSelectionCost) >= SkillNode.NodeData.Cost;

				ESkillNodeState NewState = (bPrereqsMet && bCanAfford)
					? ESkillNodeState::Available
					: ESkillNodeState::Locked;

				SkillNode.CurrentState = NewState;
				OnSkillStateChanged.Broadcast(SkillID);
			}
			// Keep it Selected
			break;

		case ESkillNodeState::Locked:
		case ESkillNodeState::Available:
			{
				// Skill is Available only if:
				// 1. Prerequisites are met (purchased or selected)
				// 2. Player has enough skill points to afford it
				bool bPrereqsMet = ArePrerequisitesMet(SkillNode.NodeData, true);
				bool bCanAfford = (AvailableSkillPoints - CurrentSelectionCost) >= SkillNode.NodeData.Cost;

				ESkillNodeState NewState = (bPrereqsMet && bCanAfford)
					? ESkillNodeState::Available
					: ESkillNodeState::Locked;

				// Only broadcast if state actually changed
				if (SkillNode.CurrentState != NewState)
				{
					SkillNode.CurrentState = NewState;
					OnSkillStateChanged.Broadcast(SkillID);
				}
			}
			break;
		}
	}
}

void USkillTreeComponent::UpdateSelectionCost()
{
	CurrentSelectionCost = 0;
	for (const FName& SkillID : SelectedSkills)
	{
		if (const FSkillNodeRuntime* SkillNode = AllSkills.Find(SkillID))
		{
			CurrentSelectionCost += SkillNode->NodeData.Cost;
		}
	}
}

bool USkillTreeComponent::ArePrerequisitesMet(
	const FSkillNodeData& SkillData, bool bIncludeSelected) const
{
	for (const FName& PrereqID : SkillData.Prerequisites)
	{
		// Check if prerequisite is purchased
		if (IsSkillPurchased(PrereqID))
			continue;

		// If including selected, also check if prerequisite is selected
		if (bIncludeSelected && IsSkillSelected(PrereqID))
			continue;

		// Prerequisite not met
		return false;
	}
	return true;
}

bool USkillTreeComponent::HasChildrenPurchased(FName SkillID) const
{
	for (const auto& Pair : AllSkills)
	{
		// Pair.Key is the RowName (SkillID)
		if (Pair.Value.NodeData.Prerequisites.Contains(SkillID)
			&& IsSkillPurchased(Pair.Key))
		{
			return true;
		}
	}
	return false;
}

void USkillTreeComponent::CascadeDeselectDependents(FName SkillID)
{
	// Find all selected skills that depend on this skill
	TArray<FName> ToDeselect;
	for (const FName& SelectedID : SelectedSkills)
	{
		if (SelectedID == SkillID)
			continue;

		FSkillNodeRuntime* SelectedNode = AllSkills.Find(SelectedID);
		if (SelectedNode && SelectedNode->NodeData.Prerequisites.Contains(SkillID))
		{
			ToDeselect.Add(SelectedID);
		}
	}

	// Recursively deselect dependent skills
	for (const FName& DependentID : ToDeselect)
	{
		CascadeDeselectDependents(DependentID);
		SelectedSkills.Remove(DependentID);
	}
}

void USkillTreeComponent::ApplySkillEffects(const FSkillNodeData& SkillData)
{
	// This is where you would apply the actual skill effects to the character
	// For now, we just log it
	UE_LOG(LogTemp, Log, TEXT("Applied skill effects for: %s"),
		   *SkillData.DisplayName.ToString());
}

void USkillTreeComponent::RemoveSkillEffects(const FSkillNodeData& SkillData)
{
	// This is where you would remove skill effects from the character
	// For now, we just log it
	UE_LOG(LogTemp, Log, TEXT("Removed skill effects for: %s"),
		   *SkillData.DisplayName.ToString());
}