import { world, system, ItemStack, EquipmentSlot, Player } from "@minecraft/server";
import { ModalFormData } from "@minecraft/server-ui";

const PET_ENTITY_ID = "md_bap:seat";
const HOME_BLOCKS_KEY = "md_bap:home_blocks";
const cooldowns = new Map();
const cooldownDuration = 60; // Cooldown duration in ticks (3 seconds)
const loadedHomeBlocks = new Map();

const availableTricks = new Set ([
    "Roll Over",
    "Shake Paw",
    "Flip",
    "Play Dead"
]);

const toySet = new Set ([
    "md_bap:blue_pet_candy",
    "md_bap:red_pet_candy",
    "md_bap:green_pet_candy",
    "md_bap:pet_toy_ball",
    "md_bap:pet_toy_bone",
    "md_bap:pet_toy_rope"
]);

const storagePets = new Set ([
    "md_bap:elephant",
    "md_bap:dromedary_camel",
    "md_bap:donkey",
    "md_bap:musk_ox",
    "md_bap:tree_kangaroo",
    "md_bap:red_kangaroo",
    "md_bap:llama"
])

const gatherPets = new Set ([
    "md_bap:racoon",
    "md_bap:bat_eared_fox",
    "md_bap:fox",
    "md_bap:panda",
    "md_bap:gorilla",
    "md_bap:capybara"
])

const ignoreEnts = new Set ([
    "md_bap:entity_drop",
    "md_bap:seat"
])

let trickIndex = 0;

world.beforeEvents.playerInteractWithEntity.subscribe((e) => {
    const {target, itemStack, player} = e;
    const dimension = player.dimension;
    if (target.typeId.includes("md_bap:")) {
        //check if pet is sleeping and return if so
        const bedEntity = target.dimension.getEntities({location: target.location, maxDistance: 1 }).filter(e => e.typeId === PET_ENTITY_ID)[0];
        if (bedEntity) return;
        if (storagePets.has(target.typeId) && player.isSneaking) {
            e.cancel = true
            system.runTimeout(() => {
                const state = target.getComponent("mark_variant") || 0;
                if (state){
                    target.triggerEvent("md_bap:unsit")
                } else {
                    target.triggerEvent("md_bap:sit")
                }
                return;
            },1)
                        
        }
        if (toySet.has(itemStack?.typeId)) {
            e.cancel = true
            system.runTimeout(() => {
                const tamed = target.getComponent("is_tamed") || false;
                if (!tamed) return;
                if (itemStack?.typeId.includes("md_bap:") && itemStack?.typeId.includes("pet_candy")) {
                    decrementStack(player, itemStack)
                    const loc = target.getHeadLocation()
                    dimension.playSound("random.eat", loc);
                    dimension.spawnParticle("minecraft:crop_growth_emitter", loc)
                    target.playAnimation("beg", {"blendOutTime": 0.5});
                    return;
                } else if (itemStack?.typeId === "md_bap:pet_toy_bone" || itemStack?.typeId === "md_bap:pet_toy_ball" || itemStack?.typeId === "md_bap:pet_toy_rope") {
                    const lore = itemStack.getLore();
                    if (player.isSneaking) {
                        openTrickMenu(player, lore, itemStack?.typeId);
                        return;
                    } 
                    triggerTrick(target, lore);
                }
            },1)
        }
    }
});

world.afterEvents.playerInteractWithEntity.subscribe((e) => {
    const {target, player} = e;
    if (!target.typeId.includes("md_bap:")) return;
    
    const now = Date.now();
    const lastInteraction = cooldowns.get(target.id) || 0;
    if (now - lastInteraction < cooldownDuration) return;
    cooldowns.set(target.id, now);

    //check to see if we need to remove the bed entity
    const bedEntity = target.dimension.getEntities({location: target.location, maxDistance: 1 }).filter(e => e.typeId === PET_ENTITY_ID)[0];
    if (bedEntity && !target.hasTag("on_lead") && !player.isSneaking) {
        target.triggerEvent("md_bap:wakeup");
        return;
    }
})

world.beforeEvents.worldInitialize.subscribe((e) => {
    system.runTimeout(() => {
        loadHomeBlocks();
    }, 10);
});

