import { world, system, EquipmentSlot } from "@minecraft/server";
import { ActionFormData } from "@minecraft/server-ui";
import { getAbilityInfo, getAbilityCooldown } from "./companion_abilities.js";

const weaponPrefix = "aft_soulhunters:";
const companionPrefix = "aft_soulhunters:";
const tameEvent = "aft_soulhunters:tame";
const despawnEvent = "aft_soulhunters:despawn";
const activateEvent = "aft_soulhunters:activate";
const removeActivateEvent = "aft_soulhunters:remove_activate";
const noSoulsEvent = "aft_soulhunters:no_souls";
const dontAttackEvent = "aft_soulhunters:dont_attack";
const BEHAVIOR_PROP = "soulhunters_behavior";
const COMPANION_ID_PROP = "soulhunters_companion_id";
const COMPANION_TYPE_PROP = "soulhunters_companion_type";
const COMPANION_WEAPON_SLOT_PROP = "soulhunters_companion_weapon_slot";
const ACTIVE_PROP = "active";
const SOUL_ITEM_ID = "aft_soulhunters:soul";

const MAX_SOULS = 10;

const SCORE_TO_EVENT = {
  0: "cero", 1: "one", 2: "two", 3: "three", 4: "four",
  5: "five", 6: "six", 7: "seven", 8: "eight", 9: "nine", 10: "ten"
};

const UI_OPTIONS = [
  { text: "Attack all monsters", event: "aft_soulhunters:attack_all", texture: "textures/aft/soulhunters/behaviour/follow_and_attack" },
  { text: "Attack only what I attack", event: "aft_soulhunters:attack_only_when_attacked", texture: "textures/aft/soulhunters/behaviour/guard" },
  { text: "Use Special Ability", event: "special_ability", texture: "textures/items/book_enchanted" },
  { text: "Roar Attack", event: "scream_ability", texture: "textures/items/fireworks_charge" },
];

function getAnimalFromLore(item) {
  const lore = item?.getLore();
  if (!lore) return null;
  for (const line of lore) {
    const lower = line.toLowerCase();
    if (lower.includes("dragon companion")) return "dragon";
    if (lower.includes("elephant companion")) return "elephant";
    if (lower.includes("fox companion")) return "fox";
    if (lower.includes("leopard companion")) return "leopard";
    if (lower.includes("manta ray companion")) return "manta_ray";
    if (lower.includes("panda companion")) return "panda";
    if (lower.includes("phoenix companion")) return "phoenix";
    if (lower.includes("red panda companion")) return "red_panda";
    if (lower.includes("tiger companion")) return "tiger";
    if (lower.includes("wolf companion")) return "wolf";
  }
  return null;
}

function getWeaponInfo(item) {
  if (!item || !item.typeId.startsWith(weaponPrefix)) return null;

  const companion = getAnimalFromLore(item);
  if (!companion) return null;

  const lore = item.getLore();
  const soulLine = lore.find(line => line.trim().toLowerCase().endsWith("souls"));
  let souls = 0;

  if (soulLine) {
    const match = soulLine.match(/(\d+)\s*\/\s*\d+\s*Souls/i);
    if (match) {
      souls = parseInt(match[1]);
    }
  }

  return { companion, souls };
}

function updateWeaponLore(item, companion, souls) {
  const newLore = [`${capitalize(companion)} Companion`, `${souls}/${MAX_SOULS} Souls`];
  item.setLore(newLore);
}

function capitalize(text) {   
  return text.replace(/_/g, ' ').replace(/\b\w/g, char => char. toUpperCase()); 
}

function findEntityById(id, dimension) {
  for (const entity of dimension.getEntities()) {
    if (entity.id === id) return entity;
  }
  return null;
}

