// Probe 03: per model — (d) function-calling round trip with streaming, (e) structured output via responseJsonSchema.
// Plus (h) googleSearch grounding and code execution on selected models.
import { ai, MODELS, SUFFIX, save, errInfo, shortText, withRetry, sleep, PACE_MS } from "./lib.js";

const results: any = {};
const weatherTool = {
  functionDeclarations: [{
    name: "get_weather",
    description: "Get the current weather for a city.",
    parametersJsonSchema: { type: "object", properties: { city: { type: "string" }, unit: { type: "string", enum: ["C", "F"] } }, required: ["city"] },
  }],
};

for (const model of MODELS) {
  const r: any = (results[model] = {});

  // (d) function calling, streaming, then send functionResponse
  try {
    const contents: any[] = [{ role: "user", parts: [{ text: "What is the weather in Montreal right now? Use the tool." }] }];
    const stream = await withRetry(() => ai.models.generateContentStream({ model, contents, config: { tools: [weatherTool], maxOutputTokens: 2000 } }));
    const chunks: any[] = [];
    let modelParts: any[] = [];
    for await (const c of stream) {
      const parts = c.candidates?.[0]?.content?.parts ?? [];
      chunks.push({ parts: parts.map((p: any) => ({ keys: Object.keys(p), fc: p.functionCall, thought: p.thought, sig: p.thoughtSignature ? `${p.thoughtSignature.length} chars` : undefined })), finishReason: c.candidates?.[0]?.finishReason, usage: c.usageMetadata ? Object.keys(c.usageMetadata) : null });
      modelParts.push(...parts);
    }
    const fcParts = modelParts.filter((p) => p.functionCall);
    r.functionCall = { ok: true, chunkCount: chunks.length, chunks, functionCalls: fcParts.map((p) => p.functionCall), sigOnFcPart: fcParts.map((p) => !!p.thoughtSignature) };
    if (fcParts.length) {
      // Round trip: append model content (with thoughtSignature preserved) and functionResponse parts
      contents.push({ role: "model", parts: modelParts });
      contents.push({
        role: "user",
        parts: fcParts.map((p) => ({ functionResponse: { name: p.functionCall.name, id: p.functionCall.id, response: { temperatureC: 21, condition: "sunny" } } })),
      });
      await sleep(PACE_MS);
      const res2 = await withRetry(() => ai.models.generateContent({ model, contents, config: { tools: [weatherTool], maxOutputTokens: 2000 } }));
      r.functionCall.roundTrip = { text: shortText(res2), finishReason: res2.candidates?.[0]?.finishReason, usage: res2.usageMetadata };
    }
  } catch (e) {
    r.functionCall = { ok: false, error: errInfo(e) };
  }

  // (e) structured output responseJsonSchema
  try {
    await sleep(PACE_MS);
    const res = await withRetry(() => ai.models.generateContent({
      model,
      contents: "Extract: 'Alice is 30 and lives in Paris.'",
      config: {
        maxOutputTokens: 2000,
        responseMimeType: "application/json",
        responseJsonSchema: { type: "object", properties: { name: { type: "string" }, age: { type: "integer" }, city: { type: "string" } }, required: ["name", "age", "city"], additionalProperties: false },
      },
    }));
    const txt = shortText(res);
    let parsed: any = null; try { parsed = JSON.parse(res.candidates?.[0]?.content?.parts?.filter((p: any) => !p.thought).map((p: any) => p.text).join("") ?? ""); } catch {}
    r.structured = { ok: true, text: txt, parsedOk: !!parsed, parsed, finishReason: res.candidates?.[0]?.finishReason };
  } catch (e) {
    r.structured = { ok: false, error: errInfo(e) };
  }
  await sleep(PACE_MS);
  console.log(model, "fc:", r.functionCall.ok ? `${r.functionCall.functionCalls.length} call(s) chunks=${r.functionCall.chunkCount} rt="${r.functionCall.roundTrip?.text?.slice(0, 50)}"` : r.functionCall.error.message.slice(0, 200), "| structured:", r.structured.ok ? r.structured.text.slice(0, 80) : r.structured.error.message.slice(0, 200));
}

const EXTRAS = process.argv[3] === "extras";
// (h) Google Search grounding (3.x models; 2.5 is 404 for this key)
for (const model of EXTRAS ? [MODELS[0]] : []) {
  await sleep(PACE_MS);
  try {
    const res = await withRetry(() => ai.models.generateContent({ model, contents: "Who won the most recent FIFA World Cup final and what was the score? One sentence.", config: { tools: [{ googleSearch: {} }], maxOutputTokens: 3000 } }));
    const gm: any = res.candidates?.[0]?.groundingMetadata;
    results[`search:${model}`] = {
      ok: true, text: shortText(res),
      groundingMetadataKeys: gm ? Object.keys(gm) : null,
      webSearchQueries: gm?.webSearchQueries,
      groundingChunksSample: gm?.groundingChunks?.slice(0, 2),
      groundingSupportsSample: gm?.groundingSupports?.slice(0, 2),
      searchEntryPointKeys: gm?.searchEntryPoint ? Object.keys(gm.searchEntryPoint) : null,
      renderedContentLen: gm?.searchEntryPoint?.renderedContent?.length,
      usage: res.usageMetadata,
      partKeys: res.candidates?.[0]?.content?.parts?.map((p: any) => Object.keys(p)),
    };
  } catch (e) { results[`search:${model}`] = { ok: false, error: errInfo(e) }; }
  console.log("search", model, JSON.stringify(results[`search:${model}`]).slice(0, 300));
}

// Code execution on one model
if (EXTRAS) try {
  await sleep(PACE_MS);
  const res = await withRetry(() => ai.models.generateContent({ model: MODELS[0], contents: "Compute the 30th Fibonacci number with Python and report it.", config: { tools: [{ codeExecution: {} }], maxOutputTokens: 3000 } }));
  results[`codeExecution:${MODELS[0]}`] = { ok: true, parts: res.candidates?.[0]?.content?.parts?.map((p: any) => ({ keys: Object.keys(p), executableCode: p.executableCode, codeExecutionResult: p.codeExecutionResult, text: p.text?.slice(0, 100) })), usage: res.usageMetadata };
} catch (e) { results[`codeExecution:${MODELS[0]}`] = { ok: false, error: errInfo(e) }; }
if (EXTRAS) console.log("codeExec", JSON.stringify(results[`codeExecution:${MODELS[0]}`]).slice(0, 400));

// toolConfig mode ANY (+ streamFunctionCallArguments is Vertex-only: SDK throws before sending) on 3.8
if (EXTRAS) try {
  await sleep(PACE_MS);
  const stream = await withRetry(() => ai.models.generateContentStream({ model: MODELS[0], contents: "Weather in Quebec City?", config: { tools: [weatherTool], toolConfig: { functionCallingConfig: { mode: "ANY" as any, allowedFunctionNames: ["get_weather"] } }, maxOutputTokens: 2000 } }));
  const chunks: any[] = [];
  for await (const c of stream) chunks.push(c.candidates?.[0]?.content?.parts);
  results[`fcModeAny:${MODELS[0]}`] = { ok: true, chunks };
} catch (e) { results[`fcModeAny:${MODELS[0]}`] = { ok: false, error: errInfo(e) }; }
if (EXTRAS) console.log("fcAny", JSON.stringify(results[`fcModeAny:${MODELS[0]}`]).slice(0, 500));

save(`03-tools${SUFFIX}.json`, results);