world.beforeEvents.worldInitialize.subscribe((e) => {
  e.blockComponentRegistry.registerCustomComponent(
    "md_bap:pet_bed", petBedComponent
  );
  e.itemComponentRegistry.registerCustomComponent(
    "md_bap:pet_whistle", petWhistleComponent
  );
});

const petWhistleComponent = {
    onUse(e) {
        const {source} = e
        source.playSound("whistle.use")
        const nearbyPets = source.dimension.getEntities({location: source.location, maxDistance: 25 }).filter(e => e.typeId.includes("md_bap:") && !ignoreEnts.has(e.typeId));
        //put them all in their beds if they belong to player
        const petsInBed = [];
        for (const pet of nearbyPets) {
            const ownerId = pet.hasTag(source.id);
            if (ownerId) {
                petsInBed.push(pet)
            }
        }
        for (const pet of petsInBed) {
            try {
                const bedData = loadedHomeBlocks.get(pet.id)
                const bedBlock = source.dimension.getBlock({x: bedData.x, y: bedData.y, z: bedData.z})
                petBed(bedBlock, pet, source.dimension, true);
            } catch(e) {}
        }
    }
}

const petBedComponent = {
    onStepOn(e) {
      const {entity: pet, block, dimension} = e;
      
      if (!pet?.typeId.includes("md_bap:")) return;
      const currentTime = world.getTimeOfDay() % 24000;
      if (currentTime < 12000) return;
      const petId = pet.id;
      const now = Date.now();
      const lastInteraction = cooldowns.get(petId) || 0;
      if (now - lastInteraction < cooldownDuration) return;
      cooldowns.set(petId, now);

      const home = loadedHomeBlocks.get(petId);
      if (!home) return;
      if (
          home?.x === block.location.x &&
          home?.y === block.location.y &&
          home?.z === block.location.z &&
          home?.dimension === dimension.id
      ) {
          //if (pet.hasTag("on_lead")) {
          //    pet.removeTag("on_lead")
          //}
          petBed(block, pet, dimension, false);
      }        
    },
    onPlayerDestroy(e) {
      const block = e.block;
      const blockLocation = block.location;
      const dimension = block.dimension;

      const existingEntity = dimension.getEntities({
          location: block.location,
          maxDistance: 1,
          type: "md_bap:seat"
      })[0];
      if (existingEntity) {
          existingEntity.triggerEvent("despawn")
      }
      // Find the pet associated with this bed
      let petIdToRemove = null;
      for (const [petId, bedData] of loadedHomeBlocks) {
        if (
          bedData.x === blockLocation.x &&
          bedData.y === blockLocation.y &&
          bedData.z === blockLocation.z &&
          bedData.dimension === dimension.id
        ) {
          petIdToRemove = petId;
          break;
        }
      }

      // If a pet is associated with this bed, remove the association
      if (petIdToRemove) {
        loadedHomeBlocks.delete(petIdToRemove);
        saveHomeBlocks(petIdToRemove, block.location, dimension,true); // Save the updated home block data
      }
    },
    onPlayerInteract(e) {
      const {block, dimension, player} = e;
      if (player.isSneaking) return;

      const existingSeat = dimension.getEntities({
          location: block.location,
          maxDistance: 1,
          type: "md_bap:seat"
      })[0];
      if (existingSeat) {
          existingSeat.triggerEvent("despawn")
      }
      
      const nearbyPet = dimension.getEntities({location: player.location, maxDistance: 7 }).filter(e => e.hasTag("on_lead"))[0];
      if (!nearbyPet) return;
      const ownerId = nearbyPet.hasTag(player.id);
      if (!ownerId) return;
      
      if (nearbyPet.typeId.includes("md_bap:")) {
              // Initialize loadedHomeBlocks if it doesn't exist
              if (loadedHomeBlocks.size === undefined) {
                  const newPet = respawnEntityAtBed(nearbyPet, block, dimension, player);
                  //nearbyPet.removeTag("on_lead");
                  player.onScreenDisplay.setActionBar("Pet home has been set to this new bed");
                  petBed(block, newPet, dimension, false);
                  saveHomeBlocks(nearbyPet.id, block.location, dimension);
                  return;
              }
              // Check for existing bed at this location
              let existingBed = false;
              for (const [petId, bedData] of loadedHomeBlocks) {
                  if (
                      bedData.x === block.location.x &&
                      bedData.y === block.location.y &&
                      bedData.z === block.location.z &&
                      bedData.dimension === dimension.id &&
                      petId !== nearbyPet.id
                  ) {
                      existingBed = true;
                      break;
                  }
              }

              if (existingBed) {
                  player.onScreenDisplay.setActionBar(
                      "§cThis location is already assigned, break and re-place the bed to reset the assigned pet"
                  );
                  return;
              }

              let homeBlockData = loadedHomeBlocks.get(nearbyPet.id);

              if (
                  homeBlockData?.x === block.location.x &&
                  homeBlockData?.y === block.location.y &&
                  homeBlockData?.z === block.location.z &&
                  homeBlockData?.dimension === dimension.id
              ) {
                  petBed(block, nearbyPet, dimension, true);
              } else if (homeBlockData) {
                  player.onScreenDisplay.setActionBar("Pet is being moved from their original home to this new location");
                  loadedHomeBlocks.delete(nearbyPet.id);
                  
                  // Respawn entity at the new bed location
                  const newPet = respawnEntityAtBed(nearbyPet, block, dimension, player);
                  saveHomeBlocks(newPet.id, block.location, dimension);
                  petBed(block, newPet, dimension, false);
              } else {
                  // Respawn entity at the new bed location
                  const newPet = respawnEntityAtBed(nearbyPet, block, dimension, player);
                  saveHomeBlocks(newPet.id, block.location, dimension);
                  player.onScreenDisplay.setActionBar("Pet home has been set to this new bed");
                  petBed(block, newPet, dimension, false);
              }
          }
      }
  }