function summonOrUpdateCompanion(player, weaponInfo, item) {
  const currentId = player.getDynamicProperty(COMPANION_ID_PROP);
  const currentType = player.getDynamicProperty(COMPANION_TYPE_PROP);
  const companionTypeKey = weaponInfo.companion; // No more levels

  const lastItemId = player.getDynamicProperty("soulhunters_last_item_id");
  const currentItemId = item?.typeId ?? "";

  // Find the current weapon slot
  const equippable = player.getComponent("minecraft:equippable");
  const mainHandItem = equippable?.getEquipment(EquipmentSlot.Mainhand);
  let currentWeaponSlot = -1;
  
  
  // First check if the weapon is in inventory (more reliable)
  const inventory = player.getComponent("inventory");
  if (inventory && inventory.container) {
    const itemLore = item.getLore();
    for (let i = 0; i < inventory.container.size; i++) {
      const invItem = inventory.container.getItem(i);
      if (invItem && invItem.typeId === item.typeId) {
        const invLore = invItem.getLore();
        if (JSON.stringify(invLore) === JSON.stringify(itemLore)) {
          currentWeaponSlot = i;
          break;
        }
      }
    }
  }
  
  // Only check main hand if not found in inventory
  if (currentWeaponSlot === -1 && mainHandItem && mainHandItem.typeId === item.typeId) {
    const mainHandLore = mainHandItem.getLore();
    const itemLore = item.getLore();
    const loreMatches = JSON.stringify(mainHandLore) === JSON.stringify(itemLore);
    
    
    if (loreMatches) {
      // Weapon is in main hand, slot is -1 (special case)
      currentWeaponSlot = -1;
    }
  }
  
  // If still not found, something went wrong
  if (currentWeaponSlot === -1 && (!mainHandItem || mainHandItem.typeId !== item.typeId || JSON.stringify(mainHandItem.getLore()) !== JSON.stringify(item.getLore()))) {
    currentWeaponSlot = -1; // Default to main hand as fallback
  }

  if (currentType === companionTypeKey && lastItemId === currentItemId) {
    const existing = findEntityById(currentId, player.dimension);
    if (existing?.isValid()) {
      const souls = weaponInfo.souls;
      const scoreName = SCORE_TO_EVENT[souls];
      try {
        existing.runCommand(`event entity @s aft_soulhunters:souls_bar_${scoreName}`);
      } catch (e) {}
      // Update the weapon slot tracking on both player and companion
      player.setDynamicProperty(COMPANION_WEAPON_SLOT_PROP, currentWeaponSlot);
      existing.setDynamicProperty("weapon_slot", currentWeaponSlot);
    }
    return;
  }

  if (currentId) {
    const old = findEntityById(currentId, player.dimension);
    if (old?.isValid()) old.triggerEvent(despawnEvent);
  }

  const entityId = `${companionPrefix}${weaponInfo.companion}`;
  const spawnPos = { x: player.location.x, y: player.location.y + 1, z: player.location.z };

  const companion = player.dimension.spawnEntity(entityId, spawnPos);
  if (!companion) return;

  try {
    const tameable = companion.getComponent("minecraft:tameable");
    if (tameable && !tameable.isTamed) tameable.tame(player);
    companion.triggerEvent(tameEvent);
  } catch (e) {}

  companion.setDynamicProperty(ACTIVE_PROP, false);
  
  // Store the weapon slot on the companion itself
  companion.setDynamicProperty("weapon_slot", currentWeaponSlot);
  
  // Play companion activation sound
  player.playSound("aft_soulhunters_companion_activate");

  const scoreName = SCORE_TO_EVENT[weaponInfo.souls];
  try {
    companion.runCommand(`event entity @s aft_soulhunters:souls_bar_${scoreName}`);
  } catch (e) {}

  player.setDynamicProperty(COMPANION_ID_PROP, companion.id);
  player.setDynamicProperty(COMPANION_TYPE_PROP, companionTypeKey);
  player.setDynamicProperty(COMPANION_WEAPON_SLOT_PROP, currentWeaponSlot);
  player.setDynamicProperty("soulhunters_last_item_id", currentItemId);

}

/**
 * Finds all weapons with companion lore in player's inventory
 * @param {Player} player - The player
 * @returns {Array} Array of weapon objects with slot information
 */
function findAllCompanionWeapons(player) {
  const weapons = [];
  const equippable = player.getComponent("minecraft:equippable");
  const inventory = player.getComponent("inventory");
  
  // Check main hand
  const mainHandItem = equippable?.getEquipment(EquipmentSlot.Mainhand);
  if (mainHandItem) {
    const weaponInfo = getWeaponInfo(mainHandItem);
    if (weaponInfo) {
      weapons.push({
        item: mainHandItem,
        info: weaponInfo,
        slot: -1
      });
    }
  }
  
  // Check inventory slots
  if (inventory && inventory.container) {
    for (let i = 0; i < inventory.container.size; i++) {
      const invItem = inventory.container.getItem(i);
      if (invItem) {
        const weaponInfo = getWeaponInfo(invItem);
        if (weaponInfo) {
          weapons.push({
            item: invItem,
            info: weaponInfo,
            slot: i
          });
        }
      }
    }
  }
  
  return weapons;
}

