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'
    }

    private static ensureList(element: any): any[] {
        return Array.isArray(element) ? element : [element]
    }

    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 signedKeyIDs: SignedKeyID[] = protectInfo?.KID
            ? WRMHeader.ensureList(protectInfo.KID).map(
                  (kid: any) =>
                      new SignedKeyID(
                          kid['@_ALGID'] || '',
                          kid['@_VALUE'],
                          kid['@_CHECKSUM']
                      )
              )
            : []
        return [
            signedKeyIDs,
            data.LA_URL || null,
            data.LUI_URL || null,
            data.DS_ID || null
        ]
    }
}