import {parseHeaders} from "letterparser";
import type {Maybe} from "./bones";
import {z} from "zod";
import dns from "node:dns/promises";
import crypto from "node:crypto";

type MailHeaders = {[key: string]: Maybe<string>}

const dkimPublicKey = z.object({
    v: z.literal('DKIM1').optional(),
    h: z.string().optional(),
    k: z.literal('rsa').optional(),
    n: z.string().optional(),
    p: z.string(),
    s: z.string().optional(),
    t: z.string().optional(),
})

type DKIMPublicKey = z.infer<typeof dkimPublicKey>

const dkimHeaderSchema = z.object({
    v: z.literal('1'),
    a: z.string(),
    c: z.string(),
    d: z.string(),
    s: z.string(),
    t: z.string(),
    h: z.string(),
    bh: z.string(),
    b: z.string(),
    canonicalizeHeaders: z.enum(['simple', 'relaxed']),
    canonicalizeBody: z.enum(['simple', 'relaxed']),
})

type DKIMHeader = z.infer<typeof dkimHeaderSchema>

function canonicalizeHeaders(headers: MailHeaders, mode: 'simple'|'relaxed'): MailHeaders {
    if ( mode === 'simple' ) {
        return headers
    }

    const newHeaders: MailHeaders = {}
    for ( const key of Object.keys(headers) ) {
        newHeaders[key.toLowerCase()] = headers[key]?.replace(/\s\s+/g, ' ').trimEnd()
    }
    return newHeaders
}

function hashDKIMHeaders(headers: MailHeaders, dkim: DKIMHeader): [string, string] {
    const canonicalizedHeaders = canonicalizeHeaders(headers, dkim.canonicalizeHeaders)

    let rawDKIMHeaderString = ''
    const seenHeaders: string[] = []
    for ( const header of dkim.h.split(':') ) {
        if ( seenHeaders.includes(header) ) continue
        rawDKIMHeaderString += `${header}:${canonicalizedHeaders[header]?.trim() || ''}\r\n`
        seenHeaders.push(header)
    }

    rawDKIMHeaderString += `dkim-signature:${canonicalizedHeaders['dkim-signature']!.trim().replace(/b=[\w0-9\s/+=]+/, 'b=')}\r\n`

    const hasher = new Bun.CryptoHasher('sha256')
    hasher.update(rawDKIMHeaderString)
    return [rawDKIMHeaderString, hasher.digest('base64')]
}

function parseDKIMHeader(raw: string): DKIMHeader {
    const val: Record<string, string> = {}

    for ( const part of raw.split(';') ) {
        const split = part.split('=')
        if ( split.length < 2 ) {
            throw new Error('Bad DKIM record entry: ' + part)
        }

        const key = split.shift()!
        val[key.trim()] = split.join('=').replace(/(\n|\t|\r|\s)/g, '')
    }

    val.canonicalizeHeaders = val.c?.split('/')?.[0]
    val.canonicalizeBody = val.c?.split('/')?.[1]

    return dkimHeaderSchema.parse(val)
}

function parseRawBody(raw: string): string {
    // Split the raw body by new line:
    const lines = raw.split('\n')
    for ( let i = 0; i < lines.length; i += 1 ) {
        if ( !lines[i].trim() ) {
            return lines.slice(i).join('\n').trimStart()
        }
    }

    return ''
}

function computeDKIMBodyHash(raw: string): string {
    const body = parseRawBody(raw)
    const hasher = new Bun.CryptoHasher('sha256')
    hasher.update(body)
    return hasher.digest('base64')
}

async function getDKIMPublicKey(domain: string, selector: string) {
    const query = `${selector}._domainkey.${domain}.`
    const r = await dns.resolveTxt(query)
    const str = r.map(i => i.join('')).join('')

    const val: Record<string, string> = {}
    for ( const part of str.split(';') ) {
        const split = part.split('=')
        if ( split.length < 2 ) {
            throw new Error('Bad DKIM record entry: ' + part)
        }

        const key = split.shift()!
        val[key.trim()] = split.join('=').replace(/(\n|\t|\r|\s)/g, '')
    }

    return dkimPublicKey.parse(val)
}

const file = Bun.file('sample.eml')
const arrBuf = await file.arrayBuffer()
const buf = Buffer.from(arrBuf)
const content = buf.toString('utf-8')

const bh = computeDKIMBodyHash(content)
const rawHeaders = parseHeaders(content)
const headers = canonicalizeHeaders(rawHeaders, 'relaxed')
const dkim = parseDKIMHeader(headers['dkim-signature'] || '')

const bhMatches = bh === dkim.bh

const pk = await getDKIMPublicKey(dkim.d, dkim.s)

const [dkimHeaders, dkimHeaderHash] = hashDKIMHeaders(rawHeaders, dkim)

const verifier = crypto.createVerify('RSA-SHA256')
verifier.update(dkimHeaderHash, 'base64')

const pemkey = `-----BEGIN PUBLIC KEY-----\n${pk.p}\n-----END PUBLIC KEY-----`
const result = verifier.verify(pemkey, dkim.b, 'base64')

console.log({
    bh,
    headers,
    dkim,
    bhMatches,
    pk,
    dkimHeaders,
    dkimHeaderHash,
    result,
})


type DKIMResult = {
    pass: boolean,
    bodyHashResult: boolean,
    rawMail: string,
    bodyHash: string,
    canonicalizedHeaders: MailHeaders,
    dkim: DKIMHeader,
    key: DKIMPublicKey,
    headerString: string,
    headerHash: string,
    headerHashResult: boolean,
}