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

const GENERATION_MODEL = "claude-opus-4-6";
const NLI_MODEL = "claude-haiku-4-5";


const RESPONSE_SCHEMA = {
    type: "object",
    properties: {
        answer: {
            type: "string",
            description:
                "The answer with inline numeric citations like [1], [2] matching the source document numbers provided",
        },
        citedSourceIndices: {
            type: "array",
            items: { type: "integer" },
            description: "The document numbers cited in the answer",
        },
        followUpQuestions: {
            type: "array",
            items: { type: "string" },
            description: "Exactly 3 follow-up questions the user might ask",
        },
    },
    required: ["answer", "citedSourceIndices", "followUpQuestions"],
    additionalProperties: false,
};

const NLI_SCHEMA = {
    type: "object",
    properties: {
        label: {
            type: "string",
            enum: ["supported", "contradicted", "unsupported"],
            description: "Whether the claim is entailed, contradicted, or neither by the passages",
        },
        confidence: {
            type: "number",
            description: "Confidence in the label, from 0 to 1",
        },
    },
    required: ["label", "confidence"],
    additionalProperties: false,
};

function getClient() {
    const key = process.env.ANTHROPIC_API_KEY;
    if (!key) throw new Error("ANTHROPIC_API_KEY is not set");
    return new Anthropic({ apiKey: key });
}

function textFromMessage(message) {
    const block = message.content[0];
    return block && "text" in block ? block.text : "";
}

function buildPrompt(query, sources) {
    const today = new Date().toLocaleDateString("en-US", {
        weekday: "long",
        year: "numeric",
        month: "long",
        day: "numeric",
    });

    const sourceBlock = sources
        .map((s) => `[Source ${s.docIndex}] (${s.name})\n${s.text}`)
        .join("\n\n---\n\n");

    return `You are a research assistant. Today's date is ${today}. Answer the user's question using ONLY the source documents provided below. Follow these rules strictly:

1. Ground every claim in a specific source document. Cite sources inline using numeric notation like [1], [2], etc., matching the source document numbers below.
2. If the sources do not contain enough information to answer, say so honestly.
3. After your answer, suggest exactly 3 follow-up questions the user might ask based on the sources.

--- SOURCE DOCUMENTS ---

${sourceBlock}

--- USER QUESTION ---

${query}`;
}

function stripInvalidCitations(answer, validIndices) {
    return answer
        .replace(/\[(\d+)\]/g, (match, n) => (validIndices.has(Number(n)) ? match : ""))
        .replace(/[ \t]{2,}/g, " ")
        .trim();
}

export async function generate(query, sources) {
    const client = getClient();
    const message = await client.messages.create({
        model: GENERATION_MODEL,
        max_tokens: 2048,
        messages: [{ role: "user", content: buildPrompt(query, sources) }],
        output_config: {
            format: {
                type: "json_schema",
                schema: RESPONSE_SCHEMA,
            },
        },
    });

    const parsed = JSON.parse(textFromMessage(message) || "{}");
    const validIndices = new Set(sources.map((s) => s.docIndex));
    const citedSourceIndices = [
        ...new Set((parsed.citedSourceIndices || []).filter((idx) => validIndices.has(idx))),
    ];

    return {
        answer: stripInvalidCitations(parsed.answer || "", validIndices),
        citedSourceIndices,
        followUpQuestions: parsed.followUpQuestions || [],
    };
}

export async function judgeEntailment(sentence, chunkTexts) {
    const client = getClient();
    const passages = chunkTexts.map((t, i) => `[Passage ${i + 1}]\n${t}`).join("\n\n");

    const prompt = `You are an NLI judge. Decide whether the CLAIM is supported, contradicted, or unsupported by the SOURCE PASSAGES.

- supported: the passages entail the claim
- contradicted: the passages contradict the claim
- unsupported: the passages neither support nor contradict the claim

--- CLAIM ---
${sentence}

--- SOURCE PASSAGES ---
${passages}`;

    const message = await client.messages.create({
        model: NLI_MODEL,
        max_tokens: 256,
        messages: [{ role: "user", content: prompt }],
        output_config: {
            format: {
                type: "json_schema",
                schema: NLI_SCHEMA,
            },
        },
    });

    const parsed = JSON.parse(textFromMessage(message) || "{}");
    const label = ["supported", "contradicted", "unsupported"].includes(parsed.label)
        ? parsed.label
        : "unsupported";
    const confidence = typeof parsed.confidence === "number" ? parsed.confidence : 0;

    return { label, confidence };
}