import {setImmediate as yield_batch} from 'node:timers/promises';
import {performance} from 'node:perf_hooks';
import type {Bot} from '@wp2508/mineflayer';
import type {InfoI18n} from '../types.ts';
import {metadata, ResourceInfo} from '../resource_metrics.ts';

export interface MemoryInfo {
    estimated: true;
    total_bytes: number;
    chunks_bytes: number;
    entities_bytes: number;
    inventory_bytes: number;
    other_bytes: number;
    partial: boolean;
    categories_complete: Record<Category, boolean>;
}

type Category = 'chunks' | 'entities' | 'inventory' | 'other';
const excluded = new Set(['registry', 'world', '_client', 'client', 'bot', 'global', 'globalThis',
    '_events', '_eventsCount', '_maxListeners', '__proto__', 'constructor', 'prototype']);
const own = (object: unknown, key: PropertyKey): unknown => object && typeof object === 'object'
    ? Object.getOwnPropertyDescriptor(object, key)?.value : undefined;
const typed_buffer = Object.getOwnPropertyDescriptor(Object.getPrototypeOf(Uint8Array.prototype), 'buffer')!.get!;
const data_buffer = Object.getOwnPropertyDescriptor(DataView.prototype, 'buffer')!.get!;
const array_length = Object.getOwnPropertyDescriptor(ArrayBuffer.prototype, 'byteLength')!.get!;
const shared_length = Object.getOwnPropertyDescriptor(SharedArrayBuffer.prototype, 'byteLength')!.get!;

/** Bounded live-data estimate. Objects/slots use approximate JS overhead; backing stores use real byteLength. */
export async function estimate_memory(bot: Bot, cancelled = () => false,
    limits = {nodes: 50000, properties: 200000, milliseconds: 200}): Promise<MemoryInfo> {
    const info: MemoryInfo = {estimated: true, total_bytes: 0, chunks_bytes: 0, entities_bytes: 0,
        inventory_bytes: 0, other_bytes: 0, partial: true,
        categories_complete: {chunks: false, entities: false, inventory: false, other: false}};
    const start = performance.now();
    const seen = new WeakSet<object>();
    const blocked = new WeakSet<object>();
    for (const key of ['registry', 'world', '_client']) {
        const value = own(bot, key);
        if (value && typeof value === 'object') blocked.add(value);
    }
    blocked.add(bot);
    blocked.add(globalThis);
    type Task = {value: unknown} | {object: object; keys: PropertyKey[]; index: number; columns?: boolean}
        | {load_columns: true};
    const categories: Category[] = ['chunks', 'entities', 'inventory', 'other'];
    const roots = {entities: ['entity', 'entities', 'players'], inventory: ['inventory', 'currentWindow'],
        other: ['game', 'time', 'settings', 'scoreboards', 'teams', 'tablist']};
    // Fixed category order and DFS order make first-visit ownership deterministic, including backing stores.
    const states = categories.map((category, index) => ({category,
        stack: (category === 'chunks' ? [{load_columns: true}]
            : [{object: bot, keys: roots[category], index: 0}]) as Task[],
        nodes: 0, properties: 0, elapsed: 0, incomplete: false,
        node_limit: Math.floor(limits.nodes / 4) + (index < limits.nodes % 4 ? 1 : 0),
        property_limit: Math.floor(limits.properties / 4) + (index < limits.properties % 4 ? 1 : 0),
    }));
    let steps = 0;
    let processed = 0;
    while (states.some(state => state.stack.length)) {
        if (cancelled() || performance.now() - start >= limits.milliseconds) break;
        const state = states[steps++ % states.length];
        if (!state.stack.length) continue;
        if (state.elapsed >= limits.milliseconds / 4) {
            state.incomplete = true;
            state.stack.length = 0;
            continue;
        }
        const turn = performance.now();
        const {stack, category} = state;
        const add = (bytes: number) => { info[`${category}_bytes`] += bytes; };
        const cap = () => { state.incomplete = true; stack.length = 0; };
        const task = stack.pop()!;
        try {
            if ('load_columns' in task) {
                // WorldSync exposes getColumns on its prototype; resolve descriptors without invoking getters.
                const world = own(bot, 'world');
                const world_descriptor = Object.getOwnPropertyDescriptor(bot, 'world');
                if (world_descriptor && !('value' in world_descriptor)) state.incomplete = true;
                let target = world;
                let descriptor: PropertyDescriptor | undefined;
                for (let depth = 0; target && depth < 16; depth++, target = Object.getPrototypeOf(target)) {
                    descriptor = Object.getOwnPropertyDescriptor(target, 'getColumns');
                    if (descriptor) break;
                }
                if (world && typeof descriptor?.value !== 'function') state.incomplete = true;
                else if (world) {
                    const columns = Reflect.apply(descriptor!.value, world, []);
                    const length = own(columns, 'length');
                    if (!Array.isArray(columns) || typeof length !== 'number') state.incomplete = true;
                    else {
                        if (length > 4096) state.incomplete = true;
                        // Extraction is incremental and charged to the chunks property budget.
                        stack.push({object: columns,
                            keys: Array.from({length: Math.min(length, 4096)}, (_, i) => String(i)),
                            index: 0, columns: true});
                    }
                }
            } else if ('object' in task) {
                if (task.index < task.keys.length) {
                    if (state.properties >= state.property_limit) cap();
                    else {
                        state.properties++;
                        const key = task.keys[task.index++];
                        if (task.index < task.keys.length) stack.push(task);
                        if (typeof key !== 'string' || !excluded.has(key)) {
                            const descriptor = Object.getOwnPropertyDescriptor(task.object, key);
                            if (descriptor && !('value' in descriptor)) state.incomplete = true;
                            else if (descriptor) {
                                if (task.columns) {
                                    const column = own(descriptor.value, 'column');
                                    if (column && typeof column === 'object') stack.push({value: column});
                                    else state.incomplete = true;
                                } else {
                                    if (task.object !== bot) add(8 + (typeof key === 'string' ? key.length * 2 : 0));
                                    if (typeof descriptor.value !== 'function') stack.push({value: descriptor.value});
                                }
                            }
                        }
                    }
                }
            } else {
                const {value} = task;
                if (typeof value === 'string') add(value.length * 2);
                else if (typeof value === 'number' || typeof value === 'bigint') add(8);
                else if (typeof value === 'boolean') add(4);
                else if (value && typeof value === 'object' && !seen.has(value) && !blocked.has(value)) {
                    if (state.nodes >= state.node_limit) { cap(); continue; }
                    seen.add(value);
                    state.nodes++;
                    add(32);
                    if (ArrayBuffer.isView(value)) {
                        // Includes a pooled Buffer's full backing allocation once, even for overlapping views.
                        stack.push({value: Reflect.apply(value instanceof DataView ? data_buffer : typed_buffer, value, [])});
                    } else if (value instanceof ArrayBuffer || value instanceof SharedArrayBuffer) {
                        add(Reflect.apply(value instanceof ArrayBuffer ? array_length : shared_length, value, []));
                    } else {
                        // Never read property values directly: accessors and prototypes are deliberately omitted.
                        const keys = Reflect.ownKeys(value);
                        const remaining = Math.max(0, state.property_limit - state.properties);
                        if (keys.length > remaining) { keys.length = remaining; state.incomplete = true; }
                        if (keys.length) stack.push({object: value, keys, index: 0});
                        if (value instanceof Map || value instanceof Set) state.incomplete = true;
                    }
                }
            }
        } catch { state.incomplete = true; }
        finally { state.elapsed += performance.now() - turn; }
        if (++processed % 256 === 0) await yield_batch();
    }
    for (const state of states) info.categories_complete[state.category] = !state.incomplete && !state.stack.length;
    info.partial = Object.values(info.categories_complete).some(complete => !complete);
    info.total_bytes = info.chunks_bytes + info.entities_bytes + info.inventory_bytes + info.other_bytes;
    return info;
}

