import os from "node:os";
import { parseArgs } from "node:util";
import { createOpenAICompatible } from "@ai-sdk/openai-compatible";
import { type ModelMessage, streamText } from "ai";
import chalk from "chalk";

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" },
    save: { type: "boolean" },
    s: { type: "boolean" },
    plain: { type: "boolean" },
    plainOut: { type: "boolean" },
    plainErr: { type: "boolean" },
  },
  strict: true,
  allowPositionals: true,
});

const home = os.homedir();
const convoFile = Bun.file(`${home}/.aicl_convo.json`);
const configFile = Bun.file(`${home}/.aicl_config.json`);

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;
  let saveConfig = args.save ?? args.s;

  let config: { model?: string };
  try {
    config = await configFile.json();
  } catch {
    config = {};
    saveConfig = true;
  }

  config.model = args.model ?? args.m ?? config.model;

  if (!config.model) {
    console.error("Please specify a model with --model/-m at least one time.");
    process.exit(1);
  }

  if (saveConfig) {
    await configFile.write(JSON.stringify(config, null, "  "));
    console.log(chalk.green("Configuration saved!"));
    return;
  }

  const renderer = new Renderer(plainOut, plainErr);

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

  let convo: { model: string; messages: ModelMessage[] } = {
    model: config.model,
    messages: [],
  };

  if (!newConvo) {
    try {
      convo = await convoFile.json();
    } catch {}

    if (args.model || args.m) {
      convo.model = args.model ?? args.m ?? config.model;
    }
  }

  const prompt = positionals[2] ?? "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 convoFile.write(JSON.stringify(convo));

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

await main();