/**
 * Finds existing companion in the world that matches a weapon
 * @param {Player} player - The player
 * @param {Object} weaponInfo - Weapon information
 * @returns {Entity|null} Existing companion or null
 */
function findExistingCompanion(player, weaponInfo) {
  const expectedTypeId = `${companionPrefix}${weaponInfo.companion}`;
  
  // Search all entities in the dimension for matching companions
  for (const entity of player.dimension.getEntities()) {
    if (entity.typeId === expectedTypeId) {
      // Check if this companion is ownerless or belongs to this player
      try {
        const tameable = entity.getComponent("minecraft:tameable");
        if (tameable && tameable.isTamed) {
          const owner = tameable.tamedToPlayer;
          if (owner && owner.id === player.id) {
            return entity;
          }
        }
      } catch (e) {
        // If tameable component check fails, continue searching
      }
    }
  }
  
  return null;
}

// OLD COMPANION MANAGEMENT SYSTEM - DISABLED (replaced by companion_manager.js)
/*
system.runInterval(() => {
  for (const player of world.getPlayers()) {
    const currentId = player.getDynamicProperty(COMPANION_ID_PROP);
    const currentType = player.getDynamicProperty(COMPANION_TYPE_PROP);
    const currentWeaponSlot = player.getDynamicProperty(COMPANION_WEAPON_SLOT_PROP);
    
    // Find all weapons with companion lore in inventory
    const companionWeapons = findAllCompanionWeapons(player);
    
    // Check if current companion should be deactivated due to weapon change
    if (currentId && currentType !== undefined && currentWeaponSlot !== undefined) {
      const companion = findEntityById(currentId, player.dimension);
      if (companion?.isValid()) {
        const isActive = companion.getDynamicProperty(ACTIVE_PROP);
        
        // Check if the companion's weapon is still in the correct slot and is currently held
        let shouldDeactivate = false;
        
        if (currentWeaponSlot === -1) {
          // Companion weapon should be in main hand
          const equippable = player.getComponent("minecraft:equippable");
          const mainHandItem = equippable?.getEquipment(EquipmentSlot.Mainhand);
          
          if (!mainHandItem) {
            // No item in main hand - deactivate
            shouldDeactivate = true;
          } else {
            // Check if main hand item is the companion weapon
            const weaponInfo = getWeaponInfo(mainHandItem);
            if (!weaponInfo || weaponInfo.companion !== currentType) {
              // Wrong weapon or no companion weapon in main hand - deactivate
              shouldDeactivate = true;
            }
          }
        } else {
          // Companion weapon should be in specific inventory slot
          // For inventory slots, we don't auto-deactivate as the weapon might just be moved around
          // But we should deactivate if the weapon is now in main hand and doesn't match
          const equippable = player.getComponent("minecraft:equippable");
          const mainHandItem = equippable?.getEquipment(EquipmentSlot.Mainhand);
          
          if (mainHandItem) {
            const weaponInfo = getWeaponInfo(mainHandItem);
            if (weaponInfo && weaponInfo.companion !== currentType) {
              // Different companion weapon in main hand - deactivate current companion
              shouldDeactivate = true;
            }
          }
        }
        
        // Auto-deactivate if conditions are met
        if (shouldDeactivate && isActive) {
          companion.setDynamicProperty(ACTIVE_PROP, false);
          companion.triggerEvent(removeActivateEvent);
          player.sendMessage(`§c${capitalize(currentType)} companion auto-deactivated (weapon changed)`);
        }
      }
    }
    
    // If no weapons with companion lore found, clear any existing companion
    if (companionWeapons.length === 0) {
      if (currentId) {
        const companion = findEntityById(currentId, player.dimension);
        if (companion?.isValid()) {
          const isActive = companion.getDynamicProperty(ACTIVE_PROP);
          if (isActive) {
            companion.setDynamicProperty(ACTIVE_PROP, false);
            companion.triggerEvent(removeActivateEvent);
          }
          companion.triggerEvent(despawnEvent);
        }
        // Clear companion tracking
        player.setDynamicProperty(COMPANION_ID_PROP, undefined);
        player.setDynamicProperty(COMPANION_TYPE_PROP, undefined);
        player.setDynamicProperty(COMPANION_WEAPON_SLOT_PROP, undefined);
        player.setDynamicProperty("soulhunters_last_item_id", undefined);
      }
      continue;
    }
    
    // Check if current companion is valid and matches one of the weapons
    let validCompanionFound = false;
    if (currentId && currentType) {
      const companion = findEntityById(currentId, player.dimension);
      if (companion?.isValid()) {
        // Check if this companion matches any of our weapons
        for (const weapon of companionWeapons) {
          const expectedType = weapon.info.companion; // No more levels
          if (currentType === expectedType) {
            validCompanionFound = true;
            // Update companion's weapon slot tracking
            companion.setDynamicProperty("weapon_slot", weapon.slot);
            player.setDynamicProperty(COMPANION_WEAPON_SLOT_PROP, weapon.slot);
            
            // Update soul display
            const souls = weapon.info.souls;
            const scoreName = SCORE_TO_EVENT[souls];
            try {
              companion.runCommand(`event entity @s aft_soulhunters:souls_bar_${scoreName}`);
            } catch (e) {}
            break;
          }
        }
        
        // If current companion doesn't match any weapon, despawn it
        if (!validCompanionFound) {
          const isActive = companion.getDynamicProperty(ACTIVE_PROP);
          if (isActive) {
            companion.setDynamicProperty(ACTIVE_PROP, false);
            companion.triggerEvent(removeActivateEvent);
          }
          companion.triggerEvent(despawnEvent);
          
          // Clear companion tracking
          player.setDynamicProperty(COMPANION_ID_PROP, undefined);
          player.setDynamicProperty(COMPANION_TYPE_PROP, undefined);
          player.setDynamicProperty(COMPANION_WEAPON_SLOT_PROP, undefined);
          player.setDynamicProperty("soulhunters_last_item_id", undefined);
        }
      } else {
        // Companion entity no longer exists, clear tracking
        player.setDynamicProperty(COMPANION_ID_PROP, undefined);
        player.setDynamicProperty(COMPANION_TYPE_PROP, undefined);
        player.setDynamicProperty(COMPANION_WEAPON_SLOT_PROP, undefined);
        player.setDynamicProperty("soulhunters_last_item_id", undefined);
        validCompanionFound = false;
      }
    }
    
    // If no valid companion found, try to find or summon one
    if (!validCompanionFound) {
      // First, try to find an existing companion in the world
      let reconnectedCompanion = null;
      for (const weapon of companionWeapons) {
        const existingCompanion = findExistingCompanion(player, weapon.info);
        if (existingCompanion) {
          reconnectedCompanion = { companion: existingCompanion, weapon: weapon };
          break;
        }
      }
      
      if (reconnectedCompanion) {
        // Reconnect to existing companion
        const { companion, weapon } = reconnectedCompanion;
        const companionTypeKey = weapon.info.companion; // No more levels
        
        // Update soul display
        const souls = weapon.info.souls;
        const scoreName = SCORE_TO_EVENT[souls];
        try {
          companion.runCommand(`event entity @s aft_soulhunters:souls_bar_${scoreName}`);
        } catch (e) {}
        
        // Update tracking
        companion.setDynamicProperty("weapon_slot", weapon.slot);
        player.setDynamicProperty(COMPANION_ID_PROP, companion.id);
        player.setDynamicProperty(COMPANION_TYPE_PROP, companionTypeKey);
        player.setDynamicProperty(COMPANION_WEAPON_SLOT_PROP, weapon.slot);
        player.setDynamicProperty("soulhunters_last_item_id", weapon.item.typeId);
        
        player.sendMessage(`§aReconnected to your ${weapon.info.companion} companion!`);
      } else {
        // Summon new companion for the first weapon found
        const firstWeapon = companionWeapons[0];
        summonOrUpdateCompanion(player, firstWeapon.info, firstWeapon.item);
      }
    }
  }
}, 10);
*/

