import * as fs from "fs"
import * as grpc from "@grpc/grpc-js"
import * as protoLoader from "@grpc/proto-loader"
import * as health from "grpc-health-check"
import { fileURLToPath } from "url"
import path from "path"

const OUT_FILE = path.resolve("src/standalone/server-setup.ts")
const DESCRIPTOR_SET = path.resolve("dist-standalone/proto/descriptor_set.pb")

// Load service definitions.
const clineDef = protoLoader.loadFileDescriptorSetFromBuffer(fs.readFileSync(DESCRIPTOR_SET))
const healthDef = protoLoader.loadSync(health.protoPath)
const packageDefinition = { ...clineDef, ...healthDef }
const proto = grpc.loadPackageDefinition(packageDefinition)

/**
 * Generate imports and function to add all the handlers to the server for all services defined in the proto files.
 */
function generateHandlersAndExports() {
	let imports = []
	let handlerSetup = []

	for (const [name, def] of Object.entries(proto.cline)) {
		if (!def || !("service" in def)) {
			continue
		}
		const domain = name.replace(/Service$/, "")
		const dir = domain.charAt(0).toLowerCase() + domain.slice(1)
		imports.push(`// ${domain} Service`)
		handlerSetup.push(`    // ${domain} Service`)
		handlerSetup.push(`    server.addService(proto.cline.${name}.service, {`)
		for (const [rpcName, rpc] of Object.entries(def.service)) {
			imports.push(`import { ${rpcName} } from "../core/controller/${dir}/${rpcName}"`)
			if (rpc.requestStream) {
				throw new Error("Request streaming is not supported")
			}
			if (rpc.responseStream) {
				handlerSetup.push(`        ${rpcName}: wrapStreamingResponse(${rpcName}, controller),`)
			} else {
				handlerSetup.push(`         ${rpcName}: wrapper(${rpcName}, controller),`)
			}
		}
		handlerSetup.push(`    });`)
		imports.push("")
		handlerSetup.push("")
	}
	return {
		imports: imports.join("\n"),
		handlerSetup: handlerSetup.join("\n"),
	}
}

const { imports, handlerSetup } = generateHandlersAndExports()
const scriptName = path.basename(fileURLToPath(import.meta.url))

// Create output file
let output = `// GENERATED CODE -- DO NOT EDIT!
// Generated by ${scriptName}
import * as grpc from "@grpc/grpc-js"
import { Controller } from "../core/controller"
import { GrpcHandlerWrapper, GrpcStreamingResponseHandlerWrapper } from "./grpc-types"

${imports}

export function addServices(
	server: grpc.Server,
	proto: any,
	controller: Controller,
	wrapper: GrpcHandlerWrapper,
	wrapStreamingResponse: GrpcStreamingResponseHandlerWrapper,
): void {
${handlerSetup}
}
`
// Write output file
fs.writeFileSync(OUT_FILE, output)

console.log(`Generated service handlers in ${OUT_FILE}.`)