// src/arkor/trainer.ts
import { createTrainer } from "arkor";

const WEBHOOK_URL = process.env.SLACK_WEBHOOK_URL;

async function postSlack(payload: Record<string, unknown>): Promise<void> {
  if (!WEBHOOK_URL) return;
  try {
    const res = await fetch(WEBHOOK_URL, {
      method: "POST",
      headers: { "content-type": "application/json" },
      body: JSON.stringify(payload),
    });
    if (!res.ok) {
      console.warn(`slack webhook ${res.status} ${res.statusText}`);
    }
  } catch (err) {
    // 通知失敗をコールバックの外に逃がさない。
    console.warn("slack webhook failed:", err);
  }
}

export const trainer = createTrainer({
  name: "support-bot-v1",
  model: "unsloth/gemma-4-E4B-it",
  dataset: { type: "huggingface", name: "arkorlab/triage-demo" },
  lora: { r: 16, alpha: 16 },
  maxSteps: 100,
  callbacks: {
    onCompleted: async ({ job, artifacts }) => {
      await postSlack({
        text: `:white_check_mark: *${job.name}* finished (${artifacts.length} artifact${artifacts.length === 1 ? "" : "s"}). Job \`${job.id}\`.`,
      });
    },
    onFailed: async ({ job, error }) => {
      await postSlack({
        text: `:x: <!here> *${job.name}* failed: ${error}\nJob \`${job.id}\`.`,
      });
    },
  },
});

onLog: async ({ step, loss }) => {
  if (step % 100 !== 0 || loss === null) return;
  await postSlack({ text: `step=${step} loss=${loss.toFixed(4)}` });
},

onCheckpoint: async ({ step, infer }) => {
  try {
    const res = await infer({
      messages: [{ role: "user", content: "Can't log in" }],
      stream: false,
      maxTokens: 80,
    });
    const data = (await res.json()) as {
      choices: Array<{ message: { content?: string | null } }>;
    };
    const sample = data.choices[0]?.message.content ?? "(empty)";
    await postSlack({ text: `step=${step} → ${sample}` });
  } catch (err) {
    console.warn("checkpoint sample failed:", err);
  }
}