/*!
 * Copyright (c) 2025-present, Vanilagy and contributors
 *
 * This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at https://mozilla.org/MPL/2.0/.
 */

import { aacChannelMap, aacFrequencyTable, AudioCodec } from '../codec';
import { Demuxer } from '../demuxer';
import { Input } from '../input';
import { InputAudioTrack, InputAudioTrackBacking } from '../input-track';
import { PacketRetrievalOptions } from '../media-sink';
import {
	assert,
	AsyncMutex,
	binarySearchExact,
	binarySearchLessOrEqual,
	Bitstream,
	UNDETERMINED_LANGUAGE,
} from '../misc';
import { EncodedPacket, PLACEHOLDER_DATA } from '../packet';
import { AdtsReader, FrameHeader, MAX_FRAME_HEADER_SIZE } from './adts-reader';

const SAMPLES_PER_AAC_FRAME = 1024;

type Sample = {
	timestamp: number;
	duration: number;
	dataStart: number;
	dataSize: number;
};

export class AdtsDemuxer extends Demuxer {
	reader: AdtsReader;

	metadataPromise: Promise<void> | null = null;
	firstFrameHeader: FrameHeader | null = null;
	loadedSamples: Sample[] = []; // All samples from the start of the file to lastLoadedPos

	tracks: InputAudioTrack[] = [];

	readingMutex = new AsyncMutex();
	lastLoadedPos = 0;
	fileSize = 0;
	nextTimestampInSamples = 0;

	constructor(input: Input) {
		super(input);

		this.reader = new AdtsReader(input._mainReader);
	}

	async readMetadata() {
		return this.metadataPromise ??= (async () => {
			this.fileSize = await this.input.source.getSize();

			await this.loadNextChunk();

			// There has to be a frame if this demuxer got selected
			assert(this.firstFrameHeader);

			// Create the single audio track
			this.tracks = [new InputAudioTrack(new AdtsAudioTrackBacking(this))];
		})();
	}

	async loadNextChunk() {
		assert(this.lastLoadedPos < this.fileSize);

		const chunkSize = 0.5 * 1024 * 1024; // 0.5 MiB
		const endPos = Math.min(this.lastLoadedPos + chunkSize, this.fileSize);
		await this.reader.reader.loadRange(this.lastLoadedPos, endPos);

		this.lastLoadedPos = endPos;
		assert(this.lastLoadedPos <= this.fileSize);

		this.parseFramesFromLoadedData();
	}

	private parseFramesFromLoadedData() {
		while (this.reader.pos <= this.fileSize - MAX_FRAME_HEADER_SIZE) {
			const startPos = this.reader.pos;
			const header = this.reader.readFrameHeader();
			if (!header) {
				break;
			}

			// Check if the entire frame fits in the loaded data
			if (startPos + header.frameLength > this.lastLoadedPos) {
				// Frame doesn't fit, reset positions and stop
				this.reader.pos = startPos;
				this.lastLoadedPos = startPos;
				break;
			}

			if (!this.firstFrameHeader) {
				this.firstFrameHeader = header;
			}

			const sampleRate = aacFrequencyTable[header.samplingFrequencyIndex];
			assert(sampleRate !== undefined);
			const sampleDuration = SAMPLES_PER_AAC_FRAME / sampleRate;
			const headerSize = header.crcCheck ? MAX_FRAME_HEADER_SIZE : MAX_FRAME_HEADER_SIZE - 2;

			const sample: Sample = {
				timestamp: this.nextTimestampInSamples / sampleRate,
				duration: sampleDuration,
				dataStart: startPos + headerSize,
				dataSize: header.frameLength - headerSize,
			};

			this.loadedSamples.push(sample);
			this.nextTimestampInSamples += SAMPLES_PER_AAC_FRAME;
			this.reader.pos = startPos + header.frameLength;
		}
	}

	async getMimeType() {
		return 'audio/aac';
	}

	async getTracks() {
		await this.readMetadata();
		return this.tracks;
	}

	async computeDuration() {
		await this.readMetadata();

		const track = this.tracks[0];
		assert(track);

		return track.computeDuration();
	}
}

class AdtsAudioTrackBacking implements InputAudioTrackBacking {
	constructor(public demuxer: AdtsDemuxer) {}

	getId() {
		return 1;
	}

	async getFirstTimestamp() {
		return 0;
	}

	getTimeResolution() {
		const sampleRate = this.getSampleRate();
		return sampleRate / SAMPLES_PER_AAC_FRAME;
	}

	async computeDuration() {
		const lastPacket = await this.getPacket(Infinity, { metadataOnly: true });
		return (lastPacket?.timestamp ?? 0) + (lastPacket?.duration ?? 0);
	}

	getName() {
		return null;
	}

	getLanguageCode() {
		return UNDETERMINED_LANGUAGE;
	}

	getCodec(): AudioCodec {
		return 'aac';
	}

	getNumberOfChannels() {
		assert(this.demuxer.firstFrameHeader);

		const numberOfChannels = aacChannelMap[this.demuxer.firstFrameHeader.channelConfiguration];
		assert(numberOfChannels !== undefined);

		return numberOfChannels;
	}

	getSampleRate() {
		assert(this.demuxer.firstFrameHeader);

		const sampleRate = aacFrequencyTable[this.demuxer.firstFrameHeader.samplingFrequencyIndex];
		assert(sampleRate !== undefined);

		return sampleRate;
	}

	async getDecoderConfig(): Promise<AudioDecoderConfig> {
		assert(this.demuxer.firstFrameHeader);

		const bytes = new Uint8Array(3); // 19 bits max
		const bitstream = new Bitstream(bytes);

		const { objectType, samplingFrequencyIndex, channelConfiguration } = this.demuxer.firstFrameHeader;

		if (objectType > 31) {
			bitstream.writeBits(5, 31);
			bitstream.writeBits(6, objectType - 32);
		} else {
			bitstream.writeBits(5, objectType);
		}

		bitstream.writeBits(4, samplingFrequencyIndex); // samplingFrequencyIndex === 15 is forbidden

		bitstream.writeBits(4, channelConfiguration);

		return {
			codec: `mp4a.40.${this.demuxer.firstFrameHeader.objectType}`,
			numberOfChannels: this.getNumberOfChannels(),
			sampleRate: this.getSampleRate(),
			description: bytes.subarray(0, Math.ceil((bitstream.pos - 1) / 8)),
		};
	}

	getPacketAtIndex(sampleIndex: number, options: PacketRetrievalOptions) {
		if (sampleIndex === -1) {
			return null;
		}

		const rawSample = this.demuxer.loadedSamples[sampleIndex];
		if (!rawSample) {
			return null;
		}

		let data: Uint8Array;
		if (options.metadataOnly) {
			data = PLACEHOLDER_DATA;
		} else {
			this.demuxer.reader.pos = rawSample.dataStart;
			data = this.demuxer.reader.readBytes(rawSample.dataSize);
		}

		return new EncodedPacket(
			data,
			'key',
			rawSample.timestamp,
			rawSample.duration,
			sampleIndex,
			rawSample.dataSize,
		);
	}

	async getFirstPacket(options: PacketRetrievalOptions) {
		return this.getPacketAtIndex(0, options);
	}

	async getNextPacket(packet: EncodedPacket, options: PacketRetrievalOptions) {
		const release = await this.demuxer.readingMutex.acquire();

		try {
			const sampleIndex = binarySearchExact(
				this.demuxer.loadedSamples,
				packet.timestamp,
				x => x.timestamp,
			);
			if (sampleIndex === -1) {
				throw new Error('Packet was not created from this track.');
			}

			const nextIndex = sampleIndex + 1;
			// Ensure the next sample exists
			while (
				nextIndex >= this.demuxer.loadedSamples.length
				&& this.demuxer.lastLoadedPos < this.demuxer.fileSize
			) {
				await this.demuxer.loadNextChunk();
			}

			return this.getPacketAtIndex(nextIndex, options);
		} finally {
			release();
		}
	}

	async getPacket(timestamp: number, options: PacketRetrievalOptions) {
		const release = await this.demuxer.readingMutex.acquire();

		try {
			while (true) {
				const index = binarySearchLessOrEqual(
					this.demuxer.loadedSamples,
					timestamp,
					x => x.timestamp,
				);
				if (index === -1 && this.demuxer.loadedSamples.length > 0) {
					// We're before the first sample
					return null;
				}

				if (this.demuxer.lastLoadedPos === this.demuxer.fileSize) {
					// All data is loaded, return what we found
					return this.getPacketAtIndex(index, options);
				}

				if (index >= 0 && index + 1 < this.demuxer.loadedSamples.length) {
					// The next packet also exists, we're done
					return this.getPacketAtIndex(index, options);
				}

				// Otherwise, keep loading data
				await this.demuxer.loadNextChunk();
			}
		} finally {
			release();
		}
	}

	getKeyPacket(timestamp: number, options: PacketRetrievalOptions) {
		return this.getPacket(timestamp, options);
	}

	getNextKeyPacket(packet: EncodedPacket, options: PacketRetrievalOptions) {
		return this.getNextPacket(packet, options);
	}
}