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

const TRIAGE_SCHEMA: Record<string, unknown> = {
  type: "object",
  properties: {
    category: { type: "string" },
    urgency: { type: "string", enum: ["low", "medium", "high"] },
    summary: { type: "string" },
    nextAction: { type: "string" },
  },
  required: ["category", "urgency", "summary", "nextAction"],
  additionalProperties: false,
};

interface TriageOutput {
  category: string;
  urgency: "low" | "medium" | "high";
  summary: string;
  nextAction: string;
}

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: {
    onCheckpoint: async ({ step, infer }) => {
      try {
        const res = await infer({
          messages: [
            { role: "user", content: "I can't log in to my account." },
          ],
          stream: false,
          maxTokens: 200,
          responseFormat: {
            type: "json_schema",
            json_schema: {
              name: "triage",
              schema: TRIAGE_SCHEMA,
              strict: true,
            },
          },
        });
        const data = (await res.json()) as {
          choices: Array<{ message: { content: string } }>;
        };
        const content = data.choices[0]?.message.content;
        if (content === undefined || content === "") {
          throw new Error("triage check returned empty content");
        }
        const parsed = JSON.parse(content) as TriageOutput;
        console.log(`step=${step} triage=`, parsed);
      } catch (err) {
        console.error(`step=${step} triage check failed:`, err);
      }
    },
  },
});

const VALID_CATEGORIES = new Set([
  "auth",
  "billing",
  "bug",
  "feature_request",
  "other",
]);

onCheckpoint: async ({ step, infer }) => {
  const parsed = await runTriage(infer);     // 上のコードの呼び出し
  if (parsed && !VALID_CATEGORIES.has(parsed.category)) {
    console.warn(
      `step=${step} category=${parsed.category} not in label set, aborting`,
    );
    controller.abort();
    await trainer.cancel().catch(() => {});
  }
},

onCheckpoint: async ({ step, infer }) => {
  const res = await infer({
    messages: [
      { role: "user", content: "What's the status of order #4821?" },
    ],
    tools: [
      {
        type: "function",
        function: {
          name: "get_order_status",
          description: "Look up the current status of a customer order.",
          parameters: {
            type: "object",
            properties: { orderId: { type: "string" } },
            required: ["orderId"],
          },
        },
      },
    ],
    toolChoice: "auto",
    stream: false,
  });
  const data = (await res.json()) as {
    choices: Array<{ message: { tool_calls?: Array<{ function: { name: string; arguments: string } }> } }>;
  };
  const call = data.choices[0]?.message.tool_calls?.[0];
  if (call) {
    const args = JSON.parse(call.function.arguments) as { orderId: string };
    console.log(`step=${step} tool=${call.function.name} args=`, args);
  }
};

const res = await infer({
  messages: [{ role: "user", content: "Classify urgency: I can't log in." }],
  structuredOutputs: { choice: ["low", "medium", "high"] },
  stream: false,
});
const data = (await res.json()) as {
  choices: Array<{ message: { content: string } }>;
};
const urgency = data.choices[0]?.message.content; // 必ず "low" / "medium" / "high" のいずれか