import * as fs from 'fs'
import { createHash } from 'crypto'
import { Parser } from 'binary-parser-encoder'
import ECCKey from './ecc_key'

function alignUp(length: number, alignment: number): number {
    return Math.ceil(length / alignment) * alignment
}

export class BCertStructs {
    static DrmBCertBasicInfo = new Parser()
        .buffer('cert_id', { length: 16 })
        .uint32be('security_level')
        .uint32be('flags')
        .uint32be('cert_type')
        .buffer('public_key_digest', { length: 32 })
        .uint32be('expiration_date')
        .buffer('client_id', { length: 16 })

    static DrmBCertDomainInfo = new Parser()
        .buffer('service_id', { length: 16 })
        .buffer('account_id', { length: 16 })
        .uint32be('revision_timestamp')
        .uint32be('domain_url_length')
        .buffer('domain_url', {
            length: function () {
                return alignUp((this as any).domain_url_length, 4)
            }
        })

    static DrmBCertPCInfo = new Parser().uint32be('security_version')

    static DrmBCertDeviceInfo = new Parser()
        .uint32be('max_license')
        .uint32be('max_header')
        .uint32be('max_chain_depth')

    static DrmBCertFeatureInfo = new Parser()
        .uint32be('feature_count')
        .array('features', {
            type: 'uint32be',
            length: 'feature_count'
        })

    static CertKey = new Parser()
        .uint16be('type')
        .uint16be('length')
        .uint32be('flags')
        .buffer('key', {
            length: function () {
                return (this as any).length / 8
            }
        })
        .uint32be('usages_count')
        .array('usages', {
            type: 'uint32be',
            length: 'usages_count'
        })

    static DrmBCertKeyInfo = new Parser()
        .uint32be('key_count')
        .array('cert_keys', {
            type: BCertStructs.CertKey,
            length: 'key_count'
        })

    static DrmBCertManufacturerInfo = new Parser()
        .uint32be('flags')
        .uint32be('manufacturer_name_length')
        .buffer('manufacturer_name', {
            length: function () {
                return alignUp((this as any).manufacturer_name_length, 4)
            }
        })
        .uint32be('model_name_length')
        .buffer('model_name', {
            length: function () {
                return alignUp((this as any).model_name_length, 4)
            }
        })
        .uint32be('model_number_length')
        .buffer('model_number', {
            length: function () {
                return alignUp((this as any).model_number_length, 4)
            }
        })

    static DrmBCertSignatureInfo = new Parser()
        .uint16be('signature_type')
        .uint16be('signature_size')
        .buffer('signature', { length: 'signature_size' })
        .uint32be('signature_key_size')
        .buffer('signature_key', {
            length: function () {
                return (this as any).signature_key_size / 8
            }
        })

    static DrmBCertSilverlightInfo = new Parser()
        .uint32be('security_version')
        .uint32be('platform_identifier')

    static DrmBCertMeteringInfo = new Parser()
        .buffer('metering_id', { length: 16 })
        .uint32be('metering_url_length')
        .buffer('metering_url', {
            length: function () {
                return alignUp((this as any).metering_url_length, 4)
            }
        })

    static DrmBCertExtDataSignKeyInfo = new Parser()
        .uint16be('type')
        .uint16be('length')
        .uint32be('flags')
        .buffer('key', {
            length: function () {
                return (this as any).length / 8
            }
        })

    static BCertExtDataRecord = new Parser()
        .uint32be('data_size')
        .buffer('data', {
            length: 'data_size'
        })

    static DrmBCertExtDataSignature = new Parser()
        .uint16be('signature_type')
        .uint16be('signature_size')
        .buffer('signature', {
            length: 'signature_size'
        })

    static BCertExtDataContainer = new Parser()
        .uint32be('record_count')
        .array('records', {
            length: 'record_count',
            type: BCertStructs.BCertExtDataRecord
        })
        .nest('signature', {
            type: BCertStructs.DrmBCertExtDataSignature
        })

    static DrmBCertServerInfo = new Parser().uint32be('warning_days')

    static DrmBcertSecurityVersion = new Parser()
        .uint32be('security_version')
        .uint32be('platform_identifier')