// OLD CLEANUP SYSTEM - DISABLED (replaced by companion_manager.js)
/*
// Cleanup orphaned companions every 5 seconds - simplified since main interval handles most logic
system.runInterval(() => {
  for (const player of world.getPlayers()) {
    const currentId = player.getDynamicProperty(COMPANION_ID_PROP);
    if (!currentId) continue;
    
    const companion = findEntityById(currentId, player.dimension);
    if (!companion?.isValid()) {
      // Companion no longer exists, clear the properties
      player.setDynamicProperty(COMPANION_ID_PROP, undefined);
      player.setDynamicProperty(COMPANION_TYPE_PROP, undefined);
      player.setDynamicProperty(COMPANION_WEAPON_SLOT_PROP, undefined);
      player.setDynamicProperty("soulhunters_last_item_id", undefined);
    }
  }
}, 100); // Check every 5 seconds (100 ticks)
*/

// OLD COMPANION ACTIVATION SYSTEM - DISABLED (replaced by companion_manager.js)
/*
world.afterEvents.itemStopUse.subscribe(event => {
  const { source: player } = event;
  
  if (!player.isSneaking) return;

  const companionId = player.getDynamicProperty(COMPANION_ID_PROP);
  if (!companionId) return;

  const companion = findEntityById(companionId, player.dimension);
  if (!companion?.isValid()) return;

  // Find the weapon that corresponds to this companion
  const companionType = companion.typeId.replace(companionPrefix, "");
  const companionWeapons = findAllCompanionWeapons(player);
  
  let matchingWeapon = null;
  for (const weapon of companionWeapons) {
    if (weapon.info.companion === companionType) {
      matchingWeapon = weapon;
      break;
    }
  }
  
  if (!matchingWeapon) return;

  player.startItemCooldown("companion", 0);
  const isActive = companion.getDynamicProperty(ACTIVE_PROP) ?? false;

  if (isActive) {
    companion.setDynamicProperty(ACTIVE_PROP, false);
    companion.triggerEvent(removeActivateEvent);
    player.sendMessage(`§c${capitalize(companionType)} companion deactivated`);
  } else {
    const souls = matchingWeapon.info.souls ?? 0;
    if (souls > 0) {
      companion.setDynamicProperty(ACTIVE_PROP, true);
      companion.triggerEvent(activateEvent);
      startSoulDrain(companion);
      const behavior = companion.getDynamicProperty(BEHAVIOR_PROP);
      if (behavior) companion.triggerEvent(behavior);
      player.sendMessage(`§a${capitalize(companionType)} companion activated`);
    } else {
      companion.setDynamicProperty(ACTIVE_PROP, false);
      companion.triggerEvent(noSoulsEvent);
      player.sendMessage(`§c${capitalize(companionType)} companion has no souls to activate`);
    }
  }
});
*/

