tsgames/src/common/rpg/components/inventory.ts

import type { InventorySlotInput, RPGVariables, SlotId } from "../types";
import { action } from "../utils/decorators";
import { Component, type EvalContext } from "../core/world";
import { component } from "../core/registry";
import { Stackable, Usable } from "./item";
import { Equipment, Equippable } from "./equipment";
import { resolveVariables } from "../utils/variables";

interface SlotRecord {
    slotId: SlotId;
    limit: number | undefined;
    contents: { itemId: string; amount: number } | null;
}

interface InventoryState {
    infinite: boolean;
    nextSlotId: number;  // auto-increment counter for infinite mode
    slots: Record<SlotId, SlotRecord>;
}

function buildInventoryState(input?: number | InventorySlotInput[]): InventoryState {
    if (input === undefined) {
        return { infinite: true, nextSlotId: 0, slots: {} as Record<SlotId, SlotRecord> };
    }
    if (typeof input === 'number') {
        const slots = {} as Record<SlotId, SlotRecord>;
        for (let i = 0; i < input; i++) slots[i] = { slotId: i, limit: undefined, contents: null };
        return { infinite: false, nextSlotId: 0, slots };
    }
    const slots = {} as Record<SlotId, SlotRecord>;
    for (const def of input) {
        const slotId = typeof def === 'object' ? def.slotId : def;
        const limit = typeof def === 'object' ? def.limit : undefined;
        slots[slotId] = { slotId, limit, contents: null };
    }
    return { infinite: false, nextSlotId: 0, slots };
}

@component
export class Inventory extends Component<InventoryState> {
    #cachedVars: RPGVariables | null = null;

    /** Infinite inventory — grows on demand, no slot cap. */
    constructor();
    /** N generic slots, each limited only by item maxStack. */
    constructor(count: number);
    /** Named slots with optional per-slot stack limits (original behaviour). */
    constructor(slots: InventorySlotInput[]);
    constructor(input?: number | InventorySlotInput[]) {
        super(buildInventoryState(input));
    }

