#!/usr/bin/env node

import chalk from "chalk"
import * as path from "path"
import { writeFileWithMkdirs } from "./file-utils.mjs"
import { getFqn, loadServicesFromProtoDescriptor } from "./proto-utils.mjs"

// Contains the interface definitions for the host bridge clients.
const TYPES_FILE = path.resolve("src/generated/hosts/host-bridge-client-types.ts")
// Contains the ExternalHostBridgeClientManager for the external host bridge clients (using nice-grpc).
const EXTERNAL_CLIENT_FILE = path.resolve("src/generated/hosts/standalone/host-bridge-clients.ts")
// Contains the handler map for the external host bridge clients (using the custom service registry).
const VSCODE_CLIENT_FILE = path.resolve("src/generated/hosts/vscode/hostbridge-grpc-service-config.ts")

/**
 * Main function to generate the host bridge client
 */
export async function main() {
	const { hostServices } = await loadServicesFromProtoDescriptor()

	await generateTypesFile(hostServices)
	await generateExternalClientFile(hostServices)
	await generateVscodeClientFile(hostServices)

	console.log(`Generated Host Bridge client files at:`)
	console.log(`- ${TYPES_FILE}`)
	console.log(`- ${EXTERNAL_CLIENT_FILE}`)
	console.log(`- ${VSCODE_CLIENT_FILE}`)
}

/**
 * Generate the client interfaces file.
 */
async function generateTypesFile(hostServices) {
	const clientInterfaces = []
	for (const [name, def] of Object.entries(hostServices)) {
		const clientInterface = generateClientInterfaceType(name, def)
		clientInterfaces.push(clientInterface)
	}
	const content = `// GENERATED CODE -- DO NOT EDIT!
// Generated by scripts/generate-host-bridge-client.mjs
import * as proto from "@shared/proto/index"
import { StreamingCallbacks } from "@hosts/host-provider-types"

${clientInterfaces.join("\n\n")}
`
	// Write output file
	await writeFileWithMkdirs(TYPES_FILE, content)
}

/**
 * Generate a client interface for a service.
 */
function generateClientInterfaceType(serviceName, serviceDefinition) {
	// Get the methods from the service definition
	const methods = Object.entries(serviceDefinition.service)
		.map(([methodName, methodDef]) => {
			const requestType = getFqn(methodDef.requestType.type.name)
			const responseType = getFqn(methodDef.responseType.type.name)

			if (!methodDef.responseStream) {
				// Generate unary method signature.
				return `	${methodName}(request: ${requestType}): Promise<${responseType}>;`
			}
			// Generate streaming method signature.
			return `	${methodName}(request: ${requestType}, callbacks: StreamingCallbacks<${responseType}>): () => void;`
		})
		.join("\n\n")

	// Generate the interface
	return `/**
 * Interface for ${serviceName} client.
 */
export interface ${serviceName}ClientInterface {

${methods}
}`
}

/**
 * Generate the external client implementations file.
 */
async function generateExternalClientFile(hostServices) {
	// Generate imports
	const imports = []
	// Add imports for the interfaces
	for (const [name, _def] of Object.entries(hostServices)) {
		imports.push(`import { ${name}ClientInterface } from "@generated/hosts/host-bridge-client-types"`)
	}
	const clientImplementations = []
	for (const [name, def] of Object.entries(hostServices)) {
		clientImplementations.push(generateExternalClientSetup(name, def))
	}

	const content = `// GENERATED CODE -- DO NOT EDIT!
// Generated by scripts/generate-host-bridge-client.mjs
import { asyncIteratorToCallbacks } from "@/standalone/utils"
import * as niceGrpc from "@generated/nice-grpc/index"
import { StreamingCallbacks } from "@hosts/host-provider-types"
import * as proto from "@shared/proto/index"
import { Channel, createClient } from "nice-grpc"
import { BaseGrpcClient } from "@/hosts/external/grpc-types"

${imports.join("\n")}

${clientImplementations.join("\n\n")}
`
	// Write output file
	await writeFileWithMkdirs(EXTERNAL_CLIENT_FILE, content)
}

/**
 * Generate a client implementation class for a service
 */
