import { BedrockRuntimeClient, InvokeModelCommand } from "@aws-sdk/client-bedrock-runtime";
import { Pinecone } from "@pinecone-database/pinecone";

const bedrock = new BedrockRuntimeClient({});
const pc = new Pinecone({ apiKey: process.env.PINECONE_API_KEY });

export const handler = async (event) => {
    const { question } = JSON.parse(event.body);

    // 1. Embed user query
    const embedResponse = await bedrock.send(new InvokeModelCommand({
        modelId: "amazon.titan-embed-text-v1",
        contentType: "application/json",
        accept: "application/json",
        body: JSON.stringify({ inputText: question })
    }));
    const { embedding } = JSON.parse(new TextDecoder().decode(embedResponse.body));

    // 2. Query Vector DB for relevant context
    const index = pc.index(process.env.PINECONE_INDEX_NAME);
    const queryResponse = await index.query({
        vector: embedding,
        topK: 3,
        includeMetadata: true
    });

    const context = queryResponse.matches.map(match => match.metadata.text).join("\n\n");

    const systemPrompt = `Use the following context to answer the question. If you don't know, say you don't know.\n\nContext:\n${context}`;

    const llmResponse = await bedrock.send(new InvokeModelCommand({
        modelId: "anthropic.claude-3-haiku-20240307-v1:0", // Fast & Cost-effective
        contentType: "application/json",
        accept: "application/json",
        body: JSON.stringify({
            anthropic_version: "bedrock-2023-05-31",
            max_tokens: 500,
            messages: [
                { role: "user", content: `${systemPrompt}\n\nQuestion: ${question}` }
            ]
        })
    }));

    const result = JSON.parse(new TextDecoder().decode(llmResponse.body));
    const answer = result.content[0].text;

    return {
        statusCode: 200,
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({ answer, contextUsed: queryResponse.matches.length })
    };
};