// DISABLED: Now handled by unified interaction_handler.js
// This prevents conflicts between multiple playerInteractWithEntity handlers
/*
// OLD COMPANION INTERACTION SYSTEM - DISABLED (replaced by interaction_handler.js)
/*
world.afterEvents.playerInteractWithEntity.subscribe(event => {
  const { player, target: companion, itemStack } = event;
  if (!player || !companion) return;

  // Only handle companion entities - exclude altars and other entities
  if (!companion.typeId || !companion.typeId.startsWith(companionPrefix)) {
    return; // Not from our mod, ignore this interaction
  }
  
  // Exclude altar entities and altar companions - they have their own interaction systems
  if (companion.typeId.includes("altar")) {
    return; // This is an altar or altar companion, let other scripts handle it
  }

  if (!itemStack || itemStack.typeId !== SOUL_ITEM_ID) {
    showCompanionControlUI(player, companion);
    return;
  }

  // Get the stored weapon slot from the companion
  const storedWeaponSlot = companion.getDynamicProperty("weapon_slot");
  if (storedWeaponSlot === undefined) {
    return;
  }

  
  let weaponItem = null;
  let weaponInfo = null;
  
  // Get the specific weapon from the stored slot
  if (storedWeaponSlot === -1) {
    // Weapon should be in main hand
    const equippable = player.getComponent("minecraft:equippable");
    const mainHandItem = equippable?.getEquipment(EquipmentSlot.Mainhand);
    if (mainHandItem) {
      const mainHandInfo = getWeaponInfo(mainHandItem);
      if (mainHandInfo) {
        weaponItem = mainHandItem;
        weaponInfo = mainHandInfo;
      }
    }
  } else {
    // Weapon should be in specific inventory slot
    const inventory = player.getComponent("inventory");
    if (inventory && inventory.container) {
      const invItem = inventory.container.getItem(storedWeaponSlot);
      if (invItem) {
        const invInfo = getWeaponInfo(invItem);
        if (invInfo) {
          weaponItem = invItem;
          weaponInfo = invInfo;
        }
      }
    }
  }
  
  // If weapon not found in stored slot, search for it (it might have moved)
  if (!weaponItem) {
    const companionType = companion.typeId.replace(companionPrefix, "");
    
    // First check main hand
    const equippable = player.getComponent("minecraft:equippable");
    const mainHandItem = equippable?.getEquipment(EquipmentSlot.Mainhand);
    
    if (mainHandItem) {
      const mainHandInfo = getWeaponInfo(mainHandItem);
      if (mainHandInfo && mainHandInfo.companion === companionType) {
        weaponItem = mainHandItem;
        weaponInfo = mainHandInfo;
        // Update the stored slot
        companion.setDynamicProperty("weapon_slot", -1);
      }
    }
    
    // If not in main hand, search inventory
    if (!weaponItem) {
      const inventory = player.getComponent("inventory");
      if (inventory && inventory.container) {
        for (let i = 0; i < inventory.container.size; i++) {
          const invItem = inventory.container.getItem(i);
          if (invItem) {
            const invInfo = getWeaponInfo(invItem);
            if (invInfo && invInfo.companion === companionType) {
              weaponItem = invItem;
              weaponInfo = invInfo;
              // Update the stored slot
              companion.setDynamicProperty("weapon_slot", i);
              break;
            }
          }
        }
      }
    }
  }
  
  if (!weaponItem || !weaponInfo) {
    return;
  }
  
  // Verify this is the correct companion type
  const companionType = companion.typeId.replace(companionPrefix, "");
  if (weaponInfo.companion !== companionType) {
    return;
  }
  
  let souls = weaponInfo.souls ?? 0;
  if (souls >= MAX_SOULS) return;

  // Consume the soul item
  const equippable = player.getComponent("minecraft:equippable");
  const mainHandItem = equippable?.getEquipment(EquipmentSlot.Mainhand);
  if (mainHandItem && mainHandItem.typeId === SOUL_ITEM_ID) {
    if (mainHandItem.amount > 1) {
      mainHandItem.amount -= 1;
      equippable.setEquipment(EquipmentSlot.Mainhand, mainHandItem);
    } else {
      equippable.setEquipment(EquipmentSlot.Mainhand, undefined);
    }
  }

  // Update the weapon
  souls++;
  updateWeaponLore(weaponItem, weaponInfo.companion, souls);
  
  // Get the current weapon slot (might have been updated during search)
  const currentWeaponSlot = companion.getDynamicProperty("weapon_slot");
  
  // Set the weapon back to its proper location
  if (currentWeaponSlot === -1) {
    equippable.setEquipment(EquipmentSlot.Mainhand, weaponItem);
  } else {
    const inventory = player.getComponent("inventory");
    if (inventory && inventory.container) {
      inventory.container.setItem(currentWeaponSlot, weaponItem);
    }
  }

  const scoreName = SCORE_TO_EVENT[souls];
  try {
    companion.runCommand(`event entity @s aft_soulhunters:souls_bar_${scoreName}`);
    companion.triggerEvent("aft_soulhunters:eat");
  } catch (e) {}
});
*/

