import { post, rawSSE, raw, save, short } from "./lib.ts";

const results: Record<string, any> = {};
const q = { role: "user", content: "Is 221 a prime number? Answer briefly." };

// Reasoning non-stream full shape on both hybrid models + glm
for (const model of ["magistral-medium-latest", "mistral-small-latest", "glm-5-2"]) {
  const r = await post("/chat/completions", { model, messages: [q], max_tokens: 400, reasoning_effort: "high" });
  const b: any = r.body;
  const msg = b?.choices?.[0]?.message;
  console.log(`\n=== ${model} reasoning_effort=high ${r.status} finish=${b?.choices?.[0]?.finish_reason} usage=${JSON.stringify(b?.usage)}`);
  console.log("content", short(msg?.content, 1200));
  results[`${model}_high`] = r;
}

// Streaming shape with reasoning on medium: record the sequence of delta.content shapes and the raw first thinking chunk
const s = await rawSSE("/chat/completions", { model: "magistral-medium-latest", messages: [q], max_tokens: 400, reasoning_effort: "high", stream: true });
const seq: string[] = [];
let firstThink: any, transition: any, firstText: any;
for (const e of s.events) {
  if (e.data === "[DONE]") continue;
  const c = e.data?.choices?.[0]?.delta?.content;
  const shape = c === undefined ? "undef" : typeof c === "string" ? "string" : Array.isArray(c) ? "array:" + c.map((x: any) => x.type).join(",") : typeof c;
  if (seq[seq.length - 1] !== shape) seq.push(shape);
  if (!firstThink && Array.isArray(c) && c.some((x: any) => x.type === "thinking")) firstThink = e.data;
  if (!transition && Array.isArray(c) && c.length > 1) transition = e.data;
  if (!firstText && typeof c === "string" && c.length && firstThink) firstText = e.data;
}
console.log("\nstream shapes seq", seq, "events", s.events.length);
console.log("first thinking chunk", short(firstThink, 700));
console.log("transition chunk", short(transition, 700));
console.log("first text chunk after thinking", short(firstText, 400));
console.log("last", short(s.events[s.events.length - 2]?.data, 500));
results.streamReasoning = s;

// Multi-turn replay including the thinking chunk
const first: any = results["magistral-medium-latest_high"].body;
if (first?.choices?.[0]?.message) {
  const replay = await post("/chat/completions", {
    model: "magistral-medium-latest",
    reasoning_effort: "high",
    max_tokens: 200,
    messages: [q, first.choices[0].message, { role: "user", content: "And 223?" }],
  });
  console.log("\nreplay with thinking chunk", replay.status, short((replay.body as any)?.choices?.[0]?.message?.content ?? replay.body, 300));
  results.replay = replay;
  // replay with thinking stripped to text only
  const textOnly = Array.isArray(first.choices[0].message.content)
    ? first.choices[0].message.content.filter((c: any) => c.type === "text").map((c: any) => c.text).join("")
    : first.choices[0].message.content;
  const replay2 = await post("/chat/completions", {
    model: "magistral-medium-latest",
    reasoning_effort: "high",
    max_tokens: 200,
    messages: [q, { role: "assistant", content: textOnly }, { role: "user", content: "And 223?" }],
  });
  console.log("replay text-only", replay2.status, short((replay2.body as any)?.choices?.[0]?.message?.content ?? replay2.body, 200));
  results.replayTextOnly = replay2;
}

// prompt_mode reasoning on small (legacy Magistral param) — shape
const pm = await post("/chat/completions", { model: "mistral-small-latest", messages: [q], max_tokens: 300, prompt_mode: "reasoning" });
console.log("\nprompt_mode reasoning (small)", pm.status, JSON.stringify((pm.body as any)?.usage), short((pm.body as any)?.choices?.[0]?.message?.content ?? pm.body, 500));
results.promptMode = pm;

// max_tokens truncation during thinking -> finish_reason?
const trunc = await post("/chat/completions", { model: "magistral-medium-latest", messages: [q], max_tokens: 20, reasoning_effort: "high" });
console.log("truncated during thinking", trunc.status, "finish", (trunc.body as any)?.choices?.[0]?.finish_reason, JSON.stringify((trunc.body as any)?.usage), short((trunc.body as any)?.choices?.[0]?.message?.content, 300));
results.trunc = trunc;

// Labs + third-party quick chat
for (const model of ["labs-leanstral-1-5", "ministral-3b-latest", "ministral-14b-latest", "voxtral-small-latest"]) {
  const r = await post("/chat/completions", { model, messages: [{ role: "user", content: "Say OK." }], max_tokens: 20 });
  console.log(`\n${model} ${r.status} ${short((r.body as any)?.choices?.[0]?.message?.content ?? r.body, 200)} usage=${JSON.stringify((r.body as any)?.usage)} rl=${(r.headers as any)["x-ratelimit-limit-req-minute"]}/${(r.headers as any)["x-ratelimit-limit-tokens-minute"]}`);
  results[`tiny_${model}`] = r;
}

// FIM on codestral
const fim = await post("/fim/completions", { model: "codestral-latest", prompt: "def fib(n):\n", suffix: "\n    return fib(n-1) + fib(n-2)", max_tokens: 40 });
console.log("\nFIM codestral", fim.status, short((fim.body as any)?.choices?.[0]?.message?.content ?? fim.body, 200));
results.fim = fim;
const fimSmall = await post("/fim/completions", { model: "mistral-small-latest", prompt: "def fib(n):\n", max_tokens: 10 });
console.log("FIM on small", fimSmall.status, short(fimSmall.body, 200));
results.fimSmall = fimSmall;

// web_search tool on chat completions (docs list WebSearchTool in tools[])
const ws = await post("/chat/completions", { model: "mistral-medium-latest", messages: [{ role: "user", content: "What is today's date and one headline from Le Devoir today? Cite." }], tools: [{ type: "web_search" }], max_tokens: 200 });
console.log("\nweb_search on chat completions", ws.status, JSON.stringify((ws.body as any)?.usage), short(ws.body, 1200));
results.webSearchChat = ws;

// Conversations API: one call with web_search, store false
const conv = await post("/conversations", {
  model: "mistral-medium-latest",
  inputs: "In one sentence, what is the latest stable Node.js version? Cite a source.",
  tools: [{ type: "web_search" }],
  store: false,
  completion_args: { max_tokens: 200 },
});
console.log("\nconversations web_search", conv.status, short(conv.body, 1500));
results.conversation = conv;
// Conversations stream, no tools
const cs = await rawSSE("/conversations", { model: "mistral-small-latest", inputs: "Say hello in 3 words.", store: false, stream: true, completion_args: { max_tokens: 30 } });
console.log("\nconversations stream", cs.status, "events", cs.events.length, "types", [...new Set(cs.events.map((e) => e.event ?? e.data?.type))]);
console.log("first", short(cs.events[0], 400), "\nlast", short(cs.events[cs.events.length - 1], 500));
results.conversationStream = cs;

// Token counting / tokenize endpoints?
for (const p of ["/tokenize", "/chat/tokenize", "/models/mistral-small-latest/tokenize"]) {
  const t = await post(p, { model: "mistral-small-latest", messages: [{ role: "user", content: "Hello" }] });
  console.log("tokenize", p, t.status, short(t.body, 150));
}
const usage = await raw("/usage");
console.log("GET /usage", usage.status, short(usage.body, 150));
save("09-reasoning-misc.json", results);