    static Attribute = new Parser()
        .uint16be('flags')
        .uint16be('tag')
        .uint32be('length')
        .choice('attribute', {
            tag: 'tag',
            choices: {
                1: BCertStructs.DrmBCertBasicInfo,
                2: BCertStructs.DrmBCertDomainInfo,
                3: BCertStructs.DrmBCertPCInfo,
                4: BCertStructs.DrmBCertDeviceInfo,
                5: BCertStructs.DrmBCertFeatureInfo,
                6: BCertStructs.DrmBCertKeyInfo,
                7: BCertStructs.DrmBCertManufacturerInfo,
                8: BCertStructs.DrmBCertSignatureInfo,
                9: BCertStructs.DrmBCertSilverlightInfo,
                10: BCertStructs.DrmBCertMeteringInfo,
                11: BCertStructs.DrmBCertExtDataSignKeyInfo,
                12: BCertStructs.BCertExtDataContainer,
                13: BCertStructs.DrmBCertExtDataSignature,
                14: new Parser().buffer('data', {
                    length: function () {
                        return (this as any).length - 8
                    }
                }),
                15: BCertStructs.DrmBCertServerInfo,
                16: BCertStructs.DrmBcertSecurityVersion,
                17: BCertStructs.DrmBcertSecurityVersion
            },
            defaultChoice: new Parser().buffer('data', {
                length: function () {
                    return (this as any).length - 8
                }
            })
        })

    static BCert = new Parser()
        .string('signature', { length: 4, assert: 'CERT' })
        .int32be('version')
        .int32be('total_length')
        .int32be('certificate_length')
        .array('attributes', {
            type: BCertStructs.Attribute,
            lengthInBytes: function () {
                return (this as any).total_length - 16
            }
        })

    static BCertChain = new Parser()
        .string('signature', { length: 4, assert: 'CHAI' })
        .int32be('version')
        .int32be('total_length')
        .int32be('flags')
        .int32be('certificate_count')
        .array('certificates', {
            type: BCertStructs.BCert,
            length: 'certificate_count'
        })
}

export class Certificate {
    parsed: any
    _BCERT: Parser

    constructor(parsed_bcert: any, bcert_obj: Parser = BCertStructs.BCert) {
        this.parsed = parsed_bcert
        this._BCERT = bcert_obj
    }

    // UNSTABLE
    static new_key_cert(
        cert_id: Buffer,
        security_level: number,
        client_id: Buffer,
        signing_key: ECCKey,
        encryption_key: ECCKey,
        group_key: ECCKey,
        parent: CertificateChain,
        expiry: number = 0xffffffff,
        max_license: number = 10240,
        max_header: number = 15360,
        max_chain_depth: number = 2
    ): Certificate {
        if (!cert_id) {
            throw new Error('Certificate ID is required')
        }
        if (!client_id) {
            throw new Error('Client ID is required')
        }

        const basic_info = {
            cert_id: cert_id,
            security_level: security_level,
            flags: 0,
            cert_type: 2,
            public_key_digest: signing_key.publicSha256Digest(),
            expiration_date: expiry,
            client_id: client_id
        }
        const basic_info_attribute = {
            flags: 1,
            tag: 1,
            length:
                BCertStructs.DrmBCertBasicInfo.encode(basic_info).length + 8,
            attribute: basic_info
        }

        const device_info = {
            max_license: max_license,
            max_header: max_header,
            max_chain_depth: max_chain_depth
        }

        const device_info_attribute = {
            flags: 1,
            tag: 4,
            length:
                BCertStructs.DrmBCertDeviceInfo.encode(device_info).length + 8,
            attribute: device_info
        }

        const feature = {
            feature_count: 1,
            features: [4, 13]
        }
        const feature_attribute = {
            flags: 1,
            tag: 5,
            length: BCertStructs.DrmBCertFeatureInfo.encode(feature).length + 8,
            attribute: feature
        }

        const cert_key_sign = {
            type: 1,
            length: 512, // bits
            flags: 0,
            key: signing_key.privateBytes(),
            usages_count: 1,
            usages: [1]
        }
        const cert_key_encrypt = {
            type: 1,
            length: 512, // bits
            flags: 0,
            key: encryption_key.privateBytes(),
            usages_count: 1,
            usages: [2]
        }
        const key_info = {
            key_count: 2,
            cert_keys: [cert_key_sign, cert_key_encrypt]
        }
        const key_info_attribute = {
            flags: 1,
            tag: 6,
            length: BCertStructs.DrmBCertKeyInfo.encode(key_info).length + 8,
            attribute: key_info
        }

        const manufacturer_info = parent.get_certificate(0).get_attribute(7)

        const new_bcert_container = {
            signature: 'CERT',
            version: 1,
            total_length: 0,
            certificate_length: 0,
            attributes: [
                basic_info_attribute,
                device_info_attribute,
                feature_attribute,
                key_info_attribute,
                manufacturer_info
            ]
        }

        let payload = BCertStructs.BCert.encode(new_bcert_container)
        new_bcert_container.certificate_length = payload.length
        new_bcert_container.total_length = payload.length + 144
        payload = BCertStructs.BCert.encode(new_bcert_container)

        const hash = createHash('sha256')
        hash.update(payload)
        const digest = hash.digest()

        const signatureObj = group_key.keyPair.sign(digest)
        const r = Buffer.from(signatureObj.r.toArray('be', 32))
        const s = Buffer.from(signatureObj.s.toArray('be', 32))
        const signature = Buffer.concat([r, s])

        const signature_info = {
            signature_type: 1,
            signature_size: 64,
            signature: signature,
            signature_key_size: 512, // bits
            signature_key: group_key.publicBytes()
        }
        const signature_info_attribute = {
            flags: 1,
            tag: 8,
            length:
                BCertStructs.DrmBCertSignatureInfo.encode(signature_info)
                    .length + 8,
            attribute: signature_info
        }
        new_bcert_container.attributes.push(signature_info_attribute)

        return new Certificate(new_bcert_container)
    }

