import { XMLParser } from 'fast-xml-parser';

export class SignedKeyID {
	constructor(
		public alg_id: string,
		public value: string,
		public checksum?: string
	) {}
}

export type Version = '4.0.0.0' | '4.1.0.0' | '4.2.0.0' | '4.3.0.0' | 'UNKNOWN';

export type ReturnStructure = [SignedKeyID[], string | null, string | null, string | null];

interface ParsedWRMHeader {
	WRMHEADER: {
		'@_version': string;
		DATA?: any;
	};
}

export default class WRMHeader {
	private header: ParsedWRMHeader['WRMHEADER'];
	version: Version;

	constructor(data: string) {
		if (!data) throw new Error('Data must not be empty');

		const parser = new XMLParser({
			ignoreAttributes: false,
			removeNSPrefix: true,
			attributeNamePrefix: '@_'
		});
		const parsed = parser.parse(data) as ParsedWRMHeader;

		if (!parsed.WRMHEADER) throw new Error('Data is not a valid WRMHEADER');

		this.header = parsed.WRMHEADER;
		this.version = WRMHeader.fromString(this.header['@_version']);
	}

	private static fromString(value: string): Version {
		if (['4.0.0.0', '4.1.0.0', '4.2.0.0', '4.3.0.0'].includes(value)) {
			return value as Version;
		}
		return 'UNKNOWN';
	}

	to_v4_0_0_0(): string {
		const [key_ids, la_url, lui_url, ds_id] = this.readAttributes();
		if (key_ids.length === 0) throw new Error('No Key IDs available');
		const key_id = key_ids[0];
		return `<WRMHEADER xmlns="http://schemas.microsoft.com/DRM/2007/03/PlayReadyHeader" version="4.0.0.0"><DATA><PROTECTINFO><KEYLEN>16</KEYLEN><ALGID>AESCTR</ALGID></PROTECTINFO><KID>${
			key_id.value
		}</KID>${la_url ? `<LA_URL>${la_url}</LA_URL>` : ''}${lui_url ? `<LUI_URL>${lui_url}</LUI_URL>` : ''}${ds_id ? `<DS_ID>${ds_id}</DS_ID>` : ''}${
			key_id.checksum ? `<CHECKSUM>${key_id.checksum}</CHECKSUM>` : ''
		}</DATA></WRMHEADER>`;
	}

	readAttributes(): ReturnStructure {
		const data = this.header.DATA;
		if (!data) throw new Error('Not a valid PlayReady Header Record, WRMHEADER/DATA required');
		switch (this.version) {
			case '4.0.0.0':
				return WRMHeader.read_v4(data);
			case '4.1.0.0':
			case '4.2.0.0':
			case '4.3.0.0':
				return WRMHeader.read_vX(data);
			default:
				throw new Error(`Unsupported version: ${this.version}`);
		}
	}

	private static read_v4(data: any): ReturnStructure {
		const protectInfo = data.PROTECTINFO;
		return [[new SignedKeyID(protectInfo.ALGID, data.KID, data.CHECKSUM)], data.LA_URL || null, data.LUI_URL || null, data.DS_ID || null];
	}

	private static read_vX(data: any): ReturnStructure {
		const protectInfo = data.PROTECTINFO;

		const signedKeyID: SignedKeyID | undefined = protectInfo.KIDS.KID
			? new SignedKeyID(protectInfo.KIDS.KID['@_ALGID'] || '', protectInfo.KIDS.KID['@_VALUE'], protectInfo.KIDS.KID['@_CHECKSUM'])
			: undefined;
		return [signedKeyID ? [signedKeyID] : [], data.LA_URL || null, data.LUI_URL || null, data.DS_ID || null];
	}
}
