import { parseArgs } from "node:util";
import { Stream } from "misskey-js";
import type { Note } from "misskey-js/entities.js";
import type { ChatHistoryItem, LLamaChatPromptOptions } from "node-llama-cpp";
import { createHarness, type HarnessName, isHarnessName } from "./lib/harness";
import { createGrammar, getModel, LlmSession } from "./lib/llm";
import { expandReplyTree, getNotes, me, misskey } from "./lib/misskey";
import {
  formatConversationNote,
  formatTimeline,
  postJobPrompt,
  replyJobPrompt,
} from "./lib/prompt";
import { ensurePolite, isPolitenessMode } from "./lib/rephrase";
import {
  getSamplingOptions,
  isSamplingProfile,
  type SamplingProfile,
} from "./lib/sampling";
import { sanitizeText, sleep } from "./lib/util";

const { values } = parseArgs({
  args: Bun.argv,
  options: {
    test: {
      type: "boolean",
      short: "t",
      default: false,
    },
    harness: {
      type: "string",
      default: Bun.env["HARNESS"] ?? "tagged",
    },
    politeness: {
      type: "string",
      default: Bun.env["POLITENESS"] ?? "auto",
    },
    sampling: {
      type: "string",
      default: Bun.env["SAMPLING"] ?? "creative",
    },
  },
  strict: true,
  allowPositionals: true,
});

if (!isHarnessName(values.harness)) {
  throw new Error(`Unknown harness: ${values.harness}`);
}
const harnessName: HarnessName = values.harness;
if (!isPolitenessMode(values.politeness)) {
  throw new Error(`Unknown politeness mode: ${values.politeness}`);
}
const politenessMode = values.politeness;
if (!isSamplingProfile(values.sampling)) {
  throw new Error(`Unknown sampling profile: ${values.sampling}`);
}
const samplingProfile: SamplingProfile = values.sampling;

const modelName =
  Bun.env["MODEL"] ?? "mradermacher/gemma-2-baku-2b-it-GGUF:IQ4_XS";
console.log(`* loading model '${modelName}'`);
const model = await getModel(modelName);
const grammar = await createGrammar("あるびのちゃん");
const harness = createHarness(harnessName, grammar);
console.log(`* output harness: ${harness.name}`);
console.log(`* sampling profile: ${samplingProfile}`);
const baseChatPromptOptions = {
  maxTokens: 192,
} as const satisfies LLamaChatPromptOptions;

type Job =
  // read posts and post a note
  | { type: "post" }
  // reply to a specific note
  | {
      type: "reply";
      id: string;
      visibility: Note["visibility"];
      last: Note;
      history: Note[];
    };

async function processPostJob() {
  const notes = await getNotes(10, 0, 5);
  const input = formatTimeline(notes, me.id);
  const text = await (async () => {
    await using postJobSession = new LlmSession(
      model,
      postJobPrompt + harness.promptSuffix,
    );
    await postJobSession.init();
    const raw = await postJobSession.prompt(input, {
      ...baseChatPromptOptions,
      ...harness.options,
      ...getSamplingOptions(modelName, samplingProfile, "post"),
    });
    process.stderr.write("\n");
    return harness.parse(raw);
  })();
  if (text) {
    const output = sanitizeText(
      (await ensurePolite(model, text, politenessMode)).text,
    );
    if (values.test) {
      console.log(`\n${output}\n`);
      return;
    }
    await misskey.request("notes/create", {
      visibility: "public",
      text: output,
    });
  }
}

async function processReplyJob(job: Extract<Job, { type: "reply" }>) {
  const history: ChatHistoryItem[] = job.history.map((n) => {
    if (n.userId === me.id) {
      return {
        type: "model",
        response: [harness.formatHistory(n.text ?? "")],
      } as const;
    }
    return {
      type: "user",
      text: formatConversationNote(n, me.id),
    } as const;
  });
  const text = await (async () => {
    await using session = new LlmSession(
      model,
      replyJobPrompt + harness.promptSuffix,
      history,
    );
    await session.init();
    const raw = await session.prompt(formatConversationNote(job.last, me.id), {
      ...baseChatPromptOptions,
      ...harness.options,
      ...getSamplingOptions(modelName, samplingProfile, "reply"),
    });
    process.stderr.write("\n");
    return harness.parse(raw);
  })();

  if (text) {
    const output = sanitizeText(
      (await ensurePolite(model, text, politenessMode)).text,
    );
    if (values.test) {
      console.log(`\n${output}\n`);
      return;
    }
    await misskey.request("notes/create", {
      visibility: job.visibility,
      text: output,
      replyId: job.id,
    });
  }
}

/** execute a job */
async function processJob(job: Job) {
  switch (job.type) {
    case "post":
      await processPostJob();
      break;
    case "reply":
      await processReplyJob(job);
      break;
  }
}

const jobs: Job[] = [];

let stream: Stream;
let channel: ReturnType<typeof stream.useChannel<"main">>;

/** dispose stream for recreation */
function disposeStream() {
  channel.removeAllListeners();
  channel.dispose();
  stream.removeAllListeners();
  stream.close();
}

/** connect to streaming API and add handlers */
function initializeStream() {
  stream = new Stream(
    Bun.env["MISSKEY_ORIGIN"] ?? "https://misskey.cannorin.net",
    {
      token: Bun.env["MISSKEY_CREDENTIAL"] ?? "",
    },
    {
      binaryType: "arraybuffer",
    },
  );
  channel = stream.useChannel("main");

  // notify when connected
  stream.on("_connected_", () => {
    console.log("* connected");
  });

  // notify when disconnected (it will reconnect automatically)
  stream.on("_disconnected_", () => {
    console.log("* disconnected");
  });

  // push a reply job when receiving a mention
  channel.on("mention", async (e) => {
    if (e.text && e.userId !== me.id && !e.user.isBot) {
      const replyTree = await expandReplyTree(e);
      console.log(
        `* push: reply (${e.id}, ${replyTree.history.length + 1} msgs)`,
      );
      jobs.push({
        type: "reply",
        id: e.id,
        visibility: e.visibility,
        ...replyTree,
      });
    }
  });

  // follow back non-bot users
  channel.on("followed", async (e) => {
    if (!e.isBot) {
      await misskey.request("following/create", { userId: e.id });
    }
  });
}

/** take jobs from the queue in arrival order */
async function runJob() {
  while (true) {
    const job = jobs.shift();
    if (job) {
      console.log(`* take: ${job.type}`);
      try {
        await processJob(job);
        console.log("* job complete");
      } catch (e) {
        console.log(`* error: ${JSON.stringify(e)}`);
        if (e instanceof Error) console.log(e.stack);
      }
    }
    await sleep(1000); // 1sec
  }
}

/** push a job to the job queue */
async function pushJob() {
  while (true) {
    console.log("* push: post");
    jobs.push({ type: "post" });
    // random interval between 10 and 40 minutes
    const interval = Math.floor(Math.random() * 30 + 10) * 60 * 1000;
    console.log(
      `* info: next post job in ${Math.round(interval / 60000)} minutes`,
    );
    await sleep(interval);
  }
}

async function test() {
  try {
    console.log("* test a post job:");
    await processJob({ type: "post" });
    await processJob({ type: "post" });
    await processJob({ type: "post" });
    await processJob({ type: "post" });
    await processJob({ type: "post" });
  } catch (e) {
    console.error(e);
    if (e instanceof Error) console.log(e.stack);
  }
}

async function main() {
  try {
    initializeStream();
    try {
      await Promise.all([runJob(), pushJob()]);
    } catch (e) {
      console.error(e);
      if (e instanceof Error) console.log(e.stack);
    }
  } finally {
    disposeStream();
  }
}

if (values.test) await test();
else await main();