// @ts-nocheck
import { $, ProcessOutput } from "bun";
import * as fs from "fs";

// Replace input in the prompt file
async function readAndReplace(filePath: string, input: string, templateReplace: string = "{{input}}"): Promise<string> {
    const fileContent = fs.readFileSync(filePath, "utf8");
    return fileContent.replace(templateReplace, input);
}

// Chain processing of prompts and models
async function promptChain(initialInput: string, prompts: string[], models: string[]): Promise<string> {
    if (prompts.length !== models.length) {
        throw new Error("Number of prompts and models must match.");
    }

    let result = initialInput;
    for (let i = 0; i < prompts.length; i++) {
        const promptFile = prompts[i];
        const model = models[i];

        // Replace {{input}} in the prompt file
        const replacedPrompt = await readAndReplace(promptFile, result);

        // Run the model based on prefix
        console.log(`${'⭐️'.repeat(i + 1)} Prompt #${i + 1}
--------------------------------`);
        console.log(`
🤖 Model: '${model}'
📄 Input file: '${promptFile}'
📝 Prompt content: 
${replacedPrompt}
🎯 Result:
`);
        let output: ProcessOutput;

        if (model.startsWith("ollama:")) {
            const ollamaModel = model.replace("ollama:", "");
            output = await $`ollama run ${ollamaModel} "${replacedPrompt}"`;
        } else if (model.startsWith("llm:")) {
            const llmModel = model.replace("llm:", "");
            output = await $`llm -m ${llmModel} ${replacedPrompt}`;
        } else {
            throw new Error(`Unsupported model prefix in: ${model}`);
        }

        result = await output.text().trim(); // Update result with the output of the model
    }

    return result;
}

// Main logic
(async () => {
    const [initialInput, promptFilesArg, modelsArg] = process.argv.slice(2); // First argument is initial input, second is prompts, third is models
    const promptFiles = promptFilesArg.split(","); // Split prompt files on comma
    const models = modelsArg.split(","); // Split models on comma

    if (!initialInput || promptFiles.length === 0 || models.length === 0) {
        console.error("Usage: bun run chain.ts \"<first input value>\" <prompt1,prompt2,...> <model1,model2,...>");
        process.exit(1);
    }

    try {
        const finalResult = await promptChain(initialInput, promptFiles, models);
    } catch (error) {
        console.error("Error:", error.message);
        process.exit(1);
    }
})();