import {
StateGraph,
Annotation,
MessagesAnnotation,
END,
} from "@langchain/langgraph";
import { ChatOpenAI } from "@langchain/openai";
import { HumanMessage, SystemMessage } from "@langchain/core/messages";
import { PineconeStore } from "@langchain/pinecone";
import { embeddings } from "./embeddings";
import { promptTemplate } from "./ragPrompt";

const model = new ChatOpenAI({
modelName: "gpt-4",
temperature: 0,
openAIApiKey: process.env.OPENAI_API_KEY,
});

let vectorStore: PineconeStore | null = null;

async function getVectorStore() {
console.log("[CRAG] getVectorStore called, cached:", !!vectorStore);
if (!vectorStore) {
const { pineconeIndex } = await import("./pinecone");
vectorStore = await PineconeStore.fromExistingIndex(embeddings, {
pineconeIndex,
});
console.log("[CRAG] vectorStore initialized");
}
return vectorStore;
}

const CRAGState = Annotation.Root({
...MessagesAnnotation.spec,
question: Annotation({ reducer: (, b) => b, default: () => "" }),
namespace: Annotation({ reducer: (, b) => b, default: () => "" }),
documents: Annotation({ reducer: (, b) => b, default: () => [] }),
retrievalGrade: Annotation({
reducer: (, b) => b,
default: () => "",
}),
rewrittenQuestion: Annotation({
reducer: (, b) => b,
default: () => "",
}),
answer: Annotation({ reducer: (, b) => b, default: () => "" }),
retryCount: Annotation({ reducer: (_, b) => b, default: () => 0 }),
});

async function retrieve(state: typeof CRAGState.State) {
console.log("[CRAG] retrieve node started");
const store = await getVectorStore();
const query = state.rewrittenQuestion || state.question;
console.log("[CRAG] retrieve query:", query, "namespace:", state.namespace);
const results = await store.similaritySearch(query, 5, {
namespace: state.namespace,
});
const docs = results.map((doc) => doc.pageContent);
console.log("[CRAG] retrieve got", docs.length, "docs");
return { documents: docs };
}

async function gradeDocuments(state: typeof CRAGState.State) {
console.log("[CRAG] LLM grading retrieval...");

const context = state.documents.join("\n\n");

const response = await model.invoke([
new SystemMessage(`
You are a retrieval evaluator.

Determine whether the retrieved documents
are relevant enough to answer the user's question.

Reply ONLY with:

• good

• bad
  `),

  new HumanMessage(`
  Question:
  ${state.question}

Retrieved Documents:
${context}
`),
]);

const grade =
typeof response.content === "string"
? response.content.toLowerCase().trim()
: "bad";

console.log("[CRAG] LLM grade:", grade);

return {
retrievalGrade: grade,
};
}

async function rewriteQuery(state: typeof CRAGState.State) {
console.log("[CRAG] rewriteQuery node started");
const response = await model.invoke([
new SystemMessage("Rewrite the user's query to improve vector retrieval."),
new HumanMessage(state.question),
]);
const rewritten =
typeof response.content === "string" ? response.content : state.question;
console.log("[CRAG] rewriteQuery new query:", rewritten);
return { rewrittenQuestion: rewritten, retryCount: state.retryCount + 1 };
}

async function generateAnswer(state: typeof CRAGState.State) {
console.log("[CRAG] generateAnswer node started");
const combinedContext = state.documents.join("\n\n");
const formattedPrompt = await promptTemplate.format({
context: combinedContext,
});
const response = await model.invoke([
new SystemMessage(formattedPrompt),
new HumanMessage(state.question),
]);
const answer = typeof response.content === "string" ? response.content : "";
console.log("[CRAG] generateAnswer produced answer length:", answer.length);
return { answer, messages: [response] };
}

function shouldRetry(state: typeof CRAGState.State) {
console.log(
"[CRAG] shouldRetry grade:",
state.retrievalGrade,
"retryCount:",
state.retryCount,
);
if (state.retryCount >= 2) return "generate";
if (state.retrievalGrade === "bad") return "rewrite_query";
return "generate";
}

console.log("[CRAG] Graph compiled");
export const cragGraph = new StateGraph(CRAGState)
.addNode("retrieve", retrieve)
.addNode("grade_documents", gradeDocuments)
.addNode("rewrite_query", rewriteQuery)
.addNode("generate", generateAnswer)
.addEdge("start", "retrieve")
.addEdge("retrieve", "grade_documents")
.addConditionalEdges("grade_documents", shouldRetry, {
rewrite_query: "rewrite_query",
generate: "generate",
})
.addEdge("rewrite_query", "retrieve")
.addEdge("generate", END)
.compile();