UpgradeManager

extends Node

@export var experience_manager: Node
@export var upgrade_screen_scene: PackedScene

var current_upgrades = {}
var upgrade_pool: WeightedTable = WeightedTable.new()

# Preload upgrades
var upgrade_axe = preload("res://resources/upgrades/axe.tres")
var upgrade_axe_damage = preload("res://resources/upgrades/axe_damage.tres")
var upgrade_sword_rate = preload("res://resources/upgrades/sword_rate.tres")
var upgrade_sword_damage = preload("res://resources/upgrades/sword_damage.tres")
var upgrade_sword_life_steal = preload("res://resources/upgrades/sword_life_steal.tres")
var upgrade_player_speed = preload("res://resources/upgrades/player_speed.tres")
var upgrade_anvil = preload("res://resources/upgrades/anvil.tres")
var upgrade_anvil_count = preload("res://resources/upgrades/anvil_count.tres")
var upgrade_scythe = preload("res://resources/upgrades/scythe.tres")
var upgrade_damage_aura = preload("res://resources/upgrades/damage_aura.tres")

# Mastery requirements for upgrades
const mastery_requirements = {
	"sword_life_steal": {
		"weapon_id": "sword",
		"level": WeaponMasteryManager.MasteryLevel.TEST
	},
	"anvil_count": {
		"weapon_id": "anvil",
		"level": WeaponMasteryManager.MasteryLevel.TEST
	}
}

# Base upgrade weights
const base_weights = {
	"axe": 8,
	"anvil": 800,
	"sword_rate": 10,
	"sword_damage": 10,
	"player_speed": 6,
	"scythe": 4,
	"damage_aura": 6,
	"sword_life_steal": 5
}

func _ready():
	initialize_upgrade_pool()
	if experience_manager:
		experience_manager.level_up.connect(on_level_up)
	else:
		push_error("experience_manager not assigned in upgrade_manager")

func initialize_upgrade_pool():
	for upgrade_id in base_weights:
		var upgrade = get("upgrade_" + upgrade_id)
		if upgrade:
			upgrade_pool.add_item(upgrade, base_weights[upgrade_id])

func meets_mastery_requirement(upgrade_id: String) -> bool:
	if !mastery_requirements.has(upgrade_id):
		return true
		
	var req = mastery_requirements[upgrade_id]
	var controller = get_tree().get_first_node_in_group(req.weapon_id + "_controller")
	if !controller or !controller.has_node("WeaponMasteryComponent"):
		return false
		
	var mastery = controller.get_node("WeaponMasteryComponent")
	return mastery.has_mastery_requirement(req.level)

func pick_upgrades() -> Array[AbilityUpgrade]:
	var available_upgrades: Array[Dictionary] = []
	
	for item in upgrade_pool.items:
		var upgrade = item["item"] as AbilityUpgrade
		if meets_mastery_requirement(upgrade.id):
			available_upgrades.append(item)
	
	var filtered_pool = WeightedTable.new()
	for item in available_upgrades:
		filtered_pool.add_item(item["item"], item["weight"])
	
	var chosen_upgrades: Array[AbilityUpgrade] = []
	for i in 2:
		if filtered_pool.items.size() == chosen_upgrades.size():
			break
		var chosen = filtered_pool.pick_item(chosen_upgrades)
		if chosen:
			chosen_upgrades.append(chosen)
	
	return chosen_upgrades

func update_upgrade_pool(chosen_upgrade: AbilityUpgrade):
	match chosen_upgrade.id:
		"axe":
			upgrade_pool.add_item(upgrade_axe_damage, 8)
		"anvil":
			upgrade_pool.add_item(upgrade_anvil_count, 7000)

func apply_upgrade(upgrade: AbilityUpgrade):
	if !current_upgrades.has(upgrade.id):
		current_upgrades[upgrade.id] = {
			"resource": upgrade,
			"quantity": 1
		}
	else:
		current_upgrades[upgrade.id]["quantity"] += 1
	
	if upgrade.max_quantity > 0:
		var current_quantity = current_upgrades[upgrade.id]["quantity"]
		if current_quantity == upgrade.max_quantity:
			upgrade_pool.remove_item(upgrade)
			
	update_upgrade_pool(upgrade)
	GameEvents.emit_ability_upgrade_added(upgrade, current_upgrades)

func on_upgrade_selected(upgrade: AbilityUpgrade):
	apply_upgrade(upgrade)

func on_level_up(current_level: int):
	var upgrade_screen = upgrade_screen_scene.instantiate()
	add_child(upgrade_screen)
	var upgrades = pick_upgrades()
	upgrade_screen.set_ability_upgrades(upgrades)
	upgrade_screen.upgrade_selected.connect(on_upgrade_selected)