    static loads(data: string | Buffer): Certificate {
        if (typeof data === 'string') {
            data = Buffer.from(data, 'base64')
        }
        if (!Buffer.isBuffer(data)) {
            throw new Error(`Expecting Bytes or Base64 input, got ${data}`)
        }

        const cert = BCertStructs.BCert
        const parsed_bcert = cert.parse(data)
        return new Certificate(parsed_bcert, cert)
    }

    static load(filePath: string): Certificate {
        const data = fs.readFileSync(filePath)
        return Certificate.loads(data)
    }

    get_attribute(type_: number) {
        for (const attribute of this.parsed.attributes) {
            if (attribute.tag === type_) {
                return attribute
            }
        }
    }

    get_security_level(): number {
        const basic_info_attribute = this.get_attribute(1)
        if (basic_info_attribute) {
            return basic_info_attribute.attribute.security_level
        }
        return 0
    }

    private static _unpad(name: Buffer): string {
        return name.toString('utf8').replace(/\0+$/, '')
    }

    get_name(): string {
        const manufacturer_info_attribute = this.get_attribute(7)
        if (manufacturer_info_attribute) {
            const manufacturer_info = manufacturer_info_attribute.attribute
            const manufacturer_name = Certificate._unpad(
                manufacturer_info.manufacturer_name
            )
            const model_name = Certificate._unpad(manufacturer_info.model_name)
            const model_number = Certificate._unpad(
                manufacturer_info.model_number
            )
            return `${manufacturer_name} ${model_name} ${model_number}`
        }
        return ''
    }

    dumps(): Buffer {
        return this._BCERT.encode(this.parsed)
    }

    struct(): Parser {
        return this._BCERT
    }
}

export class CertificateChain {
    parsed: any
    _BCERT_CHAIN: Parser

    constructor(
        parsed_bcert_chain: any,
        bcert_chain_obj: Parser = BCertStructs.BCertChain
    ) {
        this.parsed = parsed_bcert_chain
        this._BCERT_CHAIN = bcert_chain_obj
    }

    static loads(data: string | Buffer): CertificateChain {
        if (typeof data === 'string') {
            data = Buffer.from(data, 'base64')
        }
        if (!Buffer.isBuffer(data)) {
            throw new Error(`Expecting Bytes or Base64 input, got ${data}`)
        }

        const cert_chain = BCertStructs.BCertChain
        try {
            const parsed_bcert_chain = cert_chain.parse(data)
            return new CertificateChain(parsed_bcert_chain, cert_chain)
        } catch (error) {
            console.error('Error during parsing:', error)
            throw error
        }
    }

    static load(filePath: string): CertificateChain {
        const data = fs.readFileSync(filePath)
        return CertificateChain.loads(data)
    }

    dumps(): Buffer {
        return this._BCERT_CHAIN.encode(this.parsed)
    }

    struct(): Parser {
        return this._BCERT_CHAIN
    }

    get_certificate(index: number): Certificate {
        return new Certificate(this.parsed.certificates[index])
    }

    get_security_level(): number {
        return this.get_certificate(0).get_security_level()
    }

    get_name(): string {
        return this.get_certificate(0).get_name()
    }

    append(bcert: Certificate): void {
        this.parsed.certificate_count += 1
        this.parsed.certificates.push(bcert.parsed)
        this.parsed.total_length += bcert.dumps().length
    }

    prepend(bcert: Certificate): void {
        this.parsed.certificate_count += 1
        this.parsed.certificates.unshift(bcert.parsed)
        this.parsed.total_length += bcert.dumps().length
    }
}