function showCompanionControlUI(player, entity) {
  if (!player || !entity) return;

  const isActive = entity.getDynamicProperty(ACTIVE_PROP);
  if (!isActive) {
    player.sendMessage("Your companion must be activated first.");
    return;
  }

  // Get the weapon slot that this specific companion belongs to
  const companionWeaponSlot = entity.getDynamicProperty("weapon_slot");
  if (companionWeaponSlot === undefined) return;

  let weaponItem = null;
  
  // Get the specific weapon that this companion belongs to
  if (companionWeaponSlot === -1) {
    // Weapon is in main hand
    const equippable = player.getComponent("minecraft:equippable");
    weaponItem = equippable?.getEquipment(EquipmentSlot.Mainhand);
  } else {
    // Weapon is in inventory
    const inventory = player.getComponent("inventory");
    if (!inventory || !inventory.container) return;
    weaponItem = inventory.container.getItem(companionWeaponSlot);
  }
  
  if (!weaponItem) return;
  const weaponInfo = getWeaponInfo(weaponItem);
  if (!weaponInfo) return;
  
  const currentSouls = weaponInfo.souls ?? 0;

  if (currentSouls <= 0) {
    player.sendMessage("Your companion has no souls left to fight.");
    entity.triggerEvent(dontAttackEvent);
    return;
  }

  const form = new ActionFormData().title("Companion Control").body("Select an option for your companion:");
  UI_OPTIONS.forEach(option => form.button(option.text, option.texture));

  form.show(player).then(response => {
    if (response.canceled) return;

    const selectedOption = UI_OPTIONS[response.selection];
    if (!selectedOption) return;

    if (currentSouls <= 0) {
      player.sendMessage("Your companion has no souls left and cannot fight.");
      entity.triggerEvent(dontAttackEvent);
      return;
    }

    // Handle special abilities separately
    if (selectedOption.event === "special_ability") {
      const companionType = entity.typeId.replace(companionPrefix, "");
      const abilityInfo = getAbilityInfo(companionType);
      const cooldown = getAbilityCooldown(entity.id, "special");
      
      if (cooldown > 0) {
        player.sendMessage(`§cSpecial ability is on cooldown for ${Math.ceil(cooldown / 1000)} seconds.`);
        return;
      }
      
      if (abilityInfo) {
        entity.setProperty("aft_soulhunters:attack_state", "special");
        player.sendMessage(`§6${entity.nameTag} used ${abilityInfo.name}!`);
        player.sendMessage(`§7${abilityInfo.description}`);
      } else {
        player.sendMessage("§cSpecial ability not available for this companion type.");
      }
      return;
    }
    
    if (selectedOption.event === "scream_ability") {
      const cooldown = getAbilityCooldown(entity.id, "scream");
      
      if (cooldown > 0) {
        player.sendMessage(`§cRoar attack is on cooldown for ${Math.ceil(cooldown / 1000)} seconds.`);
        return;
      }
      
      entity.setProperty("aft_soulhunters:attack_state", "scream");
      player.sendMessage(`§c${entity.nameTag} unleashes a terrifying roar!`);
      return;
    }

    entity.triggerEvent(selectedOption.event);
    entity.setDynamicProperty(BEHAVIOR_PROP, selectedOption.event);
    player.sendMessage(`Mode changed to: ${selectedOption.text}`);
    
    // Play companion interaction sound
    player.playSound("aft_soulhunters_companion_hihi");
  });
}

