import * as grpc from "@grpc/grpc-js";
import * as protoLoader from "@grpc/proto-loader";
import { fileURLToPath } from "node:url";
import { StringDecoder } from "node:string_decoder";

import type {
  CallOptions,
  ConnectionConfig,
  Environment,
  ProcessOutput,
  ProcessRequest,
  RemoteProcess,
  SubstrateTransport,
} from "./types.js";

const WRITE_CHUNK_BYTES = 64 * 1024;

type RpcMessage = Record<string, unknown>;
type UnaryCallback = (error: grpc.ServiceError | null, response?: RpcMessage) => void;
type UnaryMethod = (
  request: RpcMessage,
  metadata: grpc.Metadata,
  options: grpc.CallOptions,
  callback: UnaryCallback,
) => grpc.ClientUnaryCall;
type ReadableMethod = (
  request: RpcMessage,
  metadata: grpc.Metadata,
  options: grpc.CallOptions,
) => grpc.ClientReadableStream<RpcMessage>;
type WritableMethod = (
  metadata: grpc.Metadata,
  options: grpc.CallOptions,
  callback: UnaryCallback,
) => grpc.ClientWritableStream<RpcMessage>;

interface EnvironmentServiceClient {
  createEnvironment: UnaryMethod;
  getEnvironment: UnaryMethod;
  suspendEnvironment: UnaryMethod;
  deleteEnvironment: UnaryMethod;
  close(): void;
}

interface ProcessServiceClient {
  startProcess: UnaryMethod;
  getProcess: UnaryMethod;
  streamProcessOutputs: ReadableMethod;
  killProcess: UnaryMethod;
  close(): void;
}

interface FileSystemServiceClient {
  readFile: ReadableMethod;
  writeFile: WritableMethod;
  close(): void;
}

interface RpcContext {
  metadata: grpc.Metadata;
  options: grpc.CallOptions;
}

const protoDirectory = fileURLToPath(new URL("../proto/", import.meta.url));
const environmentProto = fileURLToPath(new URL("../proto/env.proto", import.meta.url));
const guestProto = fileURLToPath(new URL("../proto/guest.proto", import.meta.url));

const packageDefinition = protoLoader.loadSync([environmentProto, guestProto], {
  keepCase: false,
  longs: Number,
  enums: String,
  defaults: true,
  oneofs: true,
  includeDirs: [protoDirectory],
});

const loadedProto = grpc.loadPackageDefinition(packageDefinition);

function serviceConstructor(name: string): grpc.ServiceClientConstructor {
  const ateenv = (loadedProto as unknown as Record<string, unknown>).ateenv;
  const v1alpha = ateenv && typeof ateenv === "object"
    ? (ateenv as Record<string, unknown>).v1alpha
    : undefined;
  const candidate = v1alpha && typeof v1alpha === "object"
    ? (v1alpha as Record<string, unknown>)[name]
    : undefined;

  if (typeof candidate !== "function") {
    throw new Error(`Agent Substrate protocol does not expose ${name}`);
  }

  return candidate as grpc.ServiceClientConstructor;
}

function createService<T>(
  constructor: grpc.ServiceClientConstructor,
  address: string,
  credentials: grpc.ChannelCredentials,
): T {
  return new constructor(address, credentials) as unknown as T;
}

function abortError(method: string): Error {
  const error = new Error(`${method} was aborted`);
  error.name = "AbortError";
  return error;
}

function grpcError(method: string, error: grpc.ServiceError): Error {
  const detail = typeof error.code === "number" ? grpc.status[error.code] : undefined;
  return new Error(`${method} failed${detail ? ` (${detail})` : ""}`);
}

function asMessage(value: unknown, context: string): RpcMessage {
  if (!value || typeof value !== "object") {
    throw new Error(`${context} returned an invalid response`);
  }
  return value as RpcMessage;
}

