import { generateResponse } from "./ollama.js";
import { searchDocuments, type SearchResult } from "./search.js";

export interface AnswerResult {
  answer: string;
  sources: {
    documentId: string;
    title: string;
    content: string;
    similarity: number;
  }[];
}

interface AnswerDependencies {
  search: (query: string, limit?: number) => Promise<SearchResult[]>;
  generate: (prompt: string) => Promise<string>;
}

const defaultDependencies: AnswerDependencies = {
  search: searchDocuments,
  generate: generateResponse,
};

export const answerQuestion = async (
  question: string,
  dependencies: AnswerDependencies = defaultDependencies,
): Promise<AnswerResult> => {
  const normalizedQuestion = question.trim();

  if (!normalizedQuestion) {
    throw new Error("Question cannot be empty");
  }

  const results = await dependencies.search(normalizedQuestion, 5);
  const relevantResults = results.filter((result) => result.similarity >= 0.5);

  if (relevantResults.length === 0) {
    return {
      answer:
        "I couldn't find relevant information in the available documents.",
      sources: [],
    };
  }

  const context = relevantResults
    .map(
      (result, index) =>
        `[Source ${index + 1}: ${result.title}]\n${result.content}`,
    )
    .join("\n\n");

  const prompt = `Answer the question using only the context provided below.
If the context does not contain the answer, say that you don't have enough information.
Do not invent facts or claim that the context supports something it does not.

Context:
${context}

Question: ${normalizedQuestion}

Answer:`;

  const answer = await dependencies.generate(prompt);

  return {
    answer,
    sources: relevantResults.map((result) => ({
      documentId: result.documentId,
      title: result.title,
      content: result.content,
      similarity: result.similarity,
    })),
  };
};