#!/usr/bin/env node

import readline from 'readline';

import { exec } from 'child_process';

import { promisify } from 'util';

import { chat, MODELS } from './Chat.mjs';

const execAsync = promisify(exec);

// Default model if not specified

const DEFAULT_MODEL = "s";

// Get model from environment variable or use default

const MODEL = process.argv[2] || DEFAULT_MODEL;

console.log(`Welcome to ChatSH. Model: ${MODELS[MODEL]||MODEL}\n`);

// System prompt to set the assistant's behavior

const SYSTEM_PROMPT = `You are ChatSH, an AI language model that specializes in assisting users with tasks on their system using shell commands. ChatSH operates in two modes: COMMAND MODE and CHAT MODE.

# GUIDE for COMMAND NODE:

1. The USER asks you to perform a SYSTEM TASK.

2. ChatSH answers with a SHELL SCRIPT to perform the task.

# GUIDE for CHAT MODE:

1. The USER asks an ARBITRARY QUESTION or OPEN-ENDED MESSAGE.

2. ChatSH answers it with a concise, factual response.

# EXAMPLE 1:

<USER>

Show me local files.

</USER>

<ChatSH>

\`\`\`sh

ls

\`\`\`

</ChatSH>

<SYSTEM>

example.gif example.tst example_dir/

</SYSTEM>

<USER>

Is there any text file?

</USER>

<ChatSH>

Yes, there is a file named example.txt in the current directory

</ChatSH>

<USER>

Show me the contents of example.txt.

</USER>

<ChatSH>

cat example.txt

</ChatSH>

<SYSTEM>

Na matemática, um grupo é um conjunto de elementos associados a uma operação que combina dois elementos quaisquer para formar um terceiro. Para se qualificar como grupo o conjunto e a operação devem satisfazer algumas condições chamadas axiomas de grupo: associatividade, elemento neutro e elementos inversos.

</SYSTEM>

<USER>

Can you translate the first sentence to English?

</USER>

<ChatSH>

In mathematics, a group is a set of elements associated with an operation that combines any two elements to form a third element.

</ChatSH>

<USER>

What is the monster group?

</USER>

<ChatSH>

The monster group is the largest sporadic simple group, with order approximately 8×10^53.

</ChatSH>

<USER>

What is a simple group?

</USER>

<ChatSH>

A simple group is a nontrivial group that contains no proper nontrivial normal subgroups.

</ChatSH>

<USER>

Move the text file to that dir.

</USER>

<ChatSH>

\`\`\`sh

mv example.txt example_dir/

\`\`\`

</ChatSH>

# EXAMPLE 2:

<USER>

Fix main.ts

</USER>

<ChatSH>

\`\`\`sh

cat main.ts

tsc --noEmit main.ts

\`\`\`

</ChatSH>

<SYSTEM>

import List from './list';

import map from './map';

const n_1_2: List<number> = { $: "cons", head: 1, tail: { $: "cons", head: 2, tail: { $: "nil" }}};

const n_2_4: List<number> = map(n_1_2, x => x * 2);

console.log(JSON.stringify(n_2_4));

map.ts:7:22 - error TS2345: Argument of type '(val: T) => U' is not assignable to parameter of type 'List<unknown>'.

7 let tail = map(fun, list.tail);

~~~

</SYSTEM>

<ChatSH>

\`\`\`sh

ls

\`\`\`

</ChatSH>

<SYSTEM>

list.ts main.ts map.ts

</SYSTEM>

<ChatSH>

\`\`\`sh

cat map.ts

\`\`\`

</ChatSH>

<SYSTEM>

import List from './list';

function map<T, U>(list: List<T>, fun: (val: T) => U): List<U> {

switch (list.$) {

case "cons":

let head = fun(list.head);

let tail = map(fun, list.tail);

return { $: "cons", head, tail };

case "nil":

return { $: "nil" };

}

}

export default map;

</SYSTEM>

<ChatSH>

\`\`\`sh

cat > map.ts << EOL

import List from './list';

function map<T, U>(list: List<T>, fun: (val: T) => U): List<U> {

switch (list.$) {

case "cons":

let head = fun(list.head);

let tail = map(list.tail, fun);

return { $: "cons", head, tail };

case "nil":

return { $: "nil" };

}

}

export default map;

EOL

tsc --noEmit map.ts

\`\`\`

</ChatSH>

<SYSTEM>

</SYSTEM>

<ChatSH>

Done.

</ChatSH

## NOTES:

- In COMMAND MODE, ChatSH MUST answer with A SINGLE SH BLOCK.

- In COMMAND MODE, ChatSH MUST NOT answer with ENGLISH EXPLANATION.

- In TEXT MODE, ChatSH MUST ALWAYS answer with TEXT.

- In TEXT MODE, ChatSH MUST NEVER answer with SH BLOCK.

- ChatSH MUST be CONCISE, OBJECTIVE, CORRECT and USEFUL.

- ChatSH MUST NEVER attempt to install new tools. Assume they're available.

- ChatSH's interpreter can only process one SH per answer.

- On TypeScript:

- Use 'tsc --noEmit file.ts' to type-check.

- Use 'tsx file.ts' to run.

- Never generate js files.

- When a task is completed, STOP using commands. Just answer with "Done.".

- The system shell in use is: ${await get_shell()}.`;

// Create readline interface for user input/output

const rl = readline.createInterface({

input: process.stdin,

output: process.stdout,

terminal: true

});

// Create a stateful asker

const ask = chat(MODEL);

// Utility function to prompt the user for input

async function prompt(query) {

return new Promise(resolve => {

rl.question(query, resolve);

});

}

// If there are words after the 'chatsh', set them as the initialUserMessage

var initialUserMessage = process.argv.slice(3).join(' ');

// Main interaction loop

async function main() {

let lastOutput = "";

while (true) {

let userMessage;

if (initialUserMessage) {

userMessage = initialUserMessage;

initialUserMessage = null;

} else {

process.stdout.write('\x1b[1m'); // blue color

userMessage = await prompt('λ ');

process.stdout.write('\x1b[0m'); // reset color

}

try {

const fullMessage = userMessage.trim() !== ''

? `<SYSTEM>\n${lastOutput.trim()}\n</SYSTEM>\n<USER>\n${userMessage}\n</USER>\n`

: `<SYSTEM>\n${lastOutput.trim()}\n</SYSTEM>`;

const assistantMessage = await ask(fullMessage, { system: SYSTEM_PROMPT, model: MODEL });

console.log();

const code = extractCode(assistantMessage);

lastOutput = "";

if (code) {

console.log("\x1b[31mPress enter to execute, or 'N' to cancel.\x1b[0m");

const answer = await prompt('');

// TODO: delete the warning above from the terminal

process.stdout.moveCursor(0, -2);

process.stdout.clearLine(2);

if (answer.toLowerCase() === 'n') {

console.log('Execution skipped.');

lastOutput = "Command skipped.\n";

} else {

try {

const {stdout, stderr} = await execAsync(code);

const output = `${stdout.trim()}${stderr.trim()}`;

console.log('\x1b[2m' + output.trim() + '\x1b[0m');

lastOutput = output;

} catch(error) {

const output = `${error.stdout?.trim()||''}${error.stderr?.trim()||''}`;

console.log('\x1b[2m' + output.trim() + '\x1b[0m');

lastOutput = output;

}

}

}

} catch(error) {

console.error(`Error: ${error.message}`);

}

}

}

// Utility function to extract code from the assistant's message

function extractCode(text) {

const match = text.match(/```sh([\s\S]*?)```/);

return match ? match[1].trim() : null;

}

async function get_shell() {

const shellInfo = (await execAsync('uname -a && $SHELL --version')).stdout.trim();

return shellInfo;

}

main();