function requiredString(value: unknown, field: string): string {
  if (typeof value !== "string" || value.length === 0) {
    throw new Error(`Invalid ${field} in remote response`);
  }
  return value;
}

function integer(value: unknown, field: string): number {
  if (typeof value !== "number" || !Number.isSafeInteger(value)) {
    throw new Error(`Invalid ${field} in remote response`);
  }
  return value;
}

function binary(value: unknown, field: string): Buffer {
  if (Buffer.isBuffer(value)) {
    return value;
  }
  if (value instanceof Uint8Array) {
    return Buffer.from(value);
  }
  throw new Error(`Invalid ${field} in remote response`);
}

function nonNegativeOffset(value: number | undefined, name: string): number {
  if (value === undefined) {
    return 0;
  }
  if (!Number.isSafeInteger(value) || value < 0) {
    throw new Error(`${name} must be a non-negative safe integer`);
  }
  return value;
}

function validateEndpoint(endpoint: string): { address: string; credentials: grpc.ChannelCredentials } {
  if (endpoint !== endpoint.trim()) {
    throw new Error("Endpoint must not contain leading or trailing whitespace");
  }

  let url: URL;
  try {
    url = new URL(endpoint);
  } catch {
    throw new Error("Endpoint must be an absolute http or https URL");
  }

  if (url.protocol !== "http:" && url.protocol !== "https:") {
    throw new Error("Endpoint must use http or https");
  }
  if (url.username || url.password) {
    throw new Error("Endpoint must not include user credentials");
  }
  if (url.pathname !== "/" || endpoint.includes("?") || endpoint.includes("#")) {
    throw new Error("Endpoint must not include a path, query, or fragment");
  }

  const hostname = url.hostname.toLowerCase();
  if (url.protocol === "http:" && !isLoopbackHost(hostname)) {
    throw new Error("Insecure http is allowed only for loopback endpoints");
  }

  const host = url.hostname.includes(":") && !url.hostname.startsWith("[")
    ? `[${url.hostname}]`
    : url.hostname;
  const port = url.port || (url.protocol === "https:" ? "443" : "80");
  return {
    address: `${host}:${port}`,
    credentials: url.protocol === "https:" ? grpc.credentials.createSsl() : grpc.credentials.createInsecure(),
  };
}

function isLoopbackHost(hostname: string): boolean {
  if (hostname === "localhost" || hostname === "::1" || hostname === "[::1]") {
    return true;
  }

  const octets = hostname.split(".");
  return octets.length === 4
    && octets.every((octet) => /^\d+$/.test(octet) && Number(octet) >= 0 && Number(octet) <= 255)
    && Number(octets[0]) === 127;
}

/** gRPC transport for Agent Substrate environment and guest services. */
export class SubstrateClient implements SubstrateTransport {
  private readonly environment: EnvironmentServiceClient;
  private readonly processService: ProcessServiceClient;
  private readonly fileSystem: FileSystemServiceClient;
  private closed = false;

  constructor(private readonly config: ConnectionConfig) {
    if (config.token !== undefined && (typeof config.token !== "string" || /[^\x20-\x7e]/.test(config.token))) {
      throw new Error("SUBSTRATE_TOKEN contains invalid gRPC metadata characters.");
    }
    if (!config.atespace) {
      throw new Error("atespace is required");
    }
    if (!Number.isSafeInteger(config.timeoutMs) || config.timeoutMs <= 0) {
      throw new Error("timeoutMs must be a positive safe integer");
    }
    if (!Number.isSafeInteger(config.maxBytes) || config.maxBytes < 0) {
      throw new Error("maxBytes must be a non-negative safe integer");
    }

    const { address, credentials } = validateEndpoint(config.endpoint);
    this.environment = createService<EnvironmentServiceClient>(
      serviceConstructor("EnvironmentService"),
      address,
      credentials,
    );
    this.processService = createService<ProcessServiceClient>(
      serviceConstructor("ProcessService"),
      address,
      credentials,
    );
    this.fileSystem = createService<FileSystemServiceClient>(
      serviceConstructor("FileSystemService"),
      address,
      credentials,
    );
  }

