#!/usr/bin/env node
import fs from "node:fs/promises";
import consumers from "node:stream/consumers";
import util from "node:util";
import { createOpenAICompatible } from "@ai-sdk/openai-compatible";
import { streamText } from "ai";

import { initConfig, loadConfig } from "./config";
import { Conversation } from "./conversation";
import { handleTopLevel } from "./errors";
import {
  base,
  commitDiff,
  contextFence,
  diffFence,
  fileFence,
  review,
} from "./prompts";
import Renderer from "./render";

async function main() {
  const { values: args, positionals } = util.parseArgs({
    args: process.argv.slice(2),
    options: {
      model: { type: "string" },
      m: { type: "string" },
      new: { type: "boolean" },
      n: { type: "boolean" },
      config: { type: "boolean" },
      c: { type: "boolean" },
      plain: { type: "boolean" },
      "plain-out": { type: "boolean" },
      "plain-err": { type: "boolean" },
      version: { type: "boolean" },
      v: { type: "boolean" },
      "git-message": { type: "boolean" },
      g: { type: "boolean" },
      review: { type: "boolean" },
      r: { type: "boolean" },
      file: { type: "string" },
      f: { type: "string" },
      stdin: { type: "boolean" },
      i: { type: "boolean" },
    },
    strict: true,
    allowPositionals: true,
  });

  const newConvo = args.new ?? args.n;
  const plainOut = args["plain-out"] ?? args.plain ?? !process.stdout.isTTY;
  const plainErr = args["plain-err"] ?? args.plain ?? !process.stderr.isTTY;
  const modelOverride = args.model ?? args.m;
  const showVersion = args.version ?? args.v;
  const editConfig = args.config ?? args.c;
  const generateGitMessage = args["git-message"] ?? args.g;
  const generateReview = args.review ?? args.r;
  const inputFile = args.file ?? args.f;
  const stdinContent = args.stdin ?? args.i;

  if (showVersion) {
    console.log("0.0.0");
    return;
  }

  const config = await loadConfig(editConfig);

  if (editConfig || !config) {
    await initConfig(config);
    return;
  }

  const model = modelOverride ?? config.model;

  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 (modelOverride) convo.model = modelOverride;
  }

  const input = positionals.join(" ").trim();
  let stdinText = "";
  if (stdinContent) stdinText = await consumers.text(process.stdin);
  let fileText = "";
  if (inputFile)
    fileText = fileFence(inputFile, await fs.readFile(inputFile, "utf8"));

  let prompt = base(input, [fileText, contextFence(stdinText)]);
  if (generateGitMessage) {
    prompt = commitDiff(stdinText, [fileText]);
  } else if (generateReview) {
    prompt = review([fileText, stdinText ? diffFence(stdinText) : "", input]);
  }

  if (!prompt) throw new Error("No prompt given.");

  convo.messages.push({ role: "user", content: prompt });

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

  await renderer.renderStream(fullStream);

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

  await renderer.renderStats(totalUsage);
}

main().catch(handleTopLevel);