    #slot(slotId: SlotId): SlotRecord | undefined {
        return this.state.slots[slotId];
    }

    #capFor(slot: SlotRecord, itemId: string): number {
        const limitCap = slot.limit ?? Infinity;
        const stackable = this.entity.world.getEntity(itemId)?.get(Stackable);
        const stackCap = stackable ? stackable.maxStack : 1;
        return Math.min(limitCap, stackCap);
    }

    #roomFor(slot: SlotRecord, itemId: string): number {
        if (slot.contents !== null && slot.contents.itemId !== itemId) return 0;
        return this.#capFor(slot, itemId) - (slot.contents?.amount ?? 0);
    }

    @action
    add({ itemId, amount, slotId }: { itemId: string; amount: number; slotId?: SlotId }): boolean {
        this.#cachedVars = null;
        if (amount < 0) return false;
        if (amount === 0) return true;

        if (!this.entity.world.getEntity(itemId)) {
            console.warn(`[Inventory] add: item entity '${itemId}' not found`);
            return false;
        }

        // ── Direct slot ───────────────────────────────────────────────────────
        if (slotId !== undefined) {
            const slot = this.#slot(slotId);
            if (!slot) return false;
            if (this.#roomFor(slot, itemId) < amount) return false;
            slot.contents = { itemId, amount: (slot.contents?.amount ?? 0) + amount };
            this.emit('add', { itemId, amount, slotIds: [slotId] });
            return true;
        }

        // ── Finite inventory: two-phase (check → apply) ───────────────────────
        if (!this.state.infinite) {
            let canFit = 0;
            for (const slot of Object.values(this.state.slots)) {
                if (slot.contents === null || slot.contents.itemId === itemId)
                    canFit += this.#roomFor(slot, itemId);
            }
            if (canFit < amount) return false;

            let remaining = amount;
            const slotIds: SlotId[] = [];
            for (const slot of Object.values(this.state.slots)) {
                if (slot.contents?.itemId === itemId && remaining > 0) {
                    const take = Math.min(this.#roomFor(slot, itemId), remaining);
                    slot.contents!.amount += take;
                    remaining -= take;
                    slotIds.push(slot.slotId);
                }
            }
            for (const slot of Object.values(this.state.slots)) {
                if (slot.contents === null && remaining > 0) {
                    const take = Math.min(this.#capFor(slot, itemId), remaining);
                    slot.contents = { itemId, amount: take };
                    remaining -= take;
                    slotIds.push(slot.slotId);
                }
            }
            this.emit('add', { itemId, amount, slotIds });
            return true;
        }

        // ── Infinite inventory: fill existing slots, then grow ─────────────────
        let remaining = amount;
        const slotIds: SlotId[] = [];

        for (const slot of Object.values(this.state.slots)) {
            if (slot.contents?.itemId === itemId && remaining > 0) {
                const take = Math.min(this.#roomFor(slot, itemId), remaining);
                if (take > 0) {
                    slot.contents!.amount += take;
                    remaining -= take;
                    slotIds.push(slot.slotId);
                }
            }
        }
        for (const slot of Object.values(this.state.slots)) {
            if (slot.contents === null && remaining > 0) {
                const take = Math.min(this.#capFor(slot, itemId), remaining);
                slot.contents = { itemId, amount: take };
                remaining -= take;
                slotIds.push(slot.slotId);
            }
        }
        while (remaining > 0) {
            const newSlot: SlotRecord = { slotId: this.state.nextSlotId++, limit: undefined, contents: null };
            this.state.slots[newSlot.slotId] = newSlot;
            const take = Math.min(this.#capFor(newSlot, itemId), remaining);
            newSlot.contents = { itemId, amount: take };
            remaining -= take;
            slotIds.push(newSlot.slotId);
        }

        this.emit('add', { itemId, amount, slotIds });
        return true;
    }

    @action
    remove({ itemId, amount, slotId }: { itemId: string; amount: number; slotId?: SlotId }): boolean {
        this.#cachedVars = null;
        if (amount < 0) return false;
        if (amount === 0) return true;

        if (slotId !== undefined) {
            const slot = this.#slot(slotId);
            if (!slot || slot.contents?.itemId !== itemId) return false;
            if (slot.contents.amount < amount) return false;
            slot.contents.amount -= amount;
            if (slot.contents.amount === 0) slot.contents = null;
            this.emit('remove', { itemId, amount, slotIds: [slotId] });
            return true;
        }

        if (this.getAmount(itemId) < amount) return false;

        let remaining = amount;
        const slotIds: SlotId[] = [];
        for (const slot of Object.values(this.state.slots)) {
            if (slot.contents?.itemId === itemId && remaining > 0) {
                const take = Math.min(slot.contents.amount, remaining);
                slot.contents.amount -= take;
                if (slot.contents.amount === 0) slot.contents = null;
                remaining -= take;
                slotIds.push(slot.slotId);
            }
        }

        this.emit('remove', { itemId, amount, slotIds });
        return true;
    }

    /**
     * Equip an item from this inventory onto the entity's `Equipment` component.
     * `slotId` specifies which inventory slot to equip from (otherwise any slot is used).
     * `slotName` specifies the equipment body slot (otherwise auto-detected by item type).
     */
    @action
    equip(arg: string | { itemId?: string; slotId?: SlotId; slotName?: string }): boolean {
        const resolved = this.#resolveItem(arg);
        if (!resolved) return false;
        const { itemId, slotId } = resolved;
        const slotName = typeof arg === 'object' ? arg.slotName : undefined;

        if (this.getAmount(itemId, slotId) === 0) {
            console.warn(`[Inventory] equip: item '${itemId}' not in inventory`);
            return false;
        }

        const equipment = this.entity.get(Equipment);
        if (!equipment) {
            console.warn(`[Inventory] equip: entity has no Equipment component`);
            return false;
        }

        let resolvedSlot = slotName;
        if (resolvedSlot === undefined) {
            const equippable = this.entity.world.getEntity(itemId)?.get(Equippable);
            if (!equippable) {
                console.warn(`[Inventory] equip: item '${itemId}' has no Equippable component`);
                return false;
            }
            const found = equipment.findCompatibleSlot(equippable.slotType);
            if (!found) {
                console.warn(`[Inventory] equip: no compatible slot for '${itemId}'`);
                return false;
            }
            resolvedSlot = found;
        }

        return equipment.equip({ slotName: resolvedSlot, itemId });
    }

    /**
     * Use an item from this inventory.
     * `slotId` specifies which inventory slot to use from (otherwise any slot is used).
     */
    @action
    use(arg?: string | { itemId?: string; slotId?: SlotId }, ctx?: EvalContext): boolean {
        const resolved = this.#resolveItem(arg);
        if (!resolved) return false;
        const { itemId, slotId } = resolved;

        if (this.getAmount(itemId, slotId) === 0) {
            console.warn(`[Inventory] use: item '${itemId}' not in inventory`);
            return false;
        }

        const itemEntity = this.entity.world.getEntity(itemId);
        if (!itemEntity) {
            console.warn(`[Inventory] use: item entity '${itemId}' not found`);
            return false;
        }

        const usable = itemEntity.get(Usable);
        if (!usable) {
            console.warn(`[Inventory] use: item '${itemId}' has no Usable component`);
            return false;
        }

        if (usable.consumeOnUse) this.remove({ itemId, amount: 1, slotId });

        usable.use(ctx ?? this.context);
        return true;
    }

    /** Resolve an item reference, deriving `itemId` from slot contents if only `slotId` is given. */
    #resolveItem(arg?: string | { itemId?: string; slotId?: SlotId }): { itemId: string; slotId?: SlotId } | null {
        if (!arg) return null;
        if (typeof arg === 'string') return { itemId: arg };

        if (arg.slotId !== undefined) {
            const contents = this.getSlotContents(arg.slotId);
            if (!contents) return null;
            if (arg.itemId && contents.itemId !== arg.itemId) return null;
            return { itemId: contents.itemId, slotId: arg.slotId };
        }

        return arg.itemId ? { itemId: arg.itemId } : null;
    }

    getSlotContents(slotId: SlotId): { itemId: string; amount: number } | null {
        return this.#slot(slotId)?.contents ?? null;
    }

    getAmount(itemId: string, slotId?: SlotId): number {
        if (slotId !== undefined) {
            const slot = this.#slot(slotId);
            return slot?.contents?.itemId === itemId ? slot.contents.amount : 0;
        }
        let total = 0;
        for (const slot of Object.values(this.state.slots)) {
            if (slot.contents?.itemId === itemId) total += slot.contents.amount;
        }
        return total;
    }

    getItems(): Map<string, number> {
        const result = new Map<string, number>();
        for (const slot of Object.values(this.state.slots)) {
            if (slot.contents) {
                const { itemId, amount } = slot.contents;
                result.set(itemId, (result.get(itemId) ?? 0) + amount);
            }
        }
        return result;
    }

    override getVariables(): RPGVariables {
        if (this.#cachedVars) return this.#cachedVars;
        const result: RPGVariables = {};
        for (const [itemId, amount] of this.getItems()) {
            result[itemId] = amount;
            const itemEntity = this.entity.world.getEntity(itemId);
            if (itemEntity) {
                for (const [key, value] of Object.entries(resolveVariables(itemEntity))) {
                    result[`${itemId}.${key}`] = value;
                }
            }
        }
        this.#cachedVars = result;
        return result;
    }
}