import { Injectable, Logger } from '@nestjs/common';
import { InjectRepository } from '@nestjs/typeorm';
import { Repository } from 'typeorm';
import { ConversationEmbedding, MessageRole } from '../../../database/entities';
import axios from 'axios';

/**
 * Embeddings Service
 *
 * Generates and manages vector embeddings for conversation messages using OpenAI or Azure OpenAI
 */

export interface EmbeddingGenerationResult {
  embedding: number[];
  tokenCount: number;
  model: string;
}

export interface SimilarConversation {
  conversationId: string;
  messageContent: string;
  similarity: number;
  createdAt: Date;
  topics: string[];
}

@Injectable()
export class EmbeddingsService {
  private readonly logger = new Logger(EmbeddingsService.name);

  // Configuration from environment
  private readonly OPENAI_API_KEY = process.env.AZURE_OPENAI_EMBEDDINGS_API_KEY;
  private readonly OPENAI_ENDPOINT =
    process.env.AZURE_OPENAI_EMBEDDINGS_ENDPOINT;
  private readonly OPENAI_DEPLOYMENT =
    process.env.AZURE_OPENAI_EMBEDDINGS_DEPLOYMENT || 'text-embedding-ada-002';
  private readonly OPENAI_API_VERSION =
    process.env.AZURE_OPENAI_EMBEDDINGS_API_VERSION || '2023-05-15';

  // Embedding configuration
  private readonly EMBEDDING_DIMENSION = 1536; // OpenAI text-embedding-ada-002
  private readonly BATCH_SIZE = 100; // Max embeddings per batch

  constructor(
    @InjectRepository(ConversationEmbedding)
    private embeddingRepository: Repository<ConversationEmbedding>,
  ) {}

  /**
   * Generate embedding for a single text using Azure OpenAI
   */
  async generateEmbedding(text: string): Promise<EmbeddingGenerationResult> {
    try {
      // Azure OpenAI embeddings endpoint
      const url = `${this.OPENAI_ENDPOINT}/openai/deployments/${this.OPENAI_DEPLOYMENT}/embeddings?api-version=${this.OPENAI_API_VERSION}`;

      const response = await axios.post(
        url,
        {
          input: text,
          model: this.OPENAI_DEPLOYMENT,
        },
        {
          headers: {
            'api-key': this.OPENAI_API_KEY,
            'Content-Type': 'application/json',
          },
          timeout: 30000, // 30s timeout
        },
      );

      const embedding = response.data.data[0].embedding;
      const tokenCount = response.data.usage.total_tokens;

      if (embedding.length !== this.EMBEDDING_DIMENSION) {
        throw new Error(
          `Expected ${this.EMBEDDING_DIMENSION} dimensions, got ${embedding.length}`,
        );
      }

      this.logger.debug(
        `Generated embedding for text (${tokenCount} tokens, ${embedding.length} dimensions)`,
      );

      return {
        embedding,
        tokenCount,
        model: this.OPENAI_DEPLOYMENT,
      };
    } catch (error) {
      this.logger.error('Failed to generate embedding', error.stack);
      throw new Error(`Embedding generation failed: ${error.message}`);
    }
  }

  /**
   * Generate embeddings for multiple texts in batch
   */
  async generateEmbeddingsBatch(
    texts: string[],
  ): Promise<EmbeddingGenerationResult[]> {
    if (texts.length === 0) {
      return [];
    }

    if (texts.length > this.BATCH_SIZE) {
      this.logger.warn(
        `Batch size ${texts.length} exceeds maximum ${this.BATCH_SIZE}, splitting into multiple requests`,
      );
      const results: EmbeddingGenerationResult[] = [];
      for (let i = 0; i < texts.length; i += this.BATCH_SIZE) {
        const batch = texts.slice(i, i + this.BATCH_SIZE);
        const batchResults = await this.generateEmbeddingsBatch(batch);
        results.push(...batchResults);
      }
      return results;
    }

    try {
      const url = `${this.OPENAI_ENDPOINT}/openai/deployments/${this.OPENAI_DEPLOYMENT}/embeddings?api-version=${this.OPENAI_API_VERSION}`;

      const response = await axios.post(
        url,
        {
          input: texts,
          model: this.OPENAI_DEPLOYMENT,
        },
        {
          headers: {
            'api-key': this.OPENAI_API_KEY,
            'Content-Type': 'application/json',
          },
          timeout: 60000, // 60s timeout for batch
        },
      );

      return response.data.data.map((item: any) => ({
        embedding: item.embedding,
        tokenCount: response.data.usage.total_tokens / texts.length, // Average
        model: this.OPENAI_DEPLOYMENT,
      }));
    } catch (error) {
      this.logger.error('Failed to generate embeddings batch', error.stack);
      throw new Error(`Batch embedding generation failed: ${error.message}`);
    }
  }

  /**
   * Store embedding for a conversation message
   */
  async storeEmbedding(
    conversationId: string,
    userId: string,
    messageIndex: number,
    messageRole: MessageRole,
    messageContent: string,
    topics: string[],
  ): Promise<ConversationEmbedding> {
    // Generate embedding
    const { embedding } = await this.generateEmbedding(messageContent);

    // Create embedding entity
    const embeddingEntity = this.embeddingRepository.create({
      conversationId,
      userId,
      messageIndex,
      messageRole,
      messageContent,
      embedding: ConversationEmbedding.vectorToString(embedding),
      topics,
    });

    // Save to database
    const saved = await this.embeddingRepository.save(embeddingEntity);

    this.logger.debug(
      `Stored embedding ${saved.id} for conversation ${conversationId}, message ${messageIndex}`,
    );

    return saved;
  }

