import { parseArgs } from "node:util";
import { createOpenAICompatible } from "@ai-sdk/openai-compatible";
import { streamText } from "ai";

import { initConfig, loadConfig } from "./src/config";
import { Conversation } from "./src/conversation";
import Renderer from "./src/render";

const { values: args, positionals } = parseArgs({
  args: Bun.argv,
  options: {
    model: { type: "string" },
    m: { type: "string" },
    new: { type: "boolean" },
    n: { type: "boolean" },
    config: { type: "boolean" },
    c: { type: "boolean" },
    plain: { type: "boolean" },
    plainOut: { type: "boolean" },
    plainErr: { type: "boolean" },
  },
  strict: true,
  allowPositionals: true,
});

async function main() {
  const newConvo = args.new ?? args.n;
  const plainOut = args.plainOut ?? args.plain ?? !process.stdout.isTTY;
  const plainErr = args.plainErr ?? args.plain ?? !process.stderr.isTTY;
  const modelFlag = args.model ?? args.m;
  const configFlag = args.config ?? args.c;

  const config = await loadConfig();

  const model = modelFlag ?? config.model;

  if (configFlag) {
    await initConfig(config);
    return;
  }

  const renderer = new Renderer(plainOut, plainErr);

  const provider = createOpenAICompatible({
    name: "llama.cpp",
    apiKey: "local",
    baseURL: config.baseURL,
    includeUsage: true, // Include usage information in streaming responses
  });

  const convo = new Conversation(model);

  if (!newConvo) {
    await convo.load();
    if (modelFlag) convo.model = modelFlag;
  }

  const prompt = positionals.slice(2).join(" ") || "Introduce yourself.";
  convo.messages.push({ role: "user", content: prompt });

  const { fullStream, text, totalUsage } = streamText({
    model: provider(convo.model),
    messages: convo.messages,
  });

  const start = Date.now();
  await renderer.renderStream(fullStream);
  const stop = Date.now();

  convo.messages.push({ role: "assistant", content: await text });
  await convo.save();

  await renderer.renderStats(totalUsage, start, stop);
}

await main();