#!/usr/bin/env -S deno run -A

import { TextLineStream } from "https://deno.land/std@0.224.0/streams/mod.ts"
import ollama, { Message } from "npm:ollama@0.5.15"

function main() {
	const names = new Names()
	const messages = new SpillingQueue<Message>(500)
	listen(8080, (conn) => session(conn, names.pick_for(conn.remoteAddr.hostname), messages))
}

class SpillingQueue<T> {
	public items: T[] = []

	public constructor(
		public limit: number,
	) {}

	public push(value: T) {
		this.items.push(value)
		if (this.items.length > this.limit) return this.items.splice(0, 1)[0]
	}
}

async function listen(port: number, handler: (conn: Deno.TcpConn) => Promise<unknown>) {
	console.log("Listening on port", port)
	for await (const connection of Deno.listen({ transport: "tcp", port })) {
		handler(connection).catch((except) => console.log("Session closed", except))
	}
}

const header = (name: string) => `
  KUB

Kub is a LLM running on TCP which includes all messages in the same context.
As Kub is having several discussions at once, he will identify you as : ${name}
───────────────────────────────────────────────────────────────────────────────
`

async function session(conn: Deno.TcpConn, name: string, messages: SpillingQueue<Message>) {
	console.log("Opening session for", name, "at", conn.remoteAddr.hostname, conn.remoteAddr.port)
	conn.setNoDelay(true)
	await conn.write(new TextEncoder().encode(header(name)))
	async function prompt() {
		await conn.write(new TextEncoder().encode("> "))
	}
	const lines = conn.readable
		.pipeThrough(new TextDecoderStream())
		.pipeThrough(new TextLineStream())
	await prompt()
	for await (const line of lines) {
		await conn.write(new TextEncoder().encode("\n< "))
		const user_content = `(message from:${name}) ${line}`
		const user_message = { role: "user", content: user_content }
		const response = await ollama.chat({
			model: "llama3.2",
			messages: request_messages(messages.items, user_message),
			stream: true,
			options: {},
		})
		let content = "", role = "assistant"
		for await (const part of response) {
			await conn.write(new TextEncoder().encode(part.message.content))
			content += part.message.content
			role = part.message.role
		}
		await conn.write(new TextEncoder().encode("\n\n"))
		await prompt()
		messages.push(user_message)
		messages.push({ role, content })
	}
}

function request_messages(history_messages: Message[], user_message: Message) {
	return [
		{ role: "system", content: "You are 'Kub', a bot which talks to people on the internet. Be concise and helpful." },
		...history_messages,
		user_message,
	] as Message[]
}

class Names {
	options = ["Adam", "Bastien", "Corentin", "Dominique", "Ethan"]

	public pick_for(key: string) {
		let sum = 0
		for (const letter of key) sum += letter.charCodeAt(0)
		const index = Math.floor((sum + Math.random() * 100) % this.options.length)
		return this.options[index]
	}
}

if (import.meta.main) main()