import { z, transforms } from "@bpinternal/zui";
import JSON5 from "json5";
import { jsonrepair } from "jsonrepair";
import { chunk, isArray } from "lodash-es";
import pLimit from "p-limit";
import { ZaiContext } from "../context";
import { Response } from "../response";
import { getTokenizer } from "../tokenizer";
import { fastHash, stringify, takeUntilTokens } from "../utils";
import { Zai } from "../zai";
import { PROMPT_INPUT_BUFFER } from "./constants";
import { JsonParsingError } from "./errors";
const Options = z.object({
  instructions: z.string().optional().describe("Instructions to guide the user on how to extract the data"),
  chunkLength: z.number().min(100).max(1e5).optional().describe("The maximum number of tokens per chunk").default(16e3),
  strict: z.boolean().optional().default(true).describe("Whether to strictly follow the schema or not")
});
const SPECIAL_CHAR = "\u25A0";
const START = "\u25A0json_start\u25A0";
const END = "\u25A0json_end\u25A0";
const NO_MORE = "\u25A0NO_MORE_ELEMENT\u25A0";
const ZERO_ELEMENTS = "\u25A0ZERO_ELEMENTS\u25A0";
const extract = async (input, _schema, _options, ctx) => {
  ctx.controller.signal.throwIfAborted();
  let schema;
  try {
    schema = transforms.fromJSONSchema(transforms.toJSONSchema(_schema));
  } catch {
    schema = _schema;
  }
  const options = Options.parse(_options ?? {});
  const tokenizer = await getTokenizer();
  const model = await ctx.getModel();
  const taskId = ctx.taskId;
  const taskType = "zai.extract";
  const PROMPT_COMPONENT = Math.max(model.input.maxTokens - PROMPT_INPUT_BUFFER, 100);
  let isArrayOfObjects = false;
  let wrappedValue = false;
  const originalSchema = schema;
  const baseType = (schema.naked ? schema.naked() : schema)?.constructor?.name ?? "unknown";
  if (baseType === "ZodArray") {
    isArrayOfObjects = true;
    let elementType = schema.element;
    if (elementType.naked) {
      elementType = elementType.naked();
    }
    if (elementType?.constructor?.name === "ZodObject") {
      schema = elementType;
    } else {
      wrappedValue = true;
      schema = z.object({
        value: elementType
      });
    }
  } else if (baseType !== "ZodObject") {
    wrappedValue = true;
    schema = z.object({
      value: originalSchema
    });
  }
  if (!options.strict) {
    try {
      schema = schema.partial();
    } catch {
    }
  }
  const schemaTypescript = schema.toTypescriptType({ declaration: false, treatDefaultAsOptional: true });
  const schemaLength = tokenizer.count(schemaTypescript);
  options.chunkLength = Math.min(options.chunkLength, model.input.maxTokens - PROMPT_INPUT_BUFFER - schemaLength);
  const keys = Object.keys(schema.shape);
  const inputAsString = stringify(input);
  if (tokenizer.count(inputAsString) > options.chunkLength) {
    const limit = pLimit(10);
    const tokens = tokenizer.split(inputAsString);
    const chunks = chunk(tokens, options.chunkLength).map((x) => x.join(""));
    const all = await Promise.allSettled(
      chunks.map(
        (chunk2) => limit(
          () => extract(
            chunk2,
            originalSchema,
            {
              ...options,
              strict: false
              // We don't want to fail on strict mode for sub-chunks
            },
            ctx
          )
        )
      )
    ).then(
      (results) => results.filter((x) => x.status === "fulfilled").map((x) => x.value)
    );
    ctx.controller.signal.throwIfAborted();
    const rows = all.map((x, idx) => `<part-${idx + 1}>
${stringify(x, true)}
</part-${idx + 1}>`).join("\n");
    return extract(
      `
The result has been split into ${all.length} parts. Recursively merge the result into the final result.
When merging arrays, take unique values.
When merging conflictual (but defined) information, take the most reasonable and frequent value.
Non-defined values are OK and normal. Don't delete fields because of null values. Focus on defined values.

Here's the data:
${rows}

Merge it back into a final result.`.trim(),
      originalSchema,
      options,
      ctx
    );
  }
  const instructions = [];
  if (options.instructions) {
    instructions.push(options.instructions);
  }
  const shape = `{ ${keys.map((key) => `"${key}": ...`).join(", ")} }`;
  const abbv = "{ ... }";
  if (isArrayOfObjects) {
    instructions.push("You may have multiple elements, or zero elements in the input.");
    instructions.push("You must extract each element separately.");
    instructions.push(`Each element must be a JSON object with exactly the format: ${START}${shape}${END}`);
    instructions.push(`If there are no elements to extract, respond with ${ZERO_ELEMENTS}.`);
    instructions.push(`When you are done extracting all elements, type "${NO_MORE}" to finish.`);
    instructions.push(
      `For example, if you have zero elements, the output should look like this: ${ZERO_ELEMENTS}${NO_MORE}`
    );
    instructions.push(
      `For example, if you have two elements, the output should look like this: ${START}${abbv}${END}${START}${abbv}${END}${NO_MORE}`
    );
  } else {
    instructions.push("You may have exactly one element in the input.");
    instructions.push(`The element must be a JSON object with exactly the format: ${START}${shape}${END}`);
  }
  if (!options.strict) {
    instructions.push("You may ignore any fields that are not present in the input. All keys are optional.");
  }
  const EXAMPLES_TOKENS = PROMPT_COMPONENT - tokenizer.count(inputAsString) - tokenizer.count(instructions.join("\n"));
  const Key = fastHash(
    JSON.stringify({
      taskType,
      taskId,
      input: inputAsString,
      instructions: options.instructions
    })
  );
  const examples = taskId && ctx.adapter ? await ctx.adapter.getExamples({
    input: inputAsString,
    taskType,
    taskId
  }) : [];
  const exactMatch = examples.find((x) => x.key === Key);
  if (exactMatch) {
    return exactMatch.output;
  }
  const defaultExample = isArrayOfObjects ? {
    input: `The story goes as follow.
Once upon a time, there was a person named Alice who was 30 years old.
Then, there was a person named Bob who was 25 years old.
The end.`,
    schema: "Array<{ name: string, age: number }>",
    instructions: "Extract all people",
    extracted: [
      {
        name: "Alice",
        age: 30
      },
      {
        name: "Bob",
        age: 25
      }
    ]
  } : {
    input: `The story goes as follow.
Once upon a time, there was a person named Alice who was 30 years old.
The end.`,
    schema: "{ name: string, age: number }",
    instructions: "Extract the person",
    extracted: { name: "Alice", age: 30 }
  };
  const userExamples = examples.map((e) => ({
    input: e.input,
    extracted: e.output,
    schema: schemaTypescript,
    instructions: options.instructions
  }));
  let exampleId = 1;
  const formatInput = (input2, schema2, instructions2) => {
    const header = userExamples.length ? `Expert Example #${exampleId++}` : "Here's an example to help you understand the format:";
    return `
${header}

<|start_schema|>
${schema2}
<|end_schema|>

<|start_instructions|>
${instructions2 ?? "No specific instructions, just follow the schema above."}
<|end_instructions|>

<|start_input|>
${input2.trim()}
<|end_input|>
  `.trim();
  };
  const formatOutput = (extracted2) => {
    extracted2 = isArray(extracted2) ? extracted2 : [extracted2];
    return extracted2.map(
      (x) => `
${START}
${JSON.stringify(x, null, 2)}
${END}`.trim()
    ).join("\n") + NO_MORE;
  };
  const formatExample = (example) => [
    {
      type: "text",
      content: formatInput(stringify(example.input ?? null), example.schema, example.instructions),
      role: "user"
    },
    {
      type: "text",
      content: formatOutput(example.extracted),
      role: "assistant"
    }
  ];
  const allExamples = takeUntilTokens(
    userExamples.length ? userExamples : [defaultExample],
    EXAMPLES_TOKENS,
    (el) => tokenizer.count(stringify(el.input)) + tokenizer.count(stringify(el.extracted))
  ).map(formatExample).flat();
  const { meta, extracted } = await ctx.generateContent({
    systemPrompt: `
Extract the following information from the input:
${schemaTypescript}
====

${instructions.map((x) => `\u2022 ${x}`).join("\n")}
`.trim(),
    stopSequences: [isArrayOfObjects ? NO_MORE : END],
    messages: [
      ...allExamples,
      {
        role: "user",
        type: "text",
        content: formatInput(inputAsString, schemaTypescript, options.instructions ?? "")
      }
    ],
    transform: (text) => (text || "{}")?.split(START).filter((x) => x.trim().length > 0 && x.includes("}")).map((x) => {
      try {
        let json = x.slice(0, x.indexOf(END)).trim();
        if (json.includes(SPECIAL_CHAR)) {
          json = json.slice(0, json.indexOf(SPECIAL_CHAR)).trim();
        }
        const repairedJson = jsonrepair(json);
        const parsedJson = JSON5.parse(repairedJson);
        const safe = schema.safeParse(parsedJson);
        if (safe.success) {
          return safe.data;
        }
        if (options.strict) {
          throw new JsonParsingError(x, safe.error);
        }
        return parsedJson;
      } catch (error) {
        throw new JsonParsingError(x, error instanceof Error ? error : new Error("Unknown error"));
      }
    }).filter((x) => x !== null)
  });
  let final;
  if (isArrayOfObjects) {
    final = extracted;
  } else if (extracted.length === 0) {
    final = options.strict ? schema.parse({}) : {};
  } else {
    final = extracted[0];
  }
  if (wrappedValue) {
    if (Array.isArray(final)) {
      final = final.map((x) => "value" in x ? x.value : x);
    } else {
      final = "value" in final ? final.value : final;
    }
  }
  if (taskId && ctx.adapter && !ctx.controller.signal.aborted) {
    await ctx.adapter.saveExample({
      key: Key,
      taskId: `zai/${taskId}`,
      taskType,
      instructions: options.instructions ?? "No specific instructions",
      input: inputAsString,
      output: final,
      metadata: {
        cost: {
          input: meta.cost.input,
          output: meta.cost.output
        },
        latency: meta.latency,
        model: ctx.modelId,
        tokens: {
          input: meta.tokens.input,
          output: meta.tokens.output
        }
      }
    });
  }
  return final;
};
Zai.prototype.extract = function(input, schema, _options) {
  const context = new ZaiContext({
    client: this.client,
    modelId: this.Model,
    taskId: this.taskId,
    taskType: "zai.extract",
    adapter: this.adapter
  });
  return new Response(context, extract(input, schema, _options, context), (result) => result);
};
