import type { AudioDocument } from "@suaveplan/audio-core/document/document";
import { frameCount, frameMagnitude, readMono } from "./dsp.js";
import type { SpectrumAnalyzerOptionsInput } from "./options.schemas.js";
import type { SpectrogramResult } from "./result.types.js";
import {
    buildBase,
    ensureNotAborted,
    parseOptions,
    startRun,
} from "./runtime.js";

export interface SpectrogramRunOptions extends SpectrumAnalyzerOptionsInput {
    readonly signal?: AbortSignal;
}

function emptyResult(doc: AudioDocument): SpectrogramResult {
    return {
        ...buildBase("spectrogram", doc),
        op: "spectrogram",
        spectrogram: [],
    };
}

export class Spectrogram {
    public readonly name = "spectrogram";
    public readonly description =
        "Time-frequency magnitude matrix of an audio document.";

    public async analyze(
        doc: AudioDocument,
        raw?: SpectrogramRunOptions
    ): Promise<SpectrogramResult> {
        const opts = parseOptions(raw);
        ensureNotAborted(raw?.signal);
        const handle = startRun("spectrogram", doc, "analyze");
        try {
            const samples = await readMono(doc);
            const count = frameCount(
                samples.length,
                opts.frameSize,
                opts.hopSize
            );
            const spectrogram: Float32Array[] = [];
            for (let i = 0; i < count; i++) {
                ensureNotAborted(raw?.signal);
                const offset = i * opts.hopSize;
                spectrogram.push(
                    frameMagnitude(samples, offset, opts.frameSize, opts.window)
                );
            }
            handle.end();
            return {
                ...buildBase("spectrogram", doc),
                op: "spectrogram",
                spectrogram,
            };
        } catch (err) {
            handle.fail(err);
            throw err;
        }
    }

    public async *analyzeStream(
        doc: AudioDocument,
        raw?: SpectrogramRunOptions
    ): AsyncGenerator<SpectrogramResult> {
        const opts = parseOptions(raw);
        ensureNotAborted(raw?.signal);
        const handle = startRun("spectrogram", doc, "stream");
        try {
            const samples = await readMono(doc);
            const count = frameCount(
                samples.length,
                opts.frameSize,
                opts.hopSize
            );
            if (count === 0) {
                yield emptyResult(doc);
                handle.end();
                return;
            }
            for (let i = 0; i < count; i++) {
                ensureNotAborted(raw?.signal);
                const offset = i * opts.hopSize;
                const mag = frameMagnitude(
                    samples,
                    offset,
                    opts.frameSize,
                    opts.window
                );
                yield {
                    ...buildBase("spectrogram", doc),
                    op: "spectrogram",
                    spectrogram: [mag],
                };
            }
            handle.end();
        } catch (err) {
            handle.fail(err);
            throw err;
        }
    }
}