// Probe 02: generationConfig parameter acceptance matrix per model. One request per (model, param).
import { ai, MODELS, SUFFIX, save, errInfo, shortText, withRetry, sleep, PACE_MS } from "./lib.js";

const schema = { type: "object", properties: { answer: { type: "string" } }, required: ["answer"] };

const CASES: Record<string, Record<string, unknown>> = {
  temperature_0: { temperature: 0 },
  temperature_1_5: { temperature: 1.5 },
  temperature_2: { temperature: 2 },
  temperature_2_5_out_of_range: { temperature: 2.5 },
  topP_0_9: { topP: 0.9 },
  topK_40: { topK: 40 },
  seed_42: { seed: 42 },
  stopSequences: { stopSequences: ["DONE"] },
  frequencyPenalty_0_5: { frequencyPenalty: 0.5 },
  presencePenalty_0_5: { presencePenalty: 0.5 },
  candidateCount_2: { candidateCount: 2 },
  responseMimeType_json: { responseMimeType: "application/json" },
  responseSchema: { responseMimeType: "application/json", responseSchema: { type: "OBJECT", properties: { answer: { type: "STRING" } }, required: ["answer"] } },
  responseJsonSchema: { responseMimeType: "application/json", responseJsonSchema: schema },
  responseMimeType_enum: { responseMimeType: "text/x.enum", responseSchema: { type: "STRING", enum: ["red", "blue", "green"] } },
  thinkingBudget_0: { thinkingConfig: { thinkingBudget: 0 } },
  thinkingBudget_1024: { thinkingConfig: { thinkingBudget: 1024 } },
  thinkingBudget_minus1: { thinkingConfig: { thinkingBudget: -1 } },
  includeThoughts: { thinkingConfig: { includeThoughts: true } },
  thinkingLevel_minimal: { thinkingConfig: { thinkingLevel: "MINIMAL" } },
  thinkingLevel_low: { thinkingConfig: { thinkingLevel: "LOW" } },
  thinkingLevel_medium: { thinkingConfig: { thinkingLevel: "MEDIUM" } },
  thinkingLevel_high: { thinkingConfig: { thinkingLevel: "HIGH" } },
  thinkingLevel_and_budget: { thinkingConfig: { thinkingLevel: "LOW", thinkingBudget: 512 } },
  responseLogprobs: { responseLogprobs: true, logprobs: 2 },
};

const only = MODELS;
const onlyCases = process.argv[3] ? new Set(process.argv[3].split(",")) : null;
const results: any = {};

for (const model of only) {
  results[model] = {};
  for (const [name, cfg] of Object.entries(CASES)) {
    if (onlyCases && !onlyCases.has(name)) continue;
    const t0 = Date.now();
    try {
      const res = await withRetry(() => ai.models.generateContent({
        model,
        contents: name.startsWith("responseMimeType_enum") ? "Which colour is the sky? Answer with one of the enum values." : "Answer in <= 5 words: what is 2+2? Then write DONE.",
        config: { maxOutputTokens: 1500, ...cfg } as any,
      }));
      const usage = res.usageMetadata ?? {};
      results[model][name] = {
        ok: true,
        text: shortText(res),
        finishReason: res.candidates?.[0]?.finishReason,
        thoughtsTokenCount: usage.thoughtsTokenCount ?? 0,
        candidatesTokenCount: usage.candidatesTokenCount,
        nCandidates: res.candidates?.length,
        hasThoughtParts: res.candidates?.[0]?.content?.parts?.some((p: any) => p.thought) ?? false,
        ms: Date.now() - t0,
      };
    } catch (e) {
      results[model][name] = { ok: false, error: errInfo(e), ms: Date.now() - t0 };
    }
    await sleep(PACE_MS);
    const r = results[model][name];
    console.log(model.padEnd(24), name.padEnd(30), r.ok ? `OK  thoughts=${r.thoughtsTokenCount} fin=${r.finishReason} "${r.text.slice(0, 40).replace(/\n/g, " ")}"` : `FAIL ${r.error.httpStatus} ${r.error.status}: ${r.error.message.slice(0, 160)}`);
  }
}
save(`02-params${SUFFIX}${onlyCases ? "-subset-" + Date.now() : ""}.json`, results);