export class init extends ResourceInfo<MemoryInfo> {
    readonly name = 'used_memory';
    readonly i18n: InfoI18n = {...metadata('自有数据内存估算', '仅通过自身属性描述符分批估算当前数据，不代表堆内存或进程常驻内存。'
        + '扫描根为已加载的区块列对象（最多 4096 个）、entity/entities/players 和 inventory/currentWindow。'
        + '其他分类仅扫描本实例的 game、time、settings、scoreboards、teams、tablist 自有数据根。'
        + '排除注册表、世界对象遍历、客户端、全局对象、监听器、函数和原型；不调用属性 getter。'
        + '按区块、实体、物品栏、其他的固定顺序逐步轮转，共享对象及底层缓冲区归首次访问分类且只计一次，池化缓冲区按完整底层分配大小计入。'
        + '对象和字符串开销为近似值，不统计 Map/Set 内部数据。上限为 50000 个节点、200000 个属性和 200 毫秒，'
        + '节点、属性和同步处理时间各分类至多使用四分之一，总时限包含让出事件循环的时间。'
        + '每轮只处理一个值或属性，每 256 步让出事件循环；getColumns 和 ownKeys 枚举本身仍同步执行，单次调用可能超时。'
        + '完成状态仅针对上述有限范围；未完成分类的零值不表示没有内存占用。', {}), fields: {
        ...metadata('', '', {
            estimated: ['是否为估算', '固定为 true，表示该指标始终是估算值'],
            total_bytes: ['估算总内存', '各分类已扫描部分内存估算值之和，单位为字节；partial 为 true 时不是完整总量'],
            chunks_bytes: ['区块内存', '已加载区块列自身数据的内存估算，单位为字节'],
            entities_bytes: ['实体内存', '实体和玩家自身数据的内存估算，单位为字节'],
            inventory_bytes: ['物品栏内存', '物品栏和当前窗口自身数据的内存估算，单位为字节'],
            other_bytes: ['其他内存', '仅 game、time、settings、scoreboards、teams、tablist 自有数据的内存估算，单位为字节'],
            partial: ['结果是否不完整', '首次采样尚未完成，或因扫描上限、取消、不支持或无法访问而遗漏了数据'],
        }).fields,
        categories_complete: metadata('分类完成状态', '各分类是否完成约定范围的扫描；初始均为 false，false 时字节数仅为已扫描部分，零值不能说明实际无占用。', {
            chunks: ['区块扫描完成', '已扫描约定的全部区块列且未触及上限或遗漏数据时为 true'],
            entities: ['实体扫描完成', '已扫描 entity、entities、players 且未触及上限或遗漏数据时为 true'],
            inventory: ['物品栏扫描完成', '已扫描 inventory、currentWindow 且未触及上限或遗漏数据时为 true'],
            other: ['其他扫描完成', '已扫描 game、time、settings、scoreboards、teams、tablist 且未触及上限或遗漏数据时为 true'],
        }),
    }};
    constructor(bot: Bot) {
        super({estimated: true, total_bytes: 0, chunks_bytes: 0, entities_bytes: 0,
            inventory_bytes: 0, other_bytes: 0, partial: true,
            categories_complete: {chunks: false, entities: false, inventory: false, other: false}});
        this.start(5000, async () => { this.publish(await estimate_memory(bot, () => this.disposed)); }, true);
    }
}