  async create(
    id: string,
    template?: { name: string; atespace: string },
    options?: CallOptions,
  ): Promise<Environment> {
    const response = await this.unary(
      "EnvironmentService.CreateEnvironment",
      this.environment,
      this.environment.createEnvironment,
      { id, atespace: this.config.atespace, ...(template ? { template } : {}) },
      options,
    );
    return this.environmentFrom(response.environment, id);
  }

  async get(id: string, options?: CallOptions): Promise<Environment> {
    const response = await this.unary(
      "EnvironmentService.GetEnvironment",
      this.environment,
      this.environment.getEnvironment,
      { id, atespace: this.config.atespace },
      options,
    );
    return this.environmentFrom(response.environment, id);
  }

  async suspend(id: string, options?: CallOptions): Promise<void> {
    await this.unary(
      "EnvironmentService.SuspendEnvironment",
      this.environment,
      this.environment.suspendEnvironment,
      { id, atespace: this.config.atespace },
      options,
    );
  }

  async delete(id: string, options?: CallOptions): Promise<void> {
    await this.unary(
      "EnvironmentService.DeleteEnvironment",
      this.environment,
      this.environment.deleteEnvironment,
      { id, atespace: this.config.atespace },
      options,
    );
  }

  async start(id: string, request: ProcessRequest, options?: CallOptions): Promise<string> {
    const response = await this.unary(
      "ProcessService.StartProcess",
      this.processService,
      this.processService.startProcess,
      { command: request.command, ...(request.cwd !== undefined ? { cwd: request.cwd } : {}), ...(request.env ? { env: request.env } : {}) },
      options,
      id,
    );
    return requiredString(response.processId, "process_id");
  }

  async process(id: string, processId: string, options?: CallOptions): Promise<RemoteProcess> {
    const response = await this.unary(
      "ProcessService.GetProcess",
      this.processService,
      this.processService.getProcess,
      { processId },
      options,
      id,
    );
    const returnedId = requiredString(response.processId, "process_id");
    if (returnedId !== processId) {
      throw new Error("ProcessService.GetProcess returned a different process_id");
    }
    return {
      processId: returnedId,
      status: requiredString(response.status, "process status"),
      exitCode: integer(response.exitCode, "process exit_code"),
    };
  }

  async kill(id: string, processId: string, options?: CallOptions): Promise<number> {
    const response = await this.unary(
      "ProcessService.KillProcess",
      this.processService,
      this.processService.killProcess,
      { processId },
      options,
      id,
    );
    return integer(response.exitCode, "kill exit_code");
  }

