import { z } from "@bpinternal/zui";
import pLimit from "p-limit";
import { ZaiContext } from "../context";
import { Response } from "../response";
import { getTokenizer } from "../tokenizer";
import { stringify } from "../utils";
import { Zai } from "../zai";
import { PROMPT_INPUT_BUFFER } from "./constants";
const _Options = z.object({
  examples: z.array(z.any()).default([]).describe("Examples to help guide answer generation"),
  instructions: z.string().optional().describe("Additional instructions for answer generation"),
  chunkLength: z.number().min(250).max(1e5).optional().describe("Maximum number of tokens per document chunk").default(16e3),
  maxRefinementPasses: z.number().min(1).max(10).optional().describe("Maximum number of refinement iterations when merging chunked results").default(3)
});
const ANSWER_START = "\u25A0answer";
const AMBIGUOUS_START = "\u25A0ambiguous";
const OUT_OF_TOPIC_START = "\u25A0out_of_topic";
const INVALID_QUESTION_START = "\u25A0invalid_question";
const MISSING_KNOWLEDGE_START = "\u25A0missing_knowledge";
const END = "\u25A0end\u25A0";
const formatDocumentsWithLineNumbers = (documents) => {
  const mappings = [];
  const allLines = [];
  let globalLineNumber = 1;
  let totalLines = 0;
  documents.forEach((doc) => {
    const docString = stringify(doc);
    const lines = docString.split("\n");
    totalLines += lines.length;
  });
  const padding = Math.max(3, totalLines.toString().length);
  documents.forEach((doc, docIndex) => {
    const docString = stringify(doc);
    const lines = docString.split("\n");
    lines.forEach((line, lineInDoc) => {
      const paddedNumber = globalLineNumber.toString().padStart(padding, "0");
      const formattedLine = `\u25A0${paddedNumber} | ${line}`;
      allLines.push(formattedLine);
      mappings.push({
        lineNumber: globalLineNumber,
        documentIndex: docIndex,
        lineInDocument: lineInDoc,
        text: line,
        document: doc
      });
      globalLineNumber++;
    });
  });
  return {
    formatted: allLines.join("\n"),
    mappings
  };
};
const parseCitations = (answerText, mappings) => {
  const citations = [];
  const citationPattern = /■(\d+)(?:-(\d+))?/g;
  let match;
  const processedRanges = /* @__PURE__ */ new Set();
  let cleanAnswer = answerText;
  let offsetAdjustment = 0;
  while ((match = citationPattern.exec(answerText)) !== null) {
    const fullMatch = match[0];
    const offset = match.index - offsetAdjustment;
    const startLine = parseInt(match[1], 10);
    const endLine = match[2] ? parseInt(match[2], 10) : startLine;
    const rangeKey = `${offset}:${startLine}-${endLine}`;
    if (processedRanges.has(rangeKey)) {
      continue;
    }
    processedRanges.add(rangeKey);
    const lineNumbers = [];
    for (let i = startLine; i <= endLine; i++) {
      lineNumbers.push(i);
    }
    const relevantMappings = mappings.filter((m) => lineNumbers.includes(m.lineNumber));
    if (relevantMappings.length > 0) {
      const documentMap = /* @__PURE__ */ new Map();
      relevantMappings.forEach((mapping) => {
        const existing = documentMap.get(mapping.document) || [];
        existing.push(mapping.text);
        documentMap.set(mapping.document, existing);
      });
      documentMap.forEach((lines, document) => {
        citations.push({
          offset,
          item: document,
          snippet: lines.join("\n")
        });
      });
    }
    cleanAnswer = cleanAnswer.slice(0, offset) + cleanAnswer.slice(offset + fullMatch.length);
    offsetAdjustment += fullMatch.length;
  }
  return {
    cleanAnswer: cleanAnswer.trim(),
    citations: citations.sort((a, b) => a.offset - b.offset)
  };
};
const processSingleChunk = async (formattedDocs, mappings, question, options, ctx) => {
  ctx.controller.signal.throwIfAborted();
  const result = await callLLM(formattedDocs, question, options, mappings, ctx);
  ctx.controller.signal.throwIfAborted();
  return result;
};
const processMultipleChunks = async (documents, question, chunkTokenLimit, options, ctx) => {
  const tokenizer = await getTokenizer();
  const chunks = [];
  let currentChunk = [];
  let currentTokens = 0;
  for (const doc of documents) {
    const docString = stringify(doc);
    const docTokens = tokenizer.count(docString);
    if (currentTokens + docTokens > chunkTokenLimit && currentChunk.length > 0) {
      chunks.push([...currentChunk]);
      currentChunk = [doc];
      currentTokens = docTokens;
    } else {
      currentChunk.push(doc);
      currentTokens += docTokens;
    }
  }
  if (currentChunk.length > 0) {
    chunks.push(currentChunk);
  }
  const processChunk = async (chunk) => {
    ctx.controller.signal.throwIfAborted();
    const { formatted, mappings } = formatDocumentsWithLineNumbers(chunk);
    const result = await callLLM(formatted, question, options, mappings, ctx);
    ctx.controller.signal.throwIfAborted();
    return result;
  };
  const limit = pLimit(10);
  const chunkResults = await Promise.all(chunks.map((chunk) => limit(() => processChunk(chunk))));
  ctx.controller.signal.throwIfAborted();
  return mergeChunkResults(chunkResults, question, chunkTokenLimit, options, ctx);
};
const mergeChunkResults = async (results, question, chunkTokenLimit, options, ctx) => {
  ctx.controller.signal.throwIfAborted();
  const answers = results.filter((r) => r.type === "answer");
  if (answers.length === 0) {
    const nonAnswer = results.find((r) => r.type !== "answer");
    return nonAnswer || { type: "missing_knowledge", reason: "No relevant information found in documents." };
  }
  if (answers.length === 1) {
    return answers[0];
  }
  const citedDocs = /* @__PURE__ */ new Set();
  answers.forEach((answer2) => {
    answer2.citations.forEach((citation) => {
      citedDocs.add(citation.item);
    });
  });
  const citedDocsArray = Array.from(citedDocs);
  const tokenizer = await getTokenizer();
  const { formatted, mappings } = formatDocumentsWithLineNumbers(citedDocsArray);
  if (tokenizer.count(formatted) <= chunkTokenLimit) {
    return await callLLM(formatted, question, options, mappings, ctx);
  }
  return answers.reduce((best, current) => current.citations.length > best.citations.length ? current : best);
};
const callLLM = async (formattedDocs, question, options, mappings, ctx) => {
  const systemPrompt = `You are an expert research assistant specialized in answering questions using only the information provided in documents.

# Task
Answer the user's question based ONLY on the information in the provided documents. You MUST cite your sources using line numbers.

# Document Format
Documents are provided with line numbers:
\u25A0001 | First line of text
\u25A0002 | Second line of text
\u25A0003 | Third line of text

# Citation Format
You MUST include citations immediately after statements. Use these formats:
- Single line: \u25A0035
- Range: \u25A0005-010
- Multiple: \u25A0035\u25A0046\u25A0094

# Response Format

Choose ONE of these response types:

**TYPE 1 - ANSWER** (Use this when you can answer the question)
\u25A0answer
[Your answer with inline citations\u25A0001-003. Make sure each part is cited correctly\u25A0013. More text. \u25A0015]
\u25A0end\u25A0

**TYPE 2 - AMBIGUOUS** (Use when the question has multiple valid interpretations)
\u25A0ambiguous
[Explain the ambiguity]
\u25A0follow_up
[Ask a clarifying question]
\u25A0answer
[First interpretation with citations \u25A0001 and part 2 as well.\u25A0002]
\u25A0answer
[Second interpretation with citations \u25A0005 and part 2 of the answer.\u25A0006]
\u25A0end\u25A0

**TYPE 3 - OUT OF TOPIC** (Use when question is completely unrelated to documents)
\u25A0out_of_topic
[Explain why it's unrelated]
\u25A0end\u25A0

**TYPE 4 - INVALID QUESTION** (Use when input is not a proper question, e.g., gibberish, malformed or nonsensical)
\u25A0invalid_question
[Explain why it's invalid, e.g., "The question is incomplete" or "The question contains nonsensical terms", or "Received gibberish"]
\u25A0end\u25A0

**TYPE 5 - MISSING KNOWLEDGE** (Use ONLY when documents lack specific details needed)
\u25A0missing_knowledge
[Explain what specific information is missing]
\u25A0end\u25A0

# Important Rules
- PREFER answering when possible - only use missing_knowledge if truly no relevant info exists
- ALWAYS cite sources with line numbers
- Use ONLY information from the documents
- Be precise and factual
- Do NOT fabricate information
- Do NOT mention "According to the documents" or similar phrases \u2013 just provide a high-quality answer with citations
- Do not be too strict on the question format; assume high-level answers are acceptable unless the question clearly asks for very specific details or requests depth beyond the documents

# Additional Instructions
Here are some additional instructions to follow about how to answer the question:
${options.instructions || "Provide a clear and concise answer based on the documents."}`;
  const userPrompt = `<documents>
${formattedDocs}
</documents>

Please answer the below question using the format specified above.
Question to answer: "${question}"`;
  const { extracted } = await ctx.generateContent({
    reasoningEffort: "none",
    systemPrompt,
    stopSequences: [END],
    messages: [
      {
        type: "text",
        role: "user",
        content: userPrompt
      }
    ],
    transform: (text) => {
      return parseResponse(text || "", mappings);
    }
  });
  return extracted;
};
const parseResponse = (response, mappings) => {
  const text = response.trim();
  if (text.includes(ANSWER_START)) {
    return parseAnswerResponse(text, mappings);
  } else if (text.includes(AMBIGUOUS_START)) {
    return parseAmbiguousResponse(text, mappings);
  } else if (text.includes(OUT_OF_TOPIC_START)) {
    return parseOutOfTopicResponse(text);
  } else if (text.includes(INVALID_QUESTION_START)) {
    return parseInvalidQuestionResponse(text);
  } else if (text.includes(MISSING_KNOWLEDGE_START)) {
    return parseMissingKnowledgeResponse(text);
  }
  return {
    type: "missing_knowledge",
    reason: "Unable to determine response type from the model output."
  };
};
const parseAnswerResponse = (text, mappings) => {
  const answerMatch = text.match(new RegExp(`${ANSWER_START}(.+)$`, "s"));
  if (!answerMatch) {
    return {
      type: "missing_knowledge",
      reason: "Could not extract answer from response."
    };
  }
  const answerText = answerMatch[1].trim();
  const { cleanAnswer, citations } = parseCitations(answerText, mappings);
  if (citations.length === 0) {
    throw new Error(
      "Answer must include citations using the \u25A0XXX format (e.g., \u25A0001, \u25A0001-005). Every statement should be backed by a citation to the source documents."
    );
  }
  return {
    type: "answer",
    answer: cleanAnswer,
    citations
  };
};
const parseAmbiguousResponse = (text, mappings) => {
  const ambiguityMatch = text.match(new RegExp(`${AMBIGUOUS_START}(.+?)\u25A0follow_up`, "s"));
  const ambiguity = ambiguityMatch ? ambiguityMatch[1].trim() : "The question has multiple interpretations.";
  const followUpMatch = text.match(/■follow_up(.+?)■answer/s);
  const follow_up = followUpMatch ? followUpMatch[1].trim() : "Please clarify your question.";
  const answerPattern = /■answer(.+?)(?=■answer|$)/gs;
  const answers = [];
  let match;
  while ((match = answerPattern.exec(text)) !== null) {
    const answerText = match[1].trim();
    const { cleanAnswer, citations } = parseCitations(answerText, mappings);
    answers.push({ answer: cleanAnswer, citations });
  }
  const answersWithoutCitations = answers.filter((a) => a.citations.length === 0);
  if (answersWithoutCitations.length > 0) {
    throw new Error(
      "Each answer in an ambiguous response must include citations using the \u25A0XXX format (e.g., \u25A0001, \u25A0001-005). Every statement should be backed by a citation to the source documents."
    );
  }
  return {
    type: "ambiguous",
    ambiguity,
    follow_up,
    answers: answers.length >= 2 ? answers.slice(0, 3) : answers
  };
};
const parseOutOfTopicResponse = (text) => {
  const reasonMatch = text.match(new RegExp(`${OUT_OF_TOPIC_START}(.+)$`, "s"));
  const reason = reasonMatch ? reasonMatch[1].trim() : "The question is not related to the provided documents.";
  return {
    type: "out_of_topic",
    reason
  };
};
const parseInvalidQuestionResponse = (text) => {
  const reasonMatch = text.match(new RegExp(`${INVALID_QUESTION_START}(.+)$`, "s"));
  const reason = reasonMatch ? reasonMatch[1].trim() : "The question is invalid or malformed.";
  return {
    type: "invalid_question",
    reason
  };
};
const parseMissingKnowledgeResponse = (text) => {
  const reasonMatch = text.match(new RegExp(`${MISSING_KNOWLEDGE_START}(.+)$`, "s"));
  const reason = reasonMatch ? reasonMatch[1].trim() : "The documents do not contain sufficient information to answer the question.";
  return {
    type: "missing_knowledge",
    reason
  };
};
const answer = async (documents, question, options, ctx) => {
  ctx.controller.signal.throwIfAborted();
  if (documents.length === 0) {
    return {
      type: "missing_knowledge",
      reason: "No documents provided to answer the question."
    };
  }
  if (!question || question.trim().length === 0) {
    return {
      type: "invalid_question",
      reason: "The question is empty or contains no content."
    };
  }
  const tokenizer = await getTokenizer();
  const model = await ctx.getModel();
  const TOTAL_MAX_TOKENS = Math.min(options.chunkLength, model.input.maxTokens - PROMPT_INPUT_BUFFER);
  const { formatted: allFormattedDocs, mappings: allMappings } = formatDocumentsWithLineNumbers(documents);
  const CHUNK_DOC_TOKENS = Math.floor(TOTAL_MAX_TOKENS * 0.6);
  const totalDocTokens = tokenizer.count(allFormattedDocs);
  if (totalDocTokens <= CHUNK_DOC_TOKENS) {
    return await processSingleChunk(allFormattedDocs, allMappings, question, options, ctx);
  } else {
    return await processMultipleChunks(documents, question, CHUNK_DOC_TOKENS, options, ctx);
  }
};
Zai.prototype.answer = function(documents, question, _options) {
  const parse = _Options.safeParse(_options ?? {});
  if (!parse.success) {
    return Promise.reject(new Error(`Invalid options: ${parse.error.message}`));
  }
  const context = new ZaiContext({
    client: this.client,
    modelId: this.Model,
    taskId: this.taskId,
    taskType: "zai.answer",
    adapter: this.adapter
  });
  return new Response(
    context,
    answer(documents, question, parse.data, context),
    (result) => result
    // No simplification - return full result
  );
};
