import { readdir, readFile } from "fs/promises";
import path from "path";

interface DocumentChunk {
  id: string;
  content: string;
  metadata: {
    source: string;
    title: string;
    section: string;
    chunkIndex: number;
  };
}

async function loadDocuments(directory: string): Promise<DocumentChunk[]> {
  const files = await readdir(directory, { recursive: true });
  const mdFiles = files.filter((f) => f.toString().endsWith(".md"));
  const chunks: DocumentChunk[] = [];

  for (const file of mdFiles) {
    const filePath = path.join(directory, file.toString());
    const content = await readFile(filePath, "utf-8");
    const fileChunks = chunkDocument(content, filePath);
    chunks.push(...fileChunks);
  }

  return chunks;
}

function chunkDocument(
  content: string,
  source: string,
  maxChunkSize: number = 800,
  overlap: number = 100
): DocumentChunk[] {
  const chunks: DocumentChunk[] = [];

  // Split by headers first (semantic boundaries)
  const sections = content.split(/\n(?=#{1,3} )/);

  let chunkIndex = 0;

  for (const section of sections) {
    const title = section.match(/^#{1,3} (.+)/)?.[1] || "Untitled";

    if (section.length <= maxChunkSize) {
      // Section fits in one chunk
      chunks.push({
        id: `\${source}-\${chunkIndex}`,
        content: section.trim(),
        metadata: {
          source,
          title,
          section: title,
          chunkIndex: chunkIndex++,
        },
      });
    } else {
      // Split long sections by paragraphs with overlap
      const paragraphs = section.split(/\n\n+/);
      let currentChunk = "";

      for (const paragraph of paragraphs) {
        if ((currentChunk + paragraph).length > maxChunkSize && currentChunk) {
          chunks.push({
            id: `\${source}-\${chunkIndex}`,
            content: currentChunk.trim(),
            metadata: {
              source,
              title,
              section: title,
              chunkIndex: chunkIndex++,
            },
          });

          // Keep overlap from previous chunk
          const words = currentChunk.split(" ");
          currentChunk =
            words.slice(-Math.floor(overlap / 5)).join(" ") +
            "\n\n" +
            paragraph;
        } else {
          currentChunk += (currentChunk ? "\n\n" : "") + paragraph;
        }
      }

      // Don't forget the last chunk
      if (currentChunk.trim()) {
        chunks.push({
          id: `\${source}-\${chunkIndex}`,
          content: currentChunk.trim(),
          metadata: {
            source,
            title,
            section: title,
            chunkIndex: chunkIndex++,
          },
        });
      }
    }
  }

  return chunks;
}

import Anthropic from "@anthropic-ai/sdk";

const anthropic = new Anthropic();

async function embedText(text: string): Promise<number[]> {
  // Use an embedding model
  // Anthropic doesn't have embeddings yet, so we use OpenAI or Cohere
  const response = await fetch("https://api.openai.com/v1/embeddings", {
    method: "POST",
    headers: {
      Authorization: `Bearer \${process.env.OPENAI_API_KEY}`,
      "Content-Type": "application/json",
    },
    body: JSON.stringify({
      model: "text-embedding-3-small",
      input: text,
    }),
  });

  const data = await response.json();
  return data.data[0].embedding;
}

async function indexChunks(chunks: DocumentChunk[]) {
  for (const chunk of chunks) {
    const embedding = await embedText(chunk.content);

    await vectorDb.upsert({
      id: chunk.id,
      values: embedding,
      metadata: {
        content: chunk.content,
        ...chunk.metadata,
      },
    });
  }

  console.log(`Indexed \${chunks.length} chunks`);
}

async function retrieveRelevant(
  query: string,
  topK: number = 5
): Promise<DocumentChunk[]> {
  const queryEmbedding = await embedText(query);

  const results = await vectorDb.query({
    vector: queryEmbedding,
    topK,
    includeMetadata: true,
  });

  return results.matches.map((match) => ({
    id: match.id,
    content: match.metadata!.content as string,
    metadata: {
      source: match.metadata!.source as string,
      title: match.metadata!.title as string,
      section: match.metadata!.section as string,
      chunkIndex: match.metadata!.chunkIndex as number,
    },
    score: match.score,
  }));
}

async function answerQuestion(userQuery: string): Promise<string> {
  const relevantChunks = await retrieveRelevant(userQuery, 5);

  const context = relevantChunks
    .map(
      (chunk) =>
        `[Source: \${chunk.metadata.source} - \${chunk.metadata.title}]\n\${chunk.content}`
    )
    .join("\n\n---\n\n");

  const response = await anthropic.messages.create({
    model: "claude-sonnet-4-20250514",
    max_tokens: 1024,
    system: `You are a helpful assistant that answers questions based on the provided documentation.

Rules:
- ONLY answer based on the provided context. If the context doesn't contain the answer, say so.
- Cite your sources by mentioning the document name.
- Be concise. Answer in 2-3 sentences unless more detail is needed.
- Never make up information that isn't in the context.`,
    messages: [
      {
        role: "user",
        content: `Context documents:

\${context}

---

Question: \${userQuery}`,
      },
    ],
  });

  return response.content[0].type === "text" ? response.content[0].text : "";
}

const testCases = [
  {
    question: "What is the refund policy?",
    expectedSource: "policies/refund.md",
    expectedContent: ["30 days", "full refund", "no questions asked"],
  },
  {
    question: "How do I reset my password?",
    expectedSource: "help/account.md",
    expectedContent: ["settings", "security", "email verification"],
  },
];

async function evaluateRAG() {
  let passed = 0;

  for (const test of testCases) {
    const chunks = await retrieveRelevant(test.question);

    // Check if the right source was retrieved
    const sourceFound = chunks.some((c) =>
      c.metadata.source.includes(test.expectedSource)
    );

    // Check if key content was in retrieved chunks
    const allContent = chunks.map((c) => c.content).join(" ");
    const contentFound = test.expectedContent.every((keyword) =>
      allContent.toLowerCase().includes(keyword.toLowerCase())
    );

    if (sourceFound && contentFound) passed++;
    else {
      console.log(`FAILED: "\${test.question}"`);
      console.log(`  Source found: \${sourceFound}`);
      console.log(`  Content found: \${contentFound}`);
    }
  }

  console.log(`\nRetrieval accuracy: \${passed}/\${testCases.length} (\${Math.round((passed / testCases.length) * 100)}%)`);
}