import { TreeSummary, SummaryEntry } from "../../treeSummary";

const EXPECTED_OUTPUT_COUNT = 100;

export type AccessSummaryState = { total: number };
export type AccessTotals = { [operation: string]: { count: number; size: number } };

type AccessValue = { path: string; size: number };
type OperationStats = {
    totalCount: number;
    totalSize: number;
    countTree: TreeSummary<AccessValue, AccessSummaryState>;
    // Only exists for operations tracked with sizes; count-only operations (getInfo, del, ...) have no meaningful size breakdown.
    sizeTree?: TreeSummary<AccessValue, AccessSummaryState>;
};

const accounts = new Map<string, Map<string, OperationStats>>();

function makeTree(getValue: (value: AccessValue) => number): TreeSummary<AccessValue, AccessSummaryState> {
    return new TreeSummary<AccessValue, AccessSummaryState>({
        getPath: value => value.path,
        createSummary: () => ({ total: 0 }),
        addToSummary: (value, summary) => {
            summary.total += getValue(value);
        },
        mergeSummaries: (target, source) => {
            target.total += source.total;
        },
        getWeight: summary => summary.total,
        expectedOutputCount: EXPECTED_OUTPUT_COUNT,
    });
}

/** Counts one storage access, in memory only. size is the bytes involved (0 when the target does not exist); omit it entirely for operations that only count calls, which then only get a count tree. */
export function trackAccess(config: { account: string; operation: string; path: string; size?: number }): void {
    let operations = accounts.get(config.account);
    if (!operations) {
        operations = new Map();
        accounts.set(config.account, operations);
    }
    let stats = operations.get(config.operation);
    if (!stats) {
        stats = { totalCount: 0, totalSize: 0, countTree: makeTree(() => 1) };
        operations.set(config.operation, stats);
    }
    let value: AccessValue = { path: config.path, size: config.size || 0 };
    stats.totalCount++;
    stats.countTree.add(value);
    if (config.size !== undefined) {
        let sizeTree = stats.sizeTree;
        if (!sizeTree) {
            sizeTree = makeTree(v => v.size);
            stats.sizeTree = sizeTree;
        }
        stats.totalSize += config.size;
        sizeTree.add(value);
    }
}

/** Method decorator factory, for API methods whose single config-object argument has account and bucketName: tracks the access (as `bucketName/path`) after the method succeeds. Sizes come from the config's data (writes) or the result's data (reads); operations without either are count-only. Array results (listings - findInfo, getChangesAfter) are tracked as two breakdowns: "<op> queries" - one access per CALL, at the query prefix, sized by the number of results (so the tree shows which QUERY returns the most) - and "<op> results" - one count-only access per returned path (so the tree shows which PATHS come back most). */
export function trackAccessCall(operation: string) {
    return function (target: unknown, key: string, descriptor: PropertyDescriptor): void {
        let original = descriptor.value as (...args: unknown[]) => Promise<unknown>;
        descriptor.value = async function (...args: unknown[]): Promise<unknown> {
            let config = args[0] as { account: string; bucketName: string; path?: string; prefix?: string; data?: Buffer };
            let result = await original.apply(this, args);
            let base = `${config.bucketName}/`;
            if (Array.isArray(result)) {
                // The query, sized by how many results it returned (which query is heaviest); and each returned path, count-only (which paths come back most)
                trackAccess({ account: config.account, operation: `${operation} queries`, path: base + (config.prefix || config.path || ""), size: result.length });
                for (let entry of result as { path: string }[]) {
                    trackAccess({ account: config.account, operation: `${operation} results`, path: base + entry.path });
                }
                return result;
            }
            let size: number | undefined;
            if (config.data) {
                size = config.data.length;
            } else {
                let data = (result as { data?: Buffer } | undefined)?.data;
                if (data) size = data.length;
            }
            trackAccess({ account: config.account, operation, path: base + config.path, size });
            return result;
        };
    };
}

export function getAccessTotals(account: string): AccessTotals {
    let result: AccessTotals = {};
    let operations = accounts.get(account);
    if (!operations) return result;
    for (let [operation, stats] of operations) {
        result[operation] = { count: stats.totalCount, size: stats.totalSize };
    }
    return result;
}

export function readAccessSummaries(config: { account: string; operation: string; maxCount: number; weightBySize?: boolean }): SummaryEntry<AccessSummaryState>[] {
    let operations = accounts.get(config.account);
    let stats = operations && operations.get(config.operation);
    if (!stats) return [];
    // Count-only operations have no size tree, in which case weightBySize is ignored and the count breakdown is returned.
    let sizeTree = stats.sizeTree;
    if (config.weightBySize && sizeTree) {
        return sizeTree.getSummaries(config.maxCount);
    }
    return stats.countTree.getSummaries(config.maxCount);
}

export function clearAccountAccessStats(account: string): void {
    if (accounts.delete(account)) {
        console.log(`Cleared the in-memory access statistics for account ${account}`);
    }
}

export type BucketWriteStats = {
    /** Every set call the bucket accepted */
    originalWrites: number;
    originalBytes: number;
    /** What actually reached the sources. Fast writes coalesce repeated writes to the same key, so this is lower than the original counts (and is what the disk actually did). */
    flushedWrites: number;
    flushedBytes: number;
};
function emptyWriteStats(): BucketWriteStats {
    return { originalWrites: 0, originalBytes: 0, flushedWrites: 0, flushedBytes: 0 };
}

// In memory only, keyed `${account}/${bucketName}`: totals since this process started (or the last clearWriteStats). Persisting them to disk was more machinery than the numbers were worth.
const writeStats = new Map<string, BucketWriteStats>();

export function countBucketWrite(key: string, kind: "original" | "flushed", bytes: number): void {
    let stats = writeStats.get(key);
    if (!stats) {
        stats = emptyWriteStats();
        writeStats.set(key, stats);
    }
    if (kind === "original") {
        stats.originalWrites++;
        stats.originalBytes += bytes;
    } else {
        stats.flushedWrites++;
        stats.flushedBytes += bytes;
    }
}

export function getBucketWriteStats(key: string): BucketWriteStats {
    return writeStats.get(key) || emptyWriteStats();
}

/** Zeroes the write statistics of every bucket in the account. */
export function debugClearAccountWriteStats(account: string): number {
    let prefix = `${account}/`;
    let cleared = 0;
    for (let key of [...writeStats.keys()]) {
        if (!key.startsWith(prefix)) continue;
        writeStats.delete(key);
        cleared++;
    }
    console.log(`Cleared the write statistics of ${cleared} buckets in account ${account}`);
    return cleared;
}