/**
 * Lakebase client implementation.
 *
 * Uses the native Lakebase plugin type from @superblocksteam/types.
 * Supports server-side parameterized SQL queries via the `parameters` field.
 */

import type { PartialMessage } from "@bufbuild/protobuf";
import type { z } from "zod";

import type { Plugin as LakebasePlugin } from "@superblocksteam/types/dist/src/plugins/lakebase/v1/plugin_pb";

import { QueryValidationError } from "../../errors.js";
import { IntegrationError } from "../../runtime/errors.js";
import type { QueryExecutor, TraceMetadata } from "../registry.js";
import type { IntegrationConfig, IntegrationClientImpl } from "../types.js";
import { describeType } from "../utils.js";
import type { LakebaseClient } from "./types.js";

/**
 * Lakebase request type derived from proto definition.
 */
type LakebaseRequest = PartialMessage<LakebasePlugin>;

/**
 * Internal implementation of LakebaseClient.
 *
 * Communicates with the orchestrator to execute SQL queries.
 * Uses server-side parameterization via the `parameters` field.
 */
export class LakebaseClientImpl
  implements LakebaseClient, IntegrationClientImpl
{
  readonly name: string;
  readonly pluginId: string;
  readonly config: IntegrationConfig;

  private readonly executeQuery: QueryExecutor;

  constructor(config: IntegrationConfig, executeQuery: QueryExecutor) {
    this.name = config.name;
    this.pluginId = config.pluginId;
    this.config = config;
    this.executeQuery = executeQuery;
  }

  /**
   * Builds a Lakebase plugin request object.
   *
   * Parameters are passed via the `parameters` field as a JSON-stringified
   * array. The server performs parameter binding, avoiding client-side
   * SQL interpolation.
   */
  private buildRequest(sql: string, params?: unknown[]): LakebaseRequest {
    const hasParams = params && params.length > 0;

    return {
      body: sql,
      parameters: hasParams ? JSON.stringify(params) : undefined,
    };
  }

  async query<T>(
    sql: string,
    schema: z.ZodSchema<T>,
    params?: unknown[],
    metadata?: TraceMetadata,
  ): Promise<T[]> {
    const request = this.buildRequest(sql, params);

    try {
      const result = await this.executeQuery(
        request as Record<string, unknown>,
        undefined,
        metadata,
      );

      if (!Array.isArray(result)) {
        throw new IntegrationError(
          this.config.name,
          "query",
          `Expected array result from Lakebase query, got: ${describeType(result)}`,
        );
      }

      const validated: T[] = [];
      for (let i = 0; i < result.length; i++) {
        const row = result[i];
        const parseResult = schema.safeParse(row);

        if (!parseResult.success) {
          throw new QueryValidationError(
            `Row ${i} failed validation: ${parseResult.error.message}`,
            {
              rowIndex: i,
              errors: parseResult.error.errors,
              row,
            },
          );
        }

        validated.push(parseResult.data);
      }

      return validated;
    } catch (error) {
      if (
        error instanceof QueryValidationError ||
        error instanceof IntegrationError
      ) {
        throw error;
      }

      throw new IntegrationError(this.config.name, "query", error);
    }
  }

  async execute(
    sql: string,
    params?: unknown[],
    metadata?: TraceMetadata,
  ): Promise<{ rowCount: number }> {
    const request = this.buildRequest(sql, params);

    try {
      const result = await this.executeQuery(
        request as Record<string, unknown>,
        undefined,
        metadata,
      );

      if (
        typeof result === "object" &&
        result !== null &&
        "rowCount" in result
      ) {
        return { rowCount: Number(result.rowCount) || 0 };
      }

      if (Array.isArray(result)) {
        return { rowCount: result.length };
      }

      return { rowCount: 0 };
    } catch (error) {
      throw new IntegrationError(this.config.name, "execute", error);
    }
  }
}