  outputs(
    id: string,
    processId: string,
    request?: { follow?: boolean; stdoutOffset?: number; stderrOffset?: number },
    options?: CallOptions,
  ): Promise<ProcessOutput> {
    this.assertOpen();
    if (options?.signal?.aborted) {
      return Promise.reject(abortError("ProcessService.StreamProcessOutputs"));
    }

    const stdoutOffset = nonNegativeOffset(request?.stdoutOffset, "stdoutOffset");
    const stderrOffset = nonNegativeOffset(request?.stderrOffset, "stderrOffset");
    const context = this.rpcContext(options, id);
    let stream: grpc.ClientReadableStream<RpcMessage>;
    try {
      stream = this.processService.streamProcessOutputs(
        { processId, stdoutOffset, stderrOffset, follow: request?.follow ?? false },
        context.metadata,
        context.options,
      );
    } catch {
      return Promise.reject(new Error("ProcessService.StreamProcessOutputs failed to start"));
    }

    return new Promise<ProcessOutput>((resolve, reject) => {
      let settled = false;
      let stdoutBytes = 0;
      let stderrBytes = 0;
      let totalBytes = 0;
      let truncated = false;
      const stdout = new StringDecoder("utf8");
      const stderr = new StringDecoder("utf8");
      let stdoutText = "";
      let stderrText = "";

      const cleanup = () => {
        stream.removeListener("data", onData);
        stream.removeListener("end", onEnd);
        options?.signal?.removeEventListener("abort", onAbort);
      };
      const finish = (error?: Error) => {
        if (settled) {
          return;
        }
        settled = true;
        cleanup();
        if (error) {
          reject(error);
          return;
        }
        stdoutText += stdout.end();
        stderrText += stderr.end();
        resolve({
          stdout: stdoutText,
          stderr: stderrText,
          stdoutOffset: stdoutOffset + stdoutBytes,
          stderrOffset: stderrOffset + stderrBytes,
          truncated,
        });
      };
      const onAbort = () => {
        stream.cancel();
        finish(abortError("ProcessService.StreamProcessOutputs"));
      };
      const onError = (error: grpc.ServiceError) => finish(grpcError("ProcessService.StreamProcessOutputs", error));
      const onEnd = () => finish();
      const onData = (chunk: RpcMessage) => {
        if (settled) {
          return;
        }
        try {
          const source = requiredString(chunk.source, "output source");
          const data = binary(chunk.data, "output data");
          const remaining = this.config.maxBytes - totalBytes;
          const accepted = data.subarray(0, Math.max(remaining, 0));
          if (source === "OUTPUT_SOURCE_STDOUT") {
            stdoutText += stdout.write(accepted);
            stdoutBytes += accepted.length;
          } else if (source === "OUTPUT_SOURCE_STDERR") {
            stderrText += stderr.write(accepted);
            stderrBytes += accepted.length;
          } else {
            throw new Error("Invalid output source in remote response");
          }
          totalBytes += accepted.length;
          if (accepted.length !== data.length) {
            truncated = true;
            stream.cancel();
            finish();
          }
        } catch (error) {
          stream.cancel();
          finish(error instanceof Error ? error : new Error("Invalid process output response"));
        }
      };

      stream.on("data", onData);
      stream.once("error", onError);
      stream.once("end", onEnd);
      options?.signal?.addEventListener("abort", onAbort, { once: true });
      if (options?.signal?.aborted) {
        onAbort();
      }
    });
  }

  read(id: string, path: string, options?: CallOptions): Promise<Buffer> {
    this.assertOpen();
    if (options?.signal?.aborted) {
      return Promise.reject(abortError("FileSystemService.ReadFile"));
    }

    const context = this.rpcContext(options, id);
    let stream: grpc.ClientReadableStream<RpcMessage>;
    try {
      stream = this.fileSystem.readFile({ path }, context.metadata, context.options);
    } catch {
      return Promise.reject(new Error("FileSystemService.ReadFile failed to start"));
    }
    return new Promise<Buffer>((resolve, reject) => {
      let settled = false;
      let totalBytes = 0;
      const chunks: Buffer[] = [];

      const cleanup = () => {
        stream.removeListener("data", onData);
        stream.removeListener("end", onEnd);
        options?.signal?.removeEventListener("abort", onAbort);
      };
      const finish = (error?: Error) => {
        if (settled) {
          return;
        }
        settled = true;
        cleanup();
        if (error) {
          reject(error);
          return;
        }
        resolve(Buffer.concat(chunks, totalBytes));
      };
      const onAbort = () => {
        stream.cancel();
        finish(abortError("FileSystemService.ReadFile"));
      };
      const onError = (error: grpc.ServiceError) => finish(grpcError("FileSystemService.ReadFile", error));
      const onEnd = () => finish();
      const onData = (chunk: RpcMessage) => {
        if (settled) {
          return;
        }
        try {
          const data = binary(chunk.data, "file data");
          if (data.length > this.config.maxBytes - totalBytes) {
            stream.cancel();
            finish(new Error(`Read requested at least ${totalBytes + data.length} bytes, exceeding the ${this.config.maxBytes}-byte limit. Read a bounded range through the remote shell.`));
            return;
          }
          chunks.push(data);
          totalBytes += data.length;
        } catch (error) {
          stream.cancel();
          finish(error instanceof Error ? error : new Error("Invalid file read response"));
        }
      };

      stream.on("data", onData);
      stream.once("error", onError);
      stream.once("end", onEnd);
      options?.signal?.addEventListener("abort", onAbort, { once: true });
      if (options?.signal?.aborted) {
        onAbort();
      }
    });
  }