function generateExternalClientSetup(serviceName, serviceDefinition) {
	// Get the methods from the service definition
	const methods = Object.entries(serviceDefinition.service)
		.map(([methodName, methodDef]) => {
			// Get fully qualified type names
			const requestType = getFqn(methodDef.requestType.type.name)
			const responseType = getFqn(methodDef.responseType.type.name)
			const isStreamingResponse = methodDef.responseStream

			if (!isStreamingResponse) {
				return `    ${methodName}(request: ${requestType}): Promise<${responseType}> {
      return this.makeRequest((client) => client.${methodName}(request))
    }`
			} else {
				// Generate streaming method
				return `  ${methodName}(
		request: ${requestType},
		callbacks: StreamingCallbacks<${responseType}>,
	): () => void {
		const client = this.getClient()
		const abortController = new AbortController()
		const stream: AsyncIterable<${responseType}> = client.${methodName}(request, {
			signal: abortController.signal,
		})
		const wrappedCallbacks: StreamingCallbacks<${responseType}> = {
			...callbacks,
			onError: (error: any) => {
				if (error?.code === "UNAVAILABLE") {
					this.destroyClient()
				}
				callbacks.onError?.(error)
			},
		}
		asyncIteratorToCallbacks(stream, wrappedCallbacks)
		return () => {
			abortController.abort()
		}
	}\n`
			}
		})
		.join("\n")

	// Generate the class
	return `/**
 * Type-safe client implementation for ${serviceName}.
 */
export class ${serviceName}ClientImpl 
	extends BaseGrpcClient<niceGrpc.host.${serviceName}Client> 
	implements ${serviceName}ClientInterface {

	protected createClient(channel: Channel): niceGrpc.host.${serviceName}Client {
		return createClient(niceGrpc.host.${serviceName}Definition, channel)
	}

${methods}
}`
}

/**
 * Generate the Vscode client setup file.
 */
async function generateVscodeClientFile(hostServices) {
	const imports = []
	const clientImplementations = []
	const handlerMap = []
	for (const [serviceName, serviceDefinition] of Object.entries(hostServices)) {
		const name = serviceName.replace(/Service$/, "").toLowerCase()
		for (const [methodName, _methodDef] of Object.entries(serviceDefinition.service)) {
			imports.push(`import { ${methodName} } from "@/hosts/vscode/hostbridge/${name}/${methodName}"`)
		}
		imports.push("")

		clientImplementations.push(generateVscodeClientImplementation(name, serviceDefinition))

		handlerMap.push(`	"host.${serviceName}": {
		requestHandler: ${name}ServiceRegistry.handleRequest,
		streamingHandler: ${name}ServiceRegistry.handleStreamingRequest,
	},`)
	}

	const content = `// GENERATED CODE -- DO NOT EDIT!
// Generated by scripts/generate-host-bridge-client.mjs
import { createServiceRegistry } from "@hosts/vscode/hostbridge-grpc-service"
import { HostServiceHandlerConfig } from "@hosts/vscode/hostbridge-grpc-handler"

${imports.join("\n")}
${clientImplementations.join("\n\n")}

/**
 * Map of host service names to their handler configurations
 */
export const hostServiceHandlers: Record<string, HostServiceHandlerConfig> = {
${handlerMap.join("\n")}
}
`

	// Write output file
	await writeFileWithMkdirs(VSCODE_CLIENT_FILE, content)
}

function generateVscodeClientImplementation(serviceName, serviceDefinition) {
	// Get the methods from the service definition
	const name = serviceName.replace(/Service$/, "").toLowerCase()

	const methods = Object.entries(serviceDefinition.service)
		.map(([methodName, methodDef]) => {
			// Get fully qualified type names
			const isStreamingResponse = methodDef.responseStream
			if (!isStreamingResponse) {
				return `${name}ServiceRegistry.registerMethod("${methodName}", ${methodName})`
			} else {
				return `${name}ServiceRegistry.registerMethod("${methodName}", ${methodName}, { isStreaming: true })`
			}
		})
		.join("\n")

	// Generate the class
	return `// Setup ${name} service registry
const ${name}ServiceRegistry = createServiceRegistry("${name}")
${methods}`
}

// Only run main if this script is executed directly
if (import.meta.url === `file://${process.argv[1]}`) {
	main().catch((error) => {
		console.error(chalk.red("Error:"), error)
		process.exit(1)
	})
}