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

const OUT_FILE = path.resolve("src/generated/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(cline.${name}Service, {`)
		for (const [rpcName, rpc] of Object.entries(def.service)) {
			imports.push(`import { ${rpcName} } from "@core/controller/${dir}/${rpcName}"`)
			const requestType = "cline." + rpc.requestType.type.name
			if (rpc.requestStream) {
				throw new Error("Request streaming is not supported")
			}
			if (rpc.responseStream) {
				handlerSetup.push(`        ${rpcName}: wrapStreamingResponse<${requestType},void>(${rpcName}, controller),`)
			} else {
				const responseType = "cline." + rpc.responseType.type.name
				handlerSetup.push(`         ${rpcName}: wrapper<${requestType},${responseType}>(${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 { cline } from "@generated/grpc-js"
import { Controller } from "@core/controller"
import { GrpcHandlerWrapper, GrpcStreamingResponseHandlerWrapper } from "@/standalone/grpc-types"

${imports}
export function addProtobusServices(
	server: grpc.Server,
	controller: Controller,
	wrapper: GrpcHandlerWrapper,
	wrapStreamingResponse: GrpcStreamingResponseHandlerWrapper,
): void {
${handlerSetup}
}
`
// Write output file
fs.mkdirSync(dirname(OUT_FILE), { recursive: true })
fs.writeFileSync(OUT_FILE, output)

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