cc/src/phases/ZodifyPhase.ts

import {Phase} from './Phase'
import {ExtolloCompileConfig} from '../types'
import * as rimraf from 'rimraf'
import * as fse from 'fs-extra'
import * as path from 'path'
import {Logger} from '../Logger'

async function* walk(dir: string): any {
    for await (const d of await fse.promises.opendir(dir)) {
        const entry = path.join(dir, d.name)
        if (d.isDirectory()) {
            yield* walk(entry)
        } else if (d.isFile()) {
            yield entry
        }
    }
}

export class ZodifyPhase extends Phase {
    private static zodId = 1

    constructor(
        config: ExtolloCompileConfig,
        tsconfig: any,
        protected readonly zodPath: string,
    ) {
        super(config, tsconfig)
    }

    public async run(): Promise<void> {
        // eslint-disable-next-line @typescript-eslint/no-var-requires
        const tsz = require('ts-to-zod')
        const dir = this.config.compileDir || 'exbuild'
        await new Promise<void>((res, rej) => {
            rimraf(path.resolve(dir, this.zodPath), e => e ? rej(e) : res())
        })

        await fse.mkdirp(path.resolve(dir, this.zodPath))

        for await ( const file of walk(this.zodPath) ) {
            if ( !file.endsWith('.ts') ) {
                continue
            }

            Logger.verb('zod', file)

            const sourceText = (await fse.readFile(file)).toString('utf-8')
            const gen = tsz.generate({
                sourceText,
                getSchemaName: () => 'exZodifiedSchema',
            })

            const zodSource = gen.getZodSchemasFile(file)
            const augmentedSource = this.getAugmentedSourceText(sourceText, zodSource)

            await fse.writeFile(path.resolve(dir, file), augmentedSource)
        }
    }

    protected getAugmentedSourceText(originalSource: string, zodSource: string): string {
        const lines = originalSource.split('\n')
        let line = 0
        if ( lines[line].startsWith('#!') ) {
            if ( lines[line + 1] ) {
                line = line + 1
            }
        }

        const newlines: string[] = []
        lines.forEach((lineVal, idx) => {
            if ( idx === line ) {
                newlines.push('import { z } from "zod";')
                newlines.push('import { registerZodifiedSchema } from "@extollo/lib";')
            }

            newlines.push(lineVal)
        })

        const id = ZodifyPhase.zodId++
        const zodLines = [
            `/** @ex-zod-id ${id}@ */`,
            ...zodSource.split('\n')
                .filter(x => x.trim())
                .slice(2),
            `registerZodifiedSchema(${id}, exZodifiedSchema);`,
        ]

        return [...newlines, ...zodLines].join('\n')
    }
}