  /**
   * Search for similar conversations using vector similarity
   */
  async searchSimilarConversations(
    queryText: string,
    userId: string,
    options: {
      similarityThreshold?: number;
      limit?: number;
      topicFilter?: string;
    } = {},
  ): Promise<SimilarConversation[]> {
    const { similarityThreshold = 0.7, limit = 5, topicFilter } = options;

    // Generate embedding for query text
    const { embedding: queryEmbedding } =
      await this.generateEmbedding(queryText);
    const queryVector = ConversationEmbedding.vectorToString(queryEmbedding);

    try {
      let query;

      if (topicFilter) {
        // Use topic-filtered search function
        query = this.embeddingRepository.query(
          `
            SELECT * FROM search_conversations_by_topic(
              $1::vector,
              $2,
              $3,
              $4,
              $5
            )
          `,
          [queryVector, userId, topicFilter, similarityThreshold, limit],
        );
      } else {
        // Use general similarity search function
        query = this.embeddingRepository.query(
          `
            SELECT * FROM search_similar_conversations(
              $1::vector,
              $2,
              $3,
              $4
            )
          `,
          [queryVector, userId, similarityThreshold, limit],
        );
      }

      const results = await query;

      this.logger.debug(
        `Found ${results.length} similar conversations for user ${userId} (threshold: ${similarityThreshold})`,
      );

      return results.map((row: any) => ({
        conversationId: row.conversation_id,
        messageContent: row.message_content,
        similarity: parseFloat(row.similarity),
        createdAt: new Date(row.created_at),
        topics: row.topics,
      }));
    } catch (error) {
      this.logger.error('Failed to search similar conversations', error.stack);
      throw new Error(`Similarity search failed: ${error.message}`);
    }
  }

  /**
   * Get embeddings for a conversation
   */
  async getConversationEmbeddings(
    conversationId: string,
  ): Promise<ConversationEmbedding[]> {
    return this.embeddingRepository.find({
      where: { conversationId },
      order: { messageIndex: 'ASC' },
    });
  }

  /**
   * Delete embeddings for a conversation
   */
  async deleteConversationEmbeddings(conversationId: string): Promise<void> {
    await this.embeddingRepository.delete({ conversationId });
    this.logger.debug(`Deleted embeddings for conversation ${conversationId}`);
  }

  /**
   * Bulk create embeddings for existing conversations (migration/backfill)
   */
  async backfillEmbeddings(
    conversationId: string,
    userId: string,
    messages: Array<{
      index: number;
      role: MessageRole;
      content: string;
    }>,
    topics: string[],
  ): Promise<number> {
    if (messages.length === 0) {
      return 0;
    }

    // Check if embeddings already exist
    const existingCount = await this.embeddingRepository.count({
      where: { conversationId },
    });

    if (existingCount > 0) {
      this.logger.debug(
        `Conversation ${conversationId} already has ${existingCount} embeddings, skipping`,
      );
      return 0;
    }

    // Generate embeddings in batch
    const texts = messages.map((m) => m.content);
    const embeddingResults = await this.generateEmbeddingsBatch(texts);

    // Create embedding entities
    const entities = messages.map((msg, i) =>
      this.embeddingRepository.create({
        conversationId,
        userId,
        messageIndex: msg.index,
        messageRole: msg.role,
        messageContent: msg.content,
        embedding: ConversationEmbedding.vectorToString(
          embeddingResults[i].embedding,
        ),
        topics,
      }),
    );

    // Bulk save
    await this.embeddingRepository.save(entities);

    this.logger.log(
      `Backfilled ${entities.length} embeddings for conversation ${conversationId}`,
    );

    return entities.length;
  }

  /**
   * Get embedding statistics for a user
   */
  async getUserEmbeddingStats(userId: string): Promise<{
    totalEmbeddings: number;
    conversationsWithEmbeddings: number;
    topicsDistribution: Record<string, number>;
  }> {
    const embeddings = await this.embeddingRepository.find({
      where: { userId },
    });

    const conversationIds = new Set(embeddings.map((e) => e.conversationId));

    const topicsDistribution: Record<string, number> = {};
    for (const embedding of embeddings) {
      for (const topic of embedding.topics) {
        topicsDistribution[topic] = (topicsDistribution[topic] || 0) + 1;
      }
    }

    return {
      totalEmbeddings: embeddings.length,
      conversationsWithEmbeddings: conversationIds.size,
      topicsDistribution,
    };
  }

  /**
   * Health check: verify embeddings service is configured correctly
   */
  async healthCheck(): Promise<{ status: 'ok' | 'error'; message: string }> {
    if (!this.OPENAI_API_KEY || !this.OPENAI_ENDPOINT) {
      return {
        status: 'error',
        message: 'Azure OpenAI credentials not configured',
      };
    }

    try {
      // Test embedding generation
      await this.generateEmbedding('Health check test');
      return { status: 'ok', message: 'Embeddings service operational' };
    } catch (error) {
      return {
        status: 'error',
        message: `Health check failed: ${error.message}`,
      };
    }
  }
}