function startSoulDrain(companion) {
  const animal = companion.typeId.replace(companionPrefix, "");
  const companionWeaponSlot = companion.getDynamicProperty("weapon_slot");
  
  const intervalId = system.runInterval(() => {
    if (!companion.isValid()) return system.clearRun(intervalId);
    const isActive = companion.getDynamicProperty(ACTIVE_PROP);
    if (!isActive) return system.clearRun(intervalId);

    for (const player of world.getPlayers()) {
      // Use the weapon slot stored on the companion
      if (companionWeaponSlot === undefined) continue;

      let weaponItem = null;
      
      // Get the specific weapon that this companion belongs to
      if (companionWeaponSlot === -1) {
        // Weapon is in main hand
        const equippable = player.getComponent("minecraft:equippable");
        weaponItem = equippable?.getEquipment(EquipmentSlot.Mainhand);
      } else {
        // Weapon is in inventory
        const inventory = player.getComponent("inventory");
        if (!inventory || !inventory.container) continue;
        weaponItem = inventory.container.getItem(companionWeaponSlot);
      }
      
      if (!weaponItem) continue;
      const weaponInfo = getWeaponInfo(weaponItem);
      if (!weaponInfo || weaponInfo.companion !== animal) continue;

      let souls = weaponInfo.souls ?? 0;
      const scoreName = SCORE_TO_EVENT[souls];

      try {
        companion.runCommand(`event entity @s aft_soulhunters:souls_bar_${scoreName}`);
      } catch (e) {}

      if (souls > 0) {
        souls--;
        updateWeaponLore(weaponItem, animal, souls);
        
        // Set the weapon back to its proper location
        if (companionWeaponSlot === -1) {
          const equippable = player.getComponent("minecraft:equippable");
          equippable.setEquipment(EquipmentSlot.Mainhand, weaponItem);
        } else {
          const inventory = player.getComponent("inventory");
          if (inventory && inventory.container) {
            inventory.container.setItem(companionWeaponSlot, weaponItem);
          }
        }
      }

      if (souls <= 0) {
        companion.setDynamicProperty(ACTIVE_PROP, false);
        companion.triggerEvent(noSoulsEvent);
        companion.triggerEvent(dontAttackEvent);
        system.clearRun(intervalId);
      }

      break;
    }
  }, 400);
}