import {execSync, spawn} from 'node:child_process';
import {mkdtempSync, rmSync} from 'node:fs';
import fs from 'node:fs/promises';
import {tmpdir} from 'node:os';
import Path, {join} from 'node:path';
import {AbortablePromise, Ai} from './ai.ts';

export class Audio {
	private downloads: {[key: string]: Promise<string>} = {};
	private pyannote!: string;
	private whisperModel!: string;

	constructor(private ai: Ai) {
		if(ai.options.whisper) {
			this.whisperModel = ai.options.asr?.endsWith('.bin') ? ai.options.asr : ai.options.asr + '.bin';
			this.downloadAsrModel();
		}

		this.pyannote = `
import sys
import json
import os
from pyannote.audio import Pipeline

os.environ['TORCH_HOME'] = r"${ai.options.path}"
pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization-3.1", token="${ai.options.hfToken}")
output = pipeline(sys.argv[1])

segments = []
for turn, speaker in output.speaker_diarization:
    segments.append({"start": turn.start, "end": turn.end, "speaker": speaker})

print(json.dumps(segments))
`;
	}

	private runAsr(file: string, opts: {model?: string, diarization?: boolean} = {}): AbortablePromise<any> {
		let proc: any;
		const p = new Promise<any>((resolve, reject) => {
			this.downloadAsrModel(opts.model).then(m => {
				let output = '';
				const args = [opts.diarization ? '-owts' : '-nt', '-m', m, '-f', file];
				proc = spawn(<string>this.ai.options.whisper, args, {stdio: ['ignore', 'pipe', 'ignore']});
				proc.on('error', (err: Error) => reject(err));
				proc.stdout.on('data', (data: Buffer) => output += data.toString());
				proc.on('close', (code: number) => {
					if(code === 0) {
						if(opts.diarization) {
							try { resolve(JSON.parse(output)); }
							catch(e) { reject(new Error('Failed to parse whisper JSON')); }
						} else {
							resolve(output.trim() || null);
						}
					} else {
						reject(new Error(`Exit code ${code}`));
					}
				});
			});
		});
		return <any>Object.assign(p, {abort: () => proc?.kill('SIGTERM')});
	}

	private runDiarization(file: string): AbortablePromise<any> {
		let aborted = false, abort = () => { aborted = true; };
		const checkPython = (cmd: string) => {
			return new Promise<boolean>((resolve) => {
				const proc = spawn(cmd, ['-c', 'import pyannote.audio']);
				proc.on('close', (code: number) => resolve(code === 0));
				proc.on('error', () => resolve(false));
			});
		};
		const p = Promise.all<any>([
			checkPython('python'),
			checkPython('python3'),
		]).then(<any>(async ([p, p3]: [boolean, boolean]) => {
			if(aborted) return;
			if(!p && !p3) throw new Error('Pyannote is not installed: pip install pyannote.audio');
			const binary = p3 ? 'python3' : 'python';
			let tmp: string | null = null;
			return new Promise((resolve, reject) => {
				tmp = join(mkdtempSync(join(tmpdir(), 'audio-')), 'converted.wav');
				execSync(`ffmpeg -i "${file}" -ar 16000 -ac 1 -f wav "${tmp}"`, { stdio: 'ignore' });
				if(aborted) return;
				let output = '';
				const proc = spawn(binary, ['-c', this.pyannote, tmp]);
				proc.stdout.on('data', (data: Buffer) => output += data.toString());
				proc.stderr.on('data', (data: Buffer) => console.error(data.toString()));
				proc.on('close', (code: number) => {
					if(code === 0) {
						try { resolve(JSON.parse(output)); }
						catch (err) { reject(new Error('Failed to parse diarization output')); }
					} else {
						reject(new Error(`Python process exited with code ${code}`));
					}
				});
				proc.on('error', reject);
				abort = () => proc.kill('SIGTERM');
			}).finally(() => { if(tmp) rmSync(Path.dirname(tmp), { recursive: true, force: true }); });
		}));
		return <any>Object.assign(p, {abort});
	}

	private combineSpeakerTranscript(transcript: any, speakers: any[]): string {
		const speakerMap = new Map();
		let speakerCount = 0;
		speakers.forEach((seg: any) => {
			if(!speakerMap.has(seg.speaker)) speakerMap.set(seg.speaker, ++speakerCount);
		});

		const lines: string[] = [];
		let currentSpeaker = -1;
		let currentText = '';
		transcript.transcription.forEach((word: any) => {
			const time = word.offsets.from / 1000; // Convert ms to seconds
			const speaker = speakers.find((s: any) => time >= s.start && time <= s.end);
			const speakerNum = speaker ? speakerMap.get(speaker.speaker) : 1;
			if (speakerNum !== currentSpeaker) {
				if(currentText) lines.push(`[Speaker ${currentSpeaker}]: ${currentText.trim()}`);
				currentSpeaker = speakerNum;
				currentText = word.text;
			} else {
				currentText += ' ' + word.text;
			}
		});
		if(currentText) lines.push(`[Speaker ${currentSpeaker}]: ${currentText.trim()}`);
		return lines.join('\n');
	}

	asr(file: string, options: { model?: string; diarization?: boolean | 'id' } = {}): AbortablePromise<string | null> {
		if(!this.ai.options.whisper) throw new Error('Whisper not configured');

		const transcript = this.runAsr(file, {model: options.model, diarization: !!options.diarization});
		const diarization: any = options.diarization ? this.runDiarization(file) : Promise.resolve(null);
		const abort = () => {
			transcript.abort();
			diarization?.abort?.();
		};

		const response = Promise.all([transcript, diarization]).then(async ([t, d]) => {
			if(!options.diarization) return t;
			t = this.combineSpeakerTranscript(t, d);
			if(options.diarization === 'id') {
				if(!this.ai.language.defaultModel) throw new Error('Configure an LLM for advanced ASR speaker detection');
				let chunks = this.ai.language.chunk(t, 500, 0);
				if(chunks.length > 4) chunks = [...chunks.slice(0, 3), <string>chunks.at(-1)];
				const names = await this.ai.language.json(chunks.join('\n'), '{1: "Detected Name", 2: "Second Name"}', {
					system: 'Use the following transcript to identify speakers. Only identify speakers you are positive about, dont mention speakers you are unsure about in your response',
					temperature: 0.1,
				});
				Object.entries(names).forEach(([speaker, name]) => t = t.replaceAll(`[Speaker ${speaker}]`, `[${name}]`));
			}
			return t;
		});
		return <any>Object.assign(response, {abort});
	}

	async downloadAsrModel(model: string = this.whisperModel): Promise<string> {
		if(!this.ai.options.whisper) throw new Error('Whisper not configured');
		if(!model.endsWith('.bin')) model += '.bin';
		const p = Path.join(<string>this.ai.options.path, model);
		if(await fs.stat(p).then(() => true).catch(() => false)) return p;
		if(!!this.downloads[model]) return this.downloads[model];
		this.downloads[model] = fetch(`https://huggingface.co/ggerganov/whisper.cpp/resolve/main/${model}`)
			.then(resp => resp.arrayBuffer())
			.then(arr => Buffer.from(arr)).then(async buffer => {
				await fs.writeFile(p, buffer);
				delete this.downloads[model];
				return p;
			});
		return this.downloads[model];
	}
}