/**
 * The prompted tool-calling protocol.
 *
 * This is the part that gives tool-calling to a model that has no native
 * support for it: the schemas are injected into the system prompt, the model is
 * asked to emit a `<tool_call>` block, and we parse that back out ourselves.
 *
 * Every rule below is a response to behaviour actually observed on 2-7B models
 * served through HF Inference Providers, not a guess.
 */

/* ------------------------------------------------------------------ *
 * Prompt assembly
 * ------------------------------------------------------------------ */

export const CALL_OPEN = "<tool_call>";
export const CALL_CLOSE = "</tool_call>";

/** Stop sequence. Without it, small models emit a call and then cheerfully
 *  hallucinate its result in the same completion. */
export const STOP = [CALL_CLOSE];

export function buildSystemPrompt(schemas) {
  const fns = schemas.map((s) => s.function);
  return `You are a financial assistant with access to live market data tools.

You cannot know prices from memory. You MUST use the tools for any price,
market-cap, history or currency question.

AVAILABLE TOOLS (JSON Schema):
${JSON.stringify(fns, null, 2)}

HOW TO CALL A TOOL
Your reply must BEGIN immediately with a tool call, in exactly this form:
${CALL_OPEN}{"name": "<tool_name>", "arguments": {<args>}}${CALL_CLOSE}

RULES
1. Write NO text before the ${CALL_OPEN} tag. No preamble, no explanation, no
   restating of the question. The tag is the first thing in your reply.
2. NEVER write a tool result yourself. Stop after the closing tag and wait.
3. The result arrives as a message beginning with TOOL_RESULT.
4. To call several tools at once, emit several ${CALL_OPEN} blocks back to back.
5. Use ONLY the tool names listed above, with exactly those argument names.
6. Once the TOOL_RESULT messages give you everything you need, reply in plain
   prose with NO tool call, and state the actual numbers.
7. If a TOOL_RESULT contains "error", explain the problem or retry with
   corrected arguments.`;
}

/**
 * Few-shot examples as REAL message turns rather than dialogue text inside the
 * system prompt.
 *
 * This is not a stylistic choice. With the examples written as prose in the
 * system prompt, Qwen3.5-2B imitated the example's leading sentence ("I need the
 * live price of bitcoin.") and then emitted EOS, never reaching the tool call.
 * Delivering the pattern as genuine turns runs it through the model's own chat
 * template, and the assistant turns contain nothing but the call — so the
 * behaviour being demonstrated is exactly the behaviour we want.
 *
 * Kept to two exchanges: enough to fix the format, small enough not to crowd a
 * 2B context window.
 */
export function buildFewShot() {
  return [
    { role: "user", content: "What is Bitcoin worth?" },
    {
      role: "assistant",
      content: `${CALL_OPEN}{"name": "get_price", "arguments": {"coin_id": "bitcoin"}}${CALL_CLOSE}`,
    },
    {
      role: "user",
      content:
        'TOOL_RESULT get_price {"coin_id": "bitcoin", "vs_currency": "USD", "price": 64779, "change_24h_pct": 0.79}',
    },
    {
      role: "assistant",
      content: "Bitcoin is trading at $64,779, up 0.79% over the last 24 hours.",
    },
    { role: "user", content: "Is Ethereum pricier than Solana? Show Ethereum in euros too." },
    {
      role: "assistant",
      content:
        `${CALL_OPEN}{"name": "get_price", "arguments": {"coin_id": "ethereum"}}${CALL_CLOSE}\n` +
        `${CALL_OPEN}{"name": "get_price", "arguments": {"coin_id": "solana"}}${CALL_CLOSE}`,
    },
    {
      role: "user",
      content:
        'TOOL_RESULT get_price {"coin_id": "ethereum", "price": 1919.55, "vs_currency": "USD"}\n' +
        'TOOL_RESULT get_price {"coin_id": "solana", "price": 138.2, "vs_currency": "USD"}',
    },
    {
      role: "assistant",
      content: `${CALL_OPEN}{"name": "convert_currency", "arguments": {"amount": 1919.55, "from_currency": "USD", "to_currency": "EUR"}}${CALL_CLOSE}`,
    },
    {
      role: "user",
      content: 'TOOL_RESULT convert_currency {"result": 1775.58, "to_currency": "EUR", "rate": 0.925}',
    },
    {
      role: "assistant",
      content:
        "Ethereum ($1,919.55) is far pricier than Solana ($138.20). In euros, Ethereum is about €1,775.58.",
    },
  ];
}

/** Format a tool result as the user-role turn we feed back to the model.
 *
 *  Deliberately NOT a `role: "tool"` message: templates without native tool
 *  support have no slot for that role. gemma-2 for instance raises on it and
 *  demands strictly alternating user/assistant turns, which this satisfies.
 */
export function formatToolResult(name, result) {
  return `TOOL_RESULT ${name} ${JSON.stringify(result)}`;
}

/* ------------------------------------------------------------------ *
 * Reasoning blocks
 * ------------------------------------------------------------------ */

/**
 * Split hybrid-reasoning output into { reasoning, text }.
 *
 * Qwen3.5 emits chain-of-thought terminated by `</think>` -- often WITHOUT a
 * matching opening tag, so a naive paired regex misses it entirely and the
 * reasoning leaks into the user-visible answer.
 */
export function splitReasoning(raw) {
  let text = raw || "";
  let reasoning = "";

  const paired = text.match(/<think>([\s\S]*?)<\/think>/g);
  if (paired) {
    for (const block of paired) {
      reasoning += block.replace(/<\/?think>/g, "");
    }
    text = text.replace(/<think>[\s\S]*?<\/think>/g, "");
  }

  // Orphan close tag: everything before it was reasoning.
  const orphan = text.indexOf("</think>");
  if (orphan !== -1) {
    reasoning += text.slice(0, orphan);
    text = text.slice(orphan + "</think>".length);
  }

  // Unterminated open tag: the model was cut off mid-thought.
  const openOnly = text.indexOf("<think>");
  if (openOnly !== -1) {
    reasoning += text.slice(openOnly + "<think>".length);
    text = text.slice(0, openOnly);
  }

  return { reasoning: reasoning.trim(), text: text.trim() };
}

/* ------------------------------------------------------------------ *
 * Tolerant JSON parsing
 * ------------------------------------------------------------------ */

/** Repair the malformations small models actually produce. */
function repairJSON(s) {
  let t = s.trim();
  t = t.replace(/,\s*([}\]])/g, "$1"); // trailing commas
  t = t.replace(/([{,]\s*)'([^']+?)'(\s*:)/g, '$1"$2"$3'); // 'key':
  t = t.replace(/(:\s*)'([^']*?)'/g, '$1"$2"'); // : 'value'
  t = t.replace(/([{,]\s*)([A-Za-z_]\w*)(\s*:)/g, '$1"$2"$3'); // bare key:
  t = t.replace(/\bNone\b/g, "null").replace(/\bTrue\b/g, "true").replace(/\bFalse\b/g, "false");
  return t;
}

function tryParse(s) {
  try {
    return JSON.parse(s);
  } catch {
    try {
      return JSON.parse(repairJSON(s));
    } catch {
      return undefined;
    }
  }
}

/** Extract balanced {...} substrings, respecting strings and escapes. */
function balancedObjects(s) {
  const out = [];
  let depth = 0,
    start = -1,
    inStr = false,
    esc = false;
  for (let i = 0; i < s.length; i++) {
    const c = s[i];
    if (inStr) {
      if (esc) esc = false;
      else if (c === "\\") esc = true;
      else if (c === '"') inStr = false;
      continue;
    }
    if (c === '"') inStr = true;
    else if (c === "{") {
      if (depth === 0) start = i;
      depth++;
    } else if (c === "}") {
      depth--;
      if (depth === 0 && start !== -1) {
        out.push(s.slice(start, i + 1));
        start = -1;
      }
    }
  }
  return out;
}

/** Normalise a parsed object into { name, arguments } or null. */
function normaliseCall(obj) {
  if (!obj || typeof obj !== "object") return null;

  // Some models nest under "tool_call" / "function".
  if (obj.tool_call && typeof obj.tool_call === "object") obj = obj.tool_call;
  if (obj.function && typeof obj.function === "object" && !obj.name) obj = obj.function;

  const name = obj.name || obj.tool || obj.tool_name || obj.function_name;
  if (typeof name !== "string" || !name) return null;

  let args = obj.arguments ?? obj.args ?? obj.parameters ?? obj.params ?? obj.input ?? {};
  if (typeof args === "string") args = tryParse(args) ?? {};
  if (typeof args !== "object" || args === null || Array.isArray(args)) args = {};

  return { name, arguments: args };
}

/**
 * Coerce argument types to match the schema.
 *
 * Small models routinely send `"days": "7"` or `"amount": "1919.55"`. Without
 * this, arithmetic silently turns into string concatenation.
 */
export function coerceArgs(name, args, schemas) {
  const schema = schemas.find((s) => s.function.name === name);
  if (!schema) return args;
  const props = schema.function.parameters?.properties || {};
  const out = {};
  for (const [k, v] of Object.entries(args)) {
    const want = props[k]?.type;
    if (want === "number" && typeof v === "string" && v.trim() !== "" && isFinite(Number(v))) {
      out[k] = Number(v);
    } else if (want === "integer" && typeof v === "string" && /^-?\d+$/.test(v.trim())) {
      out[k] = parseInt(v, 10);
    } else if (want === "string" && typeof v === "number") {
      out[k] = String(v);
    } else {
      out[k] = v;
    }
  }
  return out;
}

/**
 * Parse a raw completion into { reasoning, calls, text }.
 *
 * Scans, in order of trust: explicit <tool_call> tags, ```json fences, then any
 * balanced object that looks like a call. Deduplicates, because a model that
 * both fences AND tags the same call would otherwise fire it twice.
 */
export function parseToolCalls(raw, schemas = []) {
  const { reasoning, text } = splitReasoning(raw);

  const candidates = [];
  let consumed = text;

  // 1. Tagged blocks. The stop sequence strips the closing tag, so accept an
  //    unterminated trailing block too.
  const tagged = /<tool_call>\s*([\s\S]*?)(?:<\/tool_call>|$)/g;
  let m;
  while ((m = tagged.exec(text)) !== null) {
    candidates.push(m[1]);
    consumed = consumed.replace(m[0], " ");
  }

  // 2. Fenced code blocks.
  const fenced = /```(?:json|tool_call)?\s*([\s\S]*?)```/g;
  while ((m = fenced.exec(text)) !== null) {
    candidates.push(m[1]);
    consumed = consumed.replace(m[0], " ");
  }

  // 3. Bare objects in whatever prose is left.
  for (const obj of balancedObjects(consumed)) {
    if (/"?(name|tool|function_name)"?\s*:/.test(obj)) candidates.push(obj);
  }

  const calls = [];
  const seen = new Set();
  for (const cand of candidates) {
    // A candidate may itself contain several concatenated objects.
    const pieces = balancedObjects(cand);
    const chunks = pieces.length ? pieces : [cand];
    for (const chunk of chunks) {
      const call = normaliseCall(tryParse(chunk));
      if (!call) continue;
      call.arguments = coerceArgs(call.name, call.arguments, schemas);
      const key = `${call.name}:${JSON.stringify(call.arguments)}`;
      if (seen.has(key)) continue;
      seen.add(key);
      calls.push(call);
    }
  }

  // Prose the user should see: reasoning and call blocks removed.
  const visible = consumed.replace(/\s+/g, " ").trim();

  return { reasoning, calls, text: visible };
}

/* ------------------------------------------------------------------ *
 * Template capability shims
 * ------------------------------------------------------------------ */

/**
 * Some chat templates reject a system role outright (gemma-2 raises
 * "System role not supported"). Fold the system prompt into the first user
 * message instead.
 */
export function foldSystemIntoUser(messages) {
  if (!messages.length || messages[0].role !== "system") return messages;
  const [sys, ...rest] = messages;
  const firstUser = rest.findIndex((m) => m.role === "user");
  if (firstUser === -1) return [{ role: "user", content: sys.content }, ...rest];
  const copy = [...rest];
  copy[firstUser] = {
    role: "user",
    content: `${sys.content}\n\n---\n\n${copy[firstUser].content}`,
  };
  return copy;
}

/** Merge adjacent same-role messages, for templates that demand strict
 *  user/assistant alternation (gemma-2 again). */
export function mergeAlternating(messages) {
  const out = [];
  for (const m of messages) {
    const last = out[out.length - 1];
    if (last && last.role === m.role) {
      last.content = `${last.content}\n${m.content}`;
    } else {
      out.push({ ...m });
    }
  }
  return out;
}