function decrementStack(player, selectedItem) {
    const equipment = player.getComponent("equippable");
    if (player.getGameMode() !== "creative") {
      if (selectedItem.amount === 1) {
        equipment.setEquipment("Mainhand", undefined);
      } else {
        selectedItem.amount -= 1;
        equipment.setEquipment("Mainhand", selectedItem);
      }
    }
  }

function openTrickMenu(player, storedLore, item) {
    let storedTricks = "";
    let currentLore = [];
    for(let i = 0; i < storedLore.length; i++) {
        if(storedLore[i].includes(`§r§6Trick(s)`)) {
            storedTricks = storedLore[i].replace(`§r§6Trick(s): `, "");
            currentLore = storedTricks.split(", ");
            break;
        }
    }

    const trickMenu = new ModalFormData()
        .title("Trick Selection Menu");

    // Add toggles for each available trick
    for (const trick of availableTricks) {
        trickMenu.toggle(trick, currentLore.includes(trick) || currentLore.includes("All"));
    }

    trickMenu.show(player).then(response => {
        if (response.canceled) return;
        const toggleResponses = response.formValues;
        const selectedTricks = Array.from(availableTricks).filter((_, index) => toggleResponses[index]);
        if (selectedTricks.length > 0 && selectedTricks.length < 4) {
            storedTricks = `§r§6Trick(s): ${selectedTricks.join(", ")}`;
        } else {
            storedTricks = `§r§6Trick(s): All`;
        }
        
        // Set the lore on the pet toy
        let petToy = new ItemStack(item);
        petToy.keepOnDeath = true;
        petToy.setLore([storedTricks]);
        player.getComponent('equippable').setEquipment(EquipmentSlot.Mainhand, petToy);
    }).catch(error => {
        console.error(error, error.stack);
    });
}


function triggerTrick(entity, lore) {
    let storedTricks = "";
    let currentLore = [];
    for(let i = 0; i < lore.length; i++) {
        if(lore[i].includes(`§r§6Trick(s)`)) {
            storedTricks = lore[i].replace(`§r§6Trick(s): `, "");
            currentLore = storedTricks.split(", ");
            break;
        }
    }

    let maxIndex = currentLore.length -1
    if (currentLore.length === 0 || currentLore[0] === "All") {
        currentLore = [];
        availableTricks.forEach(t => {
            currentLore.push(t)
        })
        maxIndex = 3;
    }
    if (trickIndex < maxIndex) {
        trickIndex++
    } else {
        trickIndex = 0;
    }
    entity.playAnimation(`${currentLore[trickIndex].toLowerCase().replace(" ", "_")}`, {"blendOutTime": 0.5});
};

// Function to respawn the entity at the given location
function respawnEntityAtBed(nearbyPet, block, dimension, player) {
        let lead = new ItemStack("minecraft:lead", 1);
        dimension.spawnItem(lead, nearbyPet.location)
        const clonedItemSet = [];
        if (storagePets.has(nearbyPet.typeId)) {
            const inventory = nearbyPet.getComponent("inventory").container;
            for (let i = 0; i < inventory.size; i++) {
                const item = inventory.getItem(i);
                if (item) {
                     clonedItemSet.push(item.clone())
                }
            }
        }

        //Get Data
        let health = nearbyPet.getComponent("minecraft:health").current;
        let nameTag = nearbyPet.nameTag;

        system.runTimeout(() => {
            //Remove the pet
            if (gatherPets.has(nearbyPet.typeId)) {
                nearbyPet.triggerEvent("drop")
            } else {
                nearbyPet.remove();
            }
        },1)

        // Respawn the pet at the bed's location
        let newPet = dimension.spawnEntity(`${nearbyPet.typeId}<md_bap:spawn_as_projectile>`, {
            x: block.location.x + 0.5,
            y: block.location.y + 0.5,
            z: block.location.z + 0.5,
        });
        // Set Owner
        let projectileComponent = newPet.getComponent('minecraft:projectile');
        projectileComponent.owner = player

        newPet.triggerEvent('minecraft:on_tame')

        if (clonedItemSet.length > 0) {
            const inventory = newPet.getComponent("inventory").container;
            for (let i = 0; i < clonedItemSet.length; i++) {
                inventory.setItem(i,clonedItemSet[i]);
            }
        }

        //Reapply data
        newPet.getComponent("minecraft:health").current = health;
        newPet.nameTag = nameTag;
        newPet.addTag(player.id)
        system.runTimeout(() => {
            newPet.triggerEvent('md_bap:remove_projectile')
        }, 4)

        return newPet; // Return the newly spawned entity
        
    
}

function petBed (block, pet, dimension, manual) {
    const loc = block?.location;
    if (pet instanceof Player) return;
    system.runTimeout(() => {
        const cardinalDirection = block.permutation.getState(
          "minecraft:cardinal_direction"
        );

        let seatRotation = 0;
        switch (cardinalDirection) {
          case "north":
            seatRotation = 0;
            break;
          case "west":
            seatRotation = 270;
            break;
          case "south":
            seatRotation = 180;
            break;
          case "east":
            seatRotation = 90;
            break;
          default:
            seatRotation = 0;
            break;
        }

        const seat = dimension.spawnEntity("md_bap:seat", {
          x: loc.x + 0.5,
          y: loc.y + 0.15,
          z: loc.z + 0.5,
        });

        seat.setRotation({ x: 0, y: seatRotation });
        seat.getComponent("rideable").addRider(pet);
        if (!manual) {
            pet.triggerEvent("md_bap:sleep")
        } 
      }, 5);
}

function saveHomeBlocks(petId, location, dimension, save) {
    if (!save) {
        // Add the new entry to the map
        loadedHomeBlocks.set(petId, {
            x: location.x,
            y: location.y,
            z: location.z,
            dimension: dimension.id
        });
    }
    // Convert the map to an object for JSON serialization
    const homeBlocksObject = Object.fromEntries(loadedHomeBlocks);

    // Save the updated map to dynamic properties
    world.setDynamicProperty(HOME_BLOCKS_KEY, JSON.stringify(homeBlocksObject));
}

function loadHomeBlocks() {
    try {
        const rawData = world.getDynamicProperty(HOME_BLOCKS_KEY);
        if (rawData) {
            const parsedData = JSON.parse(rawData);
            loadedHomeBlocks.clear();
            for (const [key, value] of Object.entries(parsedData)) {
                loadedHomeBlocks.set(key, value);
            }
        }
    } catch (e) {
        console.warn("§4[Error]§r Failed to load home block data: " + e);
    }
}
