import { chat, save, MODELS, short } from "./lib.ts";

const base = (model: string) => ({
  model,
  messages: [{ role: "user", content: "Reply with the single word: pong" }],
  max_completion_tokens: 40,
});

const tool = {
  type: "function",
  function: {
    name: "get_weather",
    description: "Get weather for a city",
    parameters: { type: "object", properties: { city: { type: "string" } }, required: ["city"], additionalProperties: false },
  },
};

const variants: Record<string, (m: string) => Record<string, unknown>> = {
  temperature_0_5: (m) => ({ ...base(m), temperature: 0.5 }),
  temperature_2: (m) => ({ ...base(m), temperature: 2 }),
  temperature_2_5: (m) => ({ ...base(m), temperature: 2.5 }),
  top_p: (m) => ({ ...base(m), top_p: 0.9 }),
  top_k: (m) => ({ ...base(m), top_k: 40 }),
  max_tokens: (m) => ({ model: m, messages: base(m).messages, max_tokens: 40 }),
  both_max: (m) => ({ ...base(m), max_tokens: 40 }),
  min_tokens: (m) => ({ ...base(m), min_completion_tokens: 5 }),
  stop: (m) => ({ ...base(m), stop: ["\n", "END"] }),
  seed: (m) => ({ ...base(m), seed: 42 }),
  frequency_penalty: (m) => ({ ...base(m), frequency_penalty: 0.5 }),
  presence_penalty: (m) => ({ ...base(m), presence_penalty: 0.5 }),
  logit_bias: (m) => ({ ...base(m), logit_bias: { "1234": 5 } }),
  logprobs: (m) => ({ ...base(m), logprobs: true, top_logprobs: 2 }),
  n_2: (m) => ({ ...base(m), n: 2 }),
  reasoning_effort_none: (m) => ({ ...base(m), reasoning_effort: "none" }),
  reasoning_effort_low: (m) => ({ ...base(m), reasoning_effort: "low" }),
  reasoning_effort_medium: (m) => ({ ...base(m), reasoning_effort: "medium" }),
  reasoning_effort_high: (m) => ({ ...base(m), reasoning_effort: "high" }),
  reasoning_effort_minimal: (m) => ({ ...base(m), reasoning_effort: "minimal" }),
  disable_reasoning: (m) => ({ ...base(m), disable_reasoning: true }),
  reasoning_format_parsed: (m) => ({ ...base(m), reasoning_format: "parsed" }),
  reasoning_format_raw: (m) => ({ ...base(m), reasoning_format: "raw" }),
  reasoning_format_hidden: (m) => ({ ...base(m), reasoning_format: "hidden" }),
  clear_thinking: (m) => ({ ...base(m), clear_thinking: true }),
  json_object: (m) => ({
    model: m,
    messages: [{ role: "user", content: "Return a JSON object with key answer = 4." }],
    max_completion_tokens: 60,
    response_format: { type: "json_object" },
  }),
  json_object_stream: (m) => ({
    model: m,
    messages: [{ role: "user", content: "Return a JSON object with key answer = 4." }],
    max_completion_tokens: 60,
    response_format: { type: "json_object" },
    stream: true,
  }),
  json_schema_strict: (m) => ({
    model: m,
    messages: [{ role: "user", content: "Give the sum of 2 and 2." }],
    max_completion_tokens: 60,
    response_format: {
      type: "json_schema",
      json_schema: {
        name: "sum",
        strict: true,
        schema: { type: "object", properties: { answer: { type: "integer" } }, required: ["answer"], additionalProperties: false },
      },
    },
  }),
  tools_auto: (m) => ({ ...base(m), messages: [{ role: "user", content: "Weather in Montreal?" }], max_completion_tokens: 100, tools: [tool], tool_choice: "auto" }),
  tools_required: (m) => ({ ...base(m), messages: [{ role: "user", content: "Weather in Montreal?" }], max_completion_tokens: 100, tools: [tool], tool_choice: "required" }),
  tools_named: (m) => ({ ...base(m), messages: [{ role: "user", content: "Weather in Montreal?" }], max_completion_tokens: 100, tools: [tool], tool_choice: { type: "function", function: { name: "get_weather" } } }),
  tools_parallel_false: (m) => ({ ...base(m), messages: [{ role: "user", content: "Weather in Montreal and Paris?" }], max_completion_tokens: 100, tools: [tool], parallel_tool_calls: false }),
  tools_strict: (m) => ({ ...base(m), messages: [{ role: "user", content: "Weather in Montreal?" }], max_completion_tokens: 100, tools: [{ ...tool, function: { ...tool.function, strict: true } }] }),
  tools_plus_json_schema: (m) => ({
    ...base(m),
    messages: [{ role: "user", content: "Weather in Montreal?" }],
    max_completion_tokens: 100,
    tools: [tool],
    response_format: { type: "json_schema", json_schema: { name: "x", strict: true, schema: { type: "object", properties: { a: { type: "string" } }, required: ["a"], additionalProperties: false } } },
  }),
  developer_role: (m) => ({ ...base(m), messages: [{ role: "developer", content: "Be terse." }, ...base(m).messages] }),
  system_role: (m) => ({ ...base(m), messages: [{ role: "system", content: "Be terse." }, ...base(m).messages] }),
  unknown_param: (m) => ({ ...base(m), foo_bar: 1 }),
  user_field: (m) => ({ ...base(m), user: "probe-user" }),
  service_tier_priority: (m) => ({ ...base(m), service_tier: "priority" }),
  prompt_cache_key: (m) => ({ ...base(m), prompt_cache_key: "polyllm-probe" }),
  logprobs_alone: (m) => ({ ...base(m), logprobs: true }),
  max_completion_tokens_5: (m) => ({ ...base(m), messages: [{ role: "user", content: "Write 3 sentences about Montreal." }], max_completion_tokens: 5 }),
};

