import { Pinecone } from "@pinecone-database/pinecone";
import { type BaseMessage, BaseRetriever, type RetrieveOptions } from "@voltagent/core";

// Initialize Pinecone client
const pc = new Pinecone({
  apiKey: process.env.PINECONE_API_KEY || "",
  sourceTag: "voltagent",
});

const indexName = "voltagent-knowledge-base";

// Sample documents to populate the index
const sampleRecords = [
  {
    id: "doc1",
    values: [], // Will be generated by integrated embedding
    metadata: {
      text: "VoltAgent is a TypeScript framework for building AI agents with modular components.",
      category: "frameworks",
      topic: "voltagent",
    },
  },
  {
    id: "doc2",
    values: [],
    metadata: {
      text: "Pinecone is a vector database built for machine learning applications that require fast, accurate vector search.",
      category: "databases",
      topic: "pinecone",
    },
  },
  {
    id: "doc3",
    values: [],
    metadata: {
      text: "Vector databases store high-dimensional vectors and enable semantic search capabilities for AI applications.",
      category: "databases",
      topic: "vector-search",
    },
  },
  {
    id: "doc4",
    values: [],
    metadata: {
      text: "Retrieval-Augmented Generation (RAG) combines information retrieval with language generation for better AI responses.",
      category: "techniques",
      topic: "rag",
    },
  },
  {
    id: "doc5",
    values: [],
    metadata: {
      text: "TypeScript provides static typing for JavaScript, making code more reliable and maintainable for large applications.",
      category: "programming",
      topic: "typescript",
    },
  },
];

// Initialize index with sample data
async function initializeIndex() {
  try {
    // Check if index exists
    let indexExists = false;
    try {
      await pc.describeIndex(indexName);
      indexExists = true;
      console.log(`📋 Index "${indexName}" already exists`);
    } catch (_error) {
      console.log(`📋 Creating new index "${indexName}"...`);
    }

    // Create index if it doesn't exist
    if (!indexExists) {
      await pc.createIndex({
        name: indexName,
        dimension: 1536, // OpenAI text-embedding-3-small dimension
        metric: "cosine",
        spec: {
          serverless: {
            cloud: "aws",
            region: "us-east-1",
          },
        },
        waitUntilReady: true,
      });

      console.log(`✅ Index "${indexName}" created successfully`);
    }

    // Get the index
    const index = pc.index(indexName);

    // Check if we need to populate with sample data
    const stats = await index.describeIndexStats();

    if (stats.totalRecordCount === 0) {
      console.log("📚 Populating index with sample documents...");

      // Generate embeddings for sample documents using OpenAI
      const OpenAI = await import("openai");
      const openai = new OpenAI.default({
        apiKey: process.env.OPENAI_API_KEY || "",
      });

      const recordsWithEmbeddings = [];

      for (const record of sampleRecords) {
        try {
          const embeddingResponse = await openai.embeddings.create({
            model: "text-embedding-3-small",
            input: record.metadata.text,
          });

          recordsWithEmbeddings.push({
            id: record.id,
            values: embeddingResponse.data[0].embedding,
            metadata: record.metadata,
          });
        } catch (error) {
          console.error(`Error generating embedding for ${record.id}:`, error);
        }
      }

      // Upsert records to index
      if (recordsWithEmbeddings.length > 0) {
        await index.upsert(recordsWithEmbeddings);
        console.log(`✅ Successfully upserted ${recordsWithEmbeddings.length} documents to index`);
      }
    } else {
      console.log(`📊 Index already contains ${stats.totalRecordCount} documents`);
    }
  } catch (error) {
    console.error("Error initializing Pinecone index:", error);
  }
}

// Call initialization
initializeIndex();

// Retriever function
async function retrieveDocuments(query: string, topK = 3) {
  try {
    // Generate embedding for the query
    const OpenAI = await import("openai");
    const openai = new OpenAI.default({
      apiKey: process.env.OPENAI_API_KEY || "",
    });

    const embeddingResponse = await openai.embeddings.create({
      model: "text-embedding-3-small",
      input: query,
    });

    const queryVector = embeddingResponse.data[0].embedding;

    // Get the index and perform search
    const index = pc.index(indexName);

    const searchResults = await index.query({
      vector: queryVector,
      topK,
      includeMetadata: true,
      includeValues: false,
    });

    // Format results
    return (
      searchResults.matches?.map((match) => ({
        content: match.metadata?.text || "",
        metadata: match.metadata || {},
        score: match.score || 0,
        id: match.id,
      })) || []
    );
  } catch (error) {
    console.error("Error retrieving documents from Pinecone:", error);
    return [];
  }
}

/**
 * Pinecone-based retriever implementation for VoltAgent
 */
export class PineconeRetriever extends BaseRetriever {
  /**
   * Retrieve documents from Pinecone based on semantic similarity
   * @param input - The input to use for retrieval (string or BaseMessage[])
   * @param options - Configuration and context for the retrieval
   * @returns Promise resolving to a formatted context string
   */
  async retrieve(input: string | BaseMessage[], options: RetrieveOptions): Promise<string> {
    // Convert input to searchable string
    let searchText = "";

    if (typeof input === "string") {
      searchText = input;
    } else if (Array.isArray(input) && input.length > 0) {
      const lastMessage = input[input.length - 1];

      // Handle content as array of content parts with text type
      if (Array.isArray(lastMessage.content)) {
        const textParts = lastMessage.content
          .filter((part: any) => part.type === "text")
          .map((part: any) => part.text);

        searchText = textParts.join(" ");
      } else {
        // Fallback to string content
        searchText = lastMessage.content as string;
      }
    }

    // Perform semantic search using Pinecone
    const results = await retrieveDocuments(searchText, 3);

    // Add references to context if available
    if (options.context && results.length > 0) {
      const references = results.map((doc: any, index: number) => ({
        id: doc.id,
        title: doc.metadata.topic || `Document ${index + 1}`,
        source: "Pinecone Knowledge Base",
        score: doc.score,
        category: doc.metadata.category,
      }));

      options.context.set("references", references);
    }

    // Return the concatenated content for the LLM
    if (results.length === 0) {
      return "No relevant documents found in the knowledge base.";
    }

    return results
      .map(
        (doc: any, index: number) =>
          `Document ${index + 1} (ID: ${doc.id}, Score: ${doc.score.toFixed(4)}, Category: ${doc.metadata.category}):\n${doc.content}`,
      )
      .join("\n\n---\n\n");
  }
}

// Create retriever instance
export const retriever = new PineconeRetriever();
