/**
 * CosmosDB client implementation.
 *
 * Uses the native CosmosDB plugin type from @superblocksteam/types.
 * Supports SQL queries and point operations (create, read, replace,
 * upsert, delete) against CosmosDB containers.
 */

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

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

import { RestApiValidationError } from "../../errors.js";
import { IntegrationError } from "../../runtime/errors.js";
import type { QueryExecutor, TraceMetadata } from "../registry.js";
import type { IntegrationConfig, IntegrationClientImpl } from "../types.js";
import type { CosmosDBClient, CosmosDBQueryOptions } from "./types.js";

/**
 * CosmosDB request type derived from proto definition.
 */
type CosmosDBRequest = PartialMessage<CosmosDBPlugin>;

/**
 * Internal implementation of CosmosDBClient.
 *
 * Communicates with the orchestrator using the native CosmosDB plugin
 * proto format, supporting SQL queries and point operations.
 */
export class CosmosDBClientImpl
  implements CosmosDBClient, IntegrationClientImpl
{
  readonly config: IntegrationConfig;
  private readonly executeQuery: QueryExecutor;

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

  get name(): string {
    return this.config.name;
  }

  get pluginId(): string {
    return this.config.pluginId;
  }

  /**
   * Execute a request and wrap errors.
   */
  private async exec(
    request: CosmosDBRequest,
    operation: string,
    metadata?: TraceMetadata,
  ): Promise<unknown> {
    try {
      return await this.executeQuery(
        request as Record<string, unknown>,
        undefined,
        metadata,
      );
    } catch (error) {
      if (
        error instanceof RestApiValidationError ||
        error instanceof IntegrationError
      ) {
        throw error;
      }

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

  /**
   * Validate a result against a schema.
   */
  private validate<T>(result: unknown, schema: z.ZodSchema<T>): T {
    const parseResult = schema.safeParse(result);

    if (!parseResult.success) {
      throw new RestApiValidationError(
        `Result validation failed: ${parseResult.error.message}`,
        {
          zodError: parseResult.error,
          data: result,
        },
      );
    }

    return parseResult.data;
  }

  async query<T>(
    containerId: string,
    sql: string,
    schema: z.ZodSchema<T>,
    options?: CosmosDBQueryOptions,
    metadata?: TraceMetadata,
  ): Promise<T> {
    // Proto JSON format for: oneof cosmosdb_action { Sql sql = 5; }
    // Sql.oneof action { Singleton singleton = 1; }
    const request = {
      sql: {
        singleton: {
          containerId,
          query: sql,
          crossPartition: options?.crossPartition ?? false,
          partitionKey: options?.partitionKey,
        },
      },
    };

    const result = await this.exec(
      request as unknown as Record<string, unknown>,
      "query",
      metadata,
    );
    return this.validate(result, schema);
  }

  async read<T>(
    containerId: string,
    id: string,
    schema: z.ZodSchema<T>,
    partitionKey?: string,
    metadata?: TraceMetadata,
  ): Promise<T> {
    // Proto JSON format for: oneof cosmosdb_action { PointOperation point_operation = 6; }
    const request = {
      pointOperation: {
        containerId,
        read: { id, partitionKey },
      },
    };

    const result = await this.exec(
      request as unknown as Record<string, unknown>,
      "read",
      metadata,
    );
    return this.validate(result, schema);
  }

  async create(
    containerId: string,
    body: unknown,
    partitionKey?: string,
    metadata?: TraceMetadata,
  ): Promise<unknown> {
    const request = {
      pointOperation: {
        containerId,
        create: {
          body: typeof body === "string" ? body : JSON.stringify(body),
          partitionKey,
        },
      },
    };

    return this.exec(
      request as unknown as Record<string, unknown>,
      "create",
      metadata,
    );
  }

  async replace(
    containerId: string,
    body: unknown,
    partitionKey?: string,
    metadata?: TraceMetadata,
  ): Promise<unknown> {
    const request = {
      pointOperation: {
        containerId,
        replace: {
          body: typeof body === "string" ? body : JSON.stringify(body),
          partitionKey,
        },
      },
    };

    return this.exec(
      request as unknown as Record<string, unknown>,
      "replace",
      metadata,
    );
  }

  async upsert(
    containerId: string,
    body: unknown,
    partitionKey?: string,
    metadata?: TraceMetadata,
  ): Promise<unknown> {
    const request = {
      pointOperation: {
        containerId,
        upsert: {
          body: typeof body === "string" ? body : JSON.stringify(body),
          partitionKey,
        },
      },
    };

    return this.exec(
      request as unknown as Record<string, unknown>,
      "upsert",
      metadata,
    );
  }

  async deleteItem(
    containerId: string,
    id: string,
    partitionKey?: string,
    metadata?: TraceMetadata,
  ): Promise<unknown> {
    const request = {
      pointOperation: {
        containerId,
        delete: { id, partitionKey },
      },
    };

    return this.exec(
      request as unknown as Record<string, unknown>,
      "delete",
      metadata,
    );
  }
}