  async write(id: string, path: string, content: Buffer, options?: CallOptions): Promise<number> {
    this.assertOpen();
    if (content.length > this.config.maxBytes) {
      throw new Error(`Write requested ${content.length} bytes, exceeding the ${this.config.maxBytes}-byte limit. Split the file into bounded remote operations.`);
    }
    if (options?.signal?.aborted) {
      throw abortError("FileSystemService.WriteFile");
    }

    const context = this.rpcContext(options, id);
    return new Promise<number>((resolve, reject) => {
      let settled = false;
      let stream: grpc.ClientWritableStream<RpcMessage> | undefined;

      const cleanup = () => {
        options?.signal?.removeEventListener("abort", onAbort);
      };
      const finish = (error?: Error, bytesWritten?: number) => {
        if (settled) {
          return;
        }
        settled = true;
        cleanup();
        if (error) {
          reject(error);
        } else {
          resolve(bytesWritten as number);
        }
      };
      const onAbort = () => {
        stream?.cancel();
        finish(abortError("FileSystemService.WriteFile"));
      };
      const onError = (error: grpc.ServiceError) => finish(grpcError("FileSystemService.WriteFile", error));

      try {
        stream = this.fileSystem.writeFile(context.metadata, context.options, (error, response) => {
          if (error) {
            finish(grpcError("FileSystemService.WriteFile", error));
            return;
          }
          try {
            const bytesWritten = integer(asMessage(response, "FileSystemService.WriteFile").bytesWritten, "bytes_written");
            if (bytesWritten !== content.length) {
              throw new Error("FileSystemService.WriteFile returned an unexpected bytes_written");
            }
            finish(undefined, bytesWritten);
          } catch (responseError) {
            finish(responseError instanceof Error ? responseError : new Error("Invalid file write response"));
          }
        });
        stream.once("error", onError);
        options?.signal?.addEventListener("abort", onAbort, { once: true });
        if (options?.signal?.aborted) {
          onAbort();
          return;
        }
        void this.writeChunks(stream, path, content, options?.signal).catch((error: unknown) => {
          stream?.cancel();
          finish(error instanceof Error ? error : new Error("FileSystemService.WriteFile failed"));
        });
      } catch (error) {
        finish(error instanceof Error ? error : new Error("FileSystemService.WriteFile failed"));
      }
    });
  }

  close(): void {
    if (this.closed) {
      return;
    }
    this.closed = true;
    this.environment.close();
    this.processService.close();
    this.fileSystem.close();
  }

  private async writeChunks(
    stream: grpc.ClientWritableStream<RpcMessage>,
    path: string,
    content: Buffer,
    signal: AbortSignal | undefined,
  ): Promise<void> {
    if (content.length === 0) {
      await this.writeChunk(stream, { path, chunk: Buffer.alloc(0) }, signal);
      stream.end();
      return;
    }

    for (let offset = 0; offset < content.length; offset += WRITE_CHUNK_BYTES) {
      await this.writeChunk(
        stream,
        {
          ...(offset === 0 ? { path } : {}),
          chunk: content.subarray(offset, Math.min(offset + WRITE_CHUNK_BYTES, content.length)),
        },
        signal,
      );
    }
    stream.end();
  }

  private writeChunk(
    stream: grpc.ClientWritableStream<RpcMessage>,
    message: RpcMessage,
    signal: AbortSignal | undefined,
  ): Promise<void> {
    if (signal?.aborted) {
      return Promise.reject(abortError("FileSystemService.WriteFile"));
    }
    if (stream.write(message)) {
      return Promise.resolve();
    }

    return new Promise<void>((resolve, reject) => {
      let settled = false;
      const cleanup = () => {
        stream.removeListener("drain", onDrain);
        stream.removeListener("error", onError);
        signal?.removeEventListener("abort", onAbort);
      };
      const finish = (error?: Error) => {
        if (settled) {
          return;
        }
        settled = true;
        cleanup();
        if (error) {
          reject(error);
        } else {
          resolve();
        }
      };
      const onDrain = () => finish();
      const onError = (error: grpc.ServiceError) => finish(grpcError("FileSystemService.WriteFile", error));
      const onAbort = () => finish(abortError("FileSystemService.WriteFile"));

      stream.once("drain", onDrain);
      stream.once("error", onError);
      signal?.addEventListener("abort", onAbort, { once: true });
      if (signal?.aborted) {
        onAbort();
      }
    });
  }

  private unary(
    name: string,
    client: object,
    method: UnaryMethod,
    request: RpcMessage,
    callOptions: CallOptions | undefined,
    environmentId?: string,
  ): Promise<RpcMessage> {
    this.assertOpen();
    if (callOptions?.signal?.aborted) {
      return Promise.reject(abortError(name));
    }

    const context = this.rpcContext(callOptions, environmentId);
    return new Promise<RpcMessage>((resolve, reject) => {
      let settled = false;
      let call: grpc.ClientUnaryCall | undefined;
      const cleanup = () => callOptions?.signal?.removeEventListener("abort", onAbort);
      const finish = (error?: Error, response?: RpcMessage) => {
        if (settled) {
          return;
        }
        settled = true;
        cleanup();
        if (error) {
          reject(error);
        } else {
          resolve(response as RpcMessage);
        }
      };
      const onAbort = () => {
        call?.cancel();
        finish(abortError(name));
      };

      try {
        call = method.call(client, request, context.metadata, context.options, (error, response) => {
          if (error) {
            finish(grpcError(name, error));
            return;
          }
          try {
            finish(undefined, asMessage(response, name));
          } catch (responseError) {
            finish(responseError instanceof Error ? responseError : new Error(`${name} returned an invalid response`));
          }
        });
        if (!settled) {
          callOptions?.signal?.addEventListener("abort", onAbort, { once: true });
          if (callOptions?.signal?.aborted) {
            onAbort();
          }
        }
      } catch (error) {
        finish(error instanceof Error ? error : new Error(`${name} failed`));
      }
    });
  }

  private rpcContext(options: CallOptions | undefined, environmentId?: string): RpcContext {
    const timeoutMs = options?.timeoutMs ?? this.config.timeoutMs;
    if (!Number.isSafeInteger(timeoutMs) || timeoutMs <= 0) {
      throw new Error("timeoutMs must be a positive safe integer");
    }

    const metadata = new grpc.Metadata();
    if (this.config.token) {
      metadata.set("authorization", `Bearer ${this.config.token}`);
    }
    if (environmentId !== undefined) {
      metadata.set("x-env-id", environmentId);
      metadata.set("x-env-atespace", this.config.atespace);
    }
    return { metadata, options: { deadline: new Date(Date.now() + timeoutMs) } };
  }

  private environmentFrom(value: unknown, expectedId: string): Environment {
    const environment = asMessage(value, "EnvironmentService");
    const id = requiredString(environment.id, "environment id");
    const atespace = requiredString(environment.atespace, "environment atespace");
    if (id !== expectedId || atespace !== this.config.atespace) {
      throw new Error("EnvironmentService returned an environment with unexpected identity");
    }
    return { id, atespace, status: requiredString(environment.status, "environment status") };
  }

  private assertOpen(): void {
    if (this.closed) {
      throw new Error("SubstrateClient is closed");
    }
  }
}