import { Injectable, Logger } from "@nestjs/common";
import { EntityType } from "@prisma/client";
import { EMBEDDING_DIMENSION } from "@mosaic/shared";
import { PrismaService } from "../prisma/prisma.service";

/**
 * Result from similarity search
 */
export interface SimilarEmbedding {
  id: string;
  content: string;
  similarity: number;
  entityType: EntityType | null;
  entityId: string | null;
  metadata: Record<string, unknown>;
}

/**
 * Service for managing vector embeddings using pgvector
 * Uses raw SQL for vector operations since Prisma doesn't support vector types natively
 */
@Injectable()
export class EmbeddingsService {
  private readonly logger = new Logger(EmbeddingsService.name);

  constructor(private readonly prisma: PrismaService) {}

  /**
   * Validate that an embedding array contains only finite numbers
   * @param embedding Array to validate
   * @throws Error if validation fails
   */
  private validateEmbedding(embedding: number[]): void {
    if (!Array.isArray(embedding)) {
      throw new Error("Embedding must be an array");
    }

    if (!embedding.every((val) => typeof val === "number" && Number.isFinite(val))) {
      throw new Error("Embedding array must contain only finite numbers");
    }
  }

  /**
   * Store an embedding vector for content
   * @param params Embedding parameters
   * @returns ID of the created embedding
   */
  async storeEmbedding(params: {
    workspaceId: string;
    content: string;
    embedding: number[];
    entityType?: EntityType;
    entityId?: string;
    metadata?: Record<string, unknown>;
  }): Promise<string> {
    const { workspaceId, content, embedding, entityType, entityId, metadata } = params;

    // Validate embedding array
    this.validateEmbedding(embedding);

    if (embedding.length !== EMBEDDING_DIMENSION) {
      throw new Error(
        `Invalid embedding dimension: expected EMBEDDING_DIMENSION, got ${embedding.length.toString()}`
      );
    }

    const vectorString = `[${embedding.join(",")}]`;

    try {
      const result = await this.prisma.$queryRaw<{ id: string }[]>`
        INSERT INTO memory_embeddings (
          id, workspace_id, content, embedding, entity_type, entity_id, metadata, created_at, updated_at
        )
        VALUES (
          gen_random_uuid(),
          ${workspaceId}::uuid,
          ${content},
          ${vectorString}::vector,
          ${entityType ?? null}::"EntityType",
          ${entityId ?? null}::uuid,
          ${JSON.stringify(metadata ?? {})}::jsonb,
          NOW(),
          NOW()
        )
        RETURNING id::text
      `;

      const embeddingId = result[0]?.id;
      if (!embeddingId) {
        throw new Error("Failed to get embedding ID from insert result");
      }
      this.logger.debug(`Stored embedding ${embeddingId} for workspace ${workspaceId}`);
      return embeddingId;
    } catch (error) {
      this.logger.error("Failed to store embedding", error);
      throw error;
    }
  }

  /**
   * Find similar embeddings using cosine similarity
   * @param params Search parameters
   * @returns Array of similar embeddings sorted by similarity (descending)
   */
  async findSimilar(params: {
    workspaceId: string;
    embedding: number[];
    limit?: number;
    threshold?: number;
    entityType?: EntityType;
  }): Promise<SimilarEmbedding[]> {
    const { workspaceId, embedding, limit = 10, threshold = 0.7, entityType } = params;

    // Validate embedding array
    this.validateEmbedding(embedding);

    if (embedding.length !== EMBEDDING_DIMENSION) {
      throw new Error(
        `Invalid embedding dimension: expected EMBEDDING_DIMENSION, got ${embedding.length.toString()}`
      );
    }

    const vectorString = `[${embedding.join(",")}]`;

    try {
      let results: SimilarEmbedding[];

      if (entityType) {
        results = await this.prisma.$queryRaw<SimilarEmbedding[]>`
          SELECT
            id::text,
            content,
            1 - (embedding <=> ${vectorString}::vector) as similarity,
            entity_type as "entityType",
            entity_id::text as "entityId",
            metadata
          FROM memory_embeddings
          WHERE workspace_id = ${workspaceId}::uuid
            AND embedding IS NOT NULL
            AND 1 - (embedding <=> ${vectorString}::vector) >= ${threshold}
            AND entity_type = ${entityType}::"EntityType"
          ORDER BY embedding <=> ${vectorString}::vector
          LIMIT ${limit}
        `;
      } else {
        results = await this.prisma.$queryRaw<SimilarEmbedding[]>`
          SELECT
            id::text,
            content,
            1 - (embedding <=> ${vectorString}::vector) as similarity,
            entity_type as "entityType",
            entity_id::text as "entityId",
            metadata
          FROM memory_embeddings
          WHERE workspace_id = ${workspaceId}::uuid
            AND embedding IS NOT NULL
            AND 1 - (embedding <=> ${vectorString}::vector) >= ${threshold}
          ORDER BY embedding <=> ${vectorString}::vector
          LIMIT ${limit}
        `;
      }

      this.logger.debug(
        `Found ${results.length.toString()} similar embeddings for workspace ${workspaceId}`
      );
      return results;
    } catch (error) {
      this.logger.error("Failed to find similar embeddings", error);
      throw error;
    }
  }

  /**
   * Delete embeddings for a specific entity
   * @param params Entity identifiers
   * @returns Number of embeddings deleted
   */
  async deleteByEntity(params: {
    workspaceId: string;
    entityType: EntityType;
    entityId: string;
  }): Promise<number> {
    const { workspaceId, entityType, entityId } = params;

    try {
      const result = await this.prisma.$executeRaw`
        DELETE FROM memory_embeddings
        WHERE workspace_id = ${workspaceId}::uuid
          AND entity_type = ${entityType}::"EntityType"
          AND entity_id = ${entityId}::uuid
      `;

      this.logger.debug(
        `Deleted ${result.toString()} embeddings for ${entityType}:${entityId} in workspace ${workspaceId}`
      );
      return result;
    } catch (error) {
      this.logger.error("Failed to delete embeddings", error);
      throw error;
    }
  }

  /**
   * Delete all embeddings for a workspace
   * @param workspaceId Workspace ID
   * @returns Number of embeddings deleted
   */
  async deleteByWorkspace(workspaceId: string): Promise<number> {
    try {
      const result = await this.prisma.$executeRaw`
        DELETE FROM memory_embeddings
        WHERE workspace_id = ${workspaceId}::uuid
      `;

      this.logger.debug(`Deleted ${result.toString()} embeddings for workspace ${workspaceId}`);
      return result;
    } catch (error) {
      this.logger.error("Failed to delete workspace embeddings", error);
      throw error;
    }
  }

  /**
   * Get embedding by ID
   * @param id Embedding ID
   * @returns Embedding or null if not found
   */
  async getById(id: string): Promise<SimilarEmbedding | null> {
    try {
      const results = await this.prisma.$queryRaw<SimilarEmbedding[]>`
        SELECT
          id::text,
          content,
          0 as similarity,
          entity_type as "entityType",
          entity_id::text as "entityId",
          metadata
        FROM memory_embeddings
        WHERE id = ${id}::uuid
        LIMIT 1
      `;

      return results.length > 0 ? (results[0] ?? null) : null;
    } catch (error) {
      this.logger.error(`Failed to get embedding ${id}`, error);
      throw error;
    }
  }
}