const out: Record<string, Record<string, unknown>> = {};
for (const model of MODELS) {
  out[model] = {};
  for (const [name, mk] of Object.entries(variants)) {
    const payload = mk(model);
    let r: any;
    if (payload.stream) {
      // streaming variant: just record status + first event/body
      const res = await fetch("https://api.cerebras.ai/v1/chat/completions", {
        method: "POST",
        headers: { "Content-Type": "application/json", Authorization: `Bearer ${process.env.CEREBRAS_API_KEY}` },
        body: JSON.stringify(payload),
      });
      const text = await res.text();
      let body: unknown = text;
      try {
        body = JSON.parse(text);
      } catch {}
      r = { status: res.status, body: typeof body === "string" ? short(body, 500) : body };
    } else {
      r = await chat(payload);
    }
    const b: any = r.body;
    const msg = b?.choices?.[0]?.message;
    const summary = {
      status: r.status,
      finish: b?.choices?.[0]?.finish_reason,
      content: short(msg?.content, 160),
      reasoning: msg?.reasoning !== undefined ? short(msg.reasoning, 100) : undefined,
      reasoningLen: typeof msg?.reasoning === "string" ? msg.reasoning.length : undefined,
      tool_calls: msg?.tool_calls ? short(msg.tool_calls, 200) : undefined,
      usage: b?.usage,
      logprobs: b?.choices?.[0]?.logprobs ? short(b.choices[0].logprobs, 200) : undefined,
      nChoices: b?.choices?.length,
      error: r.status !== 200 ? b : undefined,
      system_fingerprint: b?.system_fingerprint,
      service_tier: b?.service_tier,
      service_tier_used: b?.service_tier_used,
    };
    out[model][name] = summary;
    console.log(`${model} | ${name} | ${r.status} | ${summary.finish ?? ""} | ${short(summary.error ?? summary.content, 200)}${summary.reasoningLen !== undefined ? ` | reasoningLen=${summary.reasoningLen} rt=${b?.usage?.completion_tokens_details?.reasoning_tokens}` : ""}`);
  }
}
save("02-params.json", out);