"use strict";
var __awaiter = (this && this.__awaiter) || function (thisArg, _arguments, P, generator) {
    function adopt(value) { return value instanceof P ? value : new P(function (resolve) { resolve(value); }); }
    return new (P || (P = Promise))(function (resolve, reject) {
        function fulfilled(value) { try { step(generator.next(value)); } catch (e) { reject(e); } }
        function rejected(value) { try { step(generator["throw"](value)); } catch (e) { reject(e); } }
        function step(result) { result.done ? resolve(result.value) : adopt(result.value).then(fulfilled, rejected); }
        step((generator = generator.apply(thisArg, _arguments || [])).next());
    });
};
var __asyncValues = (this && this.__asyncValues) || function (o) {
    if (!Symbol.asyncIterator) throw new TypeError("Symbol.asyncIterator is not defined.");
    var m = o[Symbol.asyncIterator], i;
    return m ? m.call(o) : (o = typeof __values === "function" ? __values(o) : o[Symbol.iterator](), i = {}, verb("next"), verb("throw"), verb("return"), i[Symbol.asyncIterator] = function () { return this; }, i);
    function verb(n) { i[n] = o[n] && function (v) { return new Promise(function (resolve, reject) { v = o[n](v), settle(resolve, reject, v.done, v.value); }); }; }
    function settle(resolve, reject, d, v) { Promise.resolve(v).then(function(v) { resolve({ value: v, done: d }); }, reject); }
};
var __importDefault = (this && this.__importDefault) || function (mod) {
    return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.ingestClaudeResponseStream = ingestClaudeResponseStream;
exports.ingestAmazonMessageStream = ingestAmazonMessageStream;
exports.ingestOpenaiMessageStream = ingestOpenaiMessageStream;
exports.ingestClaudeMessages = ingestClaudeMessages;
const logger_1 = require("../utils/logger");
const chalk_1 = __importDefault(require("chalk"));
function ingestClaudeResponseStream(responseStream, res, uuid) {
    return __awaiter(this, void 0, void 0, function* () {
        var _a, responseStream_1, responseStream_1_1;
        var _b, e_1, _c, _d;
        let status = "Waiting";
        let messageText = "";
        let toolInput = null;
        // process.stdout.write(chalk.green("-----------------------------\n"))
        res.write("streamItemStart:" + JSON.stringify({ type: "temp_uuid", uuid }) + ":streamItemEnd");
        try {
            for (_a = true, responseStream_1 = __asyncValues(responseStream); responseStream_1_1 = yield responseStream_1.next(), _b = responseStream_1_1.done, !_b; _a = true) {
                _d = responseStream_1_1.value;
                _a = false;
                const streamItem = _d;
                const stringifiedItem = JSON.stringify(streamItem);
                // process.stdout.write('\rReceived from claude: ' + stringifiedItem);
                res.write("streamItemStart:" + stringifiedItem + ":streamItemEnd");
                switch (streamItem.type) {
                    case "message_start":
                        {
                            // res.write('\nClaude has started responding\n')
                        }
                        break;
                    case "message_delta":
                        {
                            // console.log("Claude is done with the chat.", streamItem.delta);
                            if (streamItem.delta.stop_reason === "end_turn") {
                                console.log("Claude is done with the chat.", streamItem);
                                status = "EndTurnReached";
                            }
                            else if (streamItem.delta.stop_reason === "stop_sequence") {
                                console.log("Claude is stopping because it encountered a stop sequence");
                            }
                            else if (streamItem.delta.stop_reason === "max_tokens") {
                                console.log("Claude is stopping because max tokens were exhausted");
                                status = "MaxTokensUsed";
                            }
                            else if (streamItem.delta.stop_reason === "tool_use") {
                                console.log("Claude is stopping because it needs a tool response");
                                status = "ToolCallRequested";
                            }
                        }
                        break;
                    case "message_stop":
                        {
                            console.log("stopping: ", streamItem);
                        }
                        break;
                    case "content_block_start":
                        {
                            switch (streamItem.content_block.type) {
                                case "text":
                                    {
                                        // res.write(streamItem.content_block.text)
                                        messageText = streamItem.content_block.text;
                                    }
                                    break;
                                case "tool_use":
                                    {
                                        // res.write(streamItem.content_block.name)
                                        toolInput = {
                                            tool: streamItem.content_block.name,
                                            id: streamItem.content_block.id,
                                            partialInput: "",
                                        };
                                    }
                                    break;
                                default:
                                    console.error("unknown delta in content_block");
                            }
                        }
                        break;
                    case "content_block_delta":
                        {
                            switch (streamItem.delta.type) {
                                case "input_json_delta":
                                    {
                                        // res.write(streamItem.delta.partial_json)
                                        toolInput.partialInput += streamItem.delta.partial_json;
                                    }
                                    break;
                                case "text_delta":
                                    {
                                        // res.write(streamItem.delta.text);
                                        messageText += streamItem.delta.text;
                                    }
                                    break;
                            }
                        }
                        break;
                    case "content_block_stop":
                        {
                            // res.write(`\n Claude has finished responding \n`);
                            // console.log("content block ended: ", streamItem);
                        }
                        break;
                    default:
                        console.error("unknown stream item: ", streamItem);
                }
            }
        }
        catch (e_1_1) { e_1 = { error: e_1_1 }; }
        finally {
            try {
                if (!_a && !_b && (_c = responseStream_1.return)) yield _c.call(responseStream_1);
            }
            finally { if (e_1) throw e_1.error; }
        }
        // process.stdout.write(chalk.green("\n--------------------------------\n"));
        if (toolInput) {
            return { messageText, toolInput: [toolInput], status, uuid };
        }
        else {
            return { messageText, status, uuid };
        }
    });
}
function ingestAmazonMessageStream(response, res, uuid) {
    return __awaiter(this, void 0, void 0, function* () {
        var _a, e_2, _b, _c;
        var _d, _e, _f;
        let blocks = [];
        let amazonUsage = null;
        if (response.stream) {
            res.write("streamItemStart:" + JSON.stringify({ type: "temp_uuid", uuid }) + ":streamItemEnd");
            try {
                for (var _g = true, _h = __asyncValues(response.stream), _j; _j = yield _h.next(), _a = _j.done, !_a; _g = true) {
                    _c = _j.value;
                    _g = false;
                    const chunk = _c;
                    if (chunk.messageStart) {
                        res.write("streamItemStart:" +
                            JSON.stringify({
                                type: "message_start",
                                delta: {},
                                content_block: {}
                            }) +
                            ":streamItemEnd");
                    }
                    if (chunk.contentBlockStart) {
                        const index = chunk.contentBlockStart.contentBlockIndex || 0;
                        if ((_d = chunk.contentBlockStart.start) === null || _d === void 0 ? void 0 : _d.toolUse) {
                            blocks[index] = {
                                name: chunk.contentBlockStart.start.toolUse.name || "",
                                id: chunk.contentBlockStart.start.toolUse.toolUseId || "",
                                input: ""
                            };
                            res.write("streamItemStart:"
                                + JSON.stringify({
                                    type: "content_block_start",
                                    index,
                                    content_block: {
                                        type: "tool_use",
                                        id: chunk.contentBlockStart.start.toolUse.toolUseId || "",
                                        name: chunk.contentBlockStart.start.toolUse.name,
                                        partialInput: ""
                                    }
                                })
                                + ":streamItemEnd");
                        }
                        else {
                            blocks[index] = "";
                        }
                    }
                    if (chunk.contentBlockDelta) {
                        const index = chunk.contentBlockDelta.contentBlockIndex || 0;
                        if ((_e = chunk.contentBlockDelta.delta) === null || _e === void 0 ? void 0 : _e.toolUse) {
                            const toolpart = chunk.contentBlockDelta.delta.toolUse.input;
                            if (typeof blocks[index] !== "string") {
                                blocks[index].input += toolpart;
                                res.write("streamItemStart:"
                                    + JSON.stringify({
                                        type: "content_block_delta",
                                        index: index,
                                        delta: {
                                            type: "input_json_delta",
                                            partial_json: toolpart
                                        }
                                    })
                                    + ":streamItemEnd");
                            }
                            else {
                                logger_1.logger.error("Failed to process chunk in aws response " + toolpart);
                            }
                        }
                        if ((_f = chunk.contentBlockDelta.delta) === null || _f === void 0 ? void 0 : _f.text) {
                            const deltext = chunk.contentBlockDelta.delta.text || "";
                            // console.log(deltext);
                            if (!blocks[index]) {
                                blocks[index] = "";
                            }
                            blocks[index] += deltext;
                            res.write("streamItemStart:" +
                                JSON.stringify({
                                    type: "content_block_delta",
                                    delta: {
                                        type: "text_delta",
                                        text: deltext
                                    },
                                    content_block: {}
                                }) +
                                ":streamItemEnd");
                        }
                    }
                    if (chunk.metadata) {
                        const usage = chunk.metadata.usage;
                        if (usage) {
                            amazonUsage = usage;
                        }
                        console.log("Usage in amazon: ", usage);
                    }
                }
            }
            catch (e_2_1) { e_2 = { error: e_2_1 }; }
            finally {
                try {
                    if (!_g && !_a && (_b = _h.return)) yield _b.call(_h);
                }
                finally { if (e_2) throw e_2.error; }
            }
            res.write("streamItemStart:" +
                JSON.stringify({
                    type: "message_stop",
                    delta: {},
                    content_block: {}
                }) +
                ":streamItemEnd");
            const toolCalls = blocks.filter(it => typeof it !== "string").map(it => {
                return {
                    id: it.id,
                    tool: it.name,
                    partialInput: it.input
                };
            });
            const messageText = blocks.find(it => typeof it === "string") || "";
            if (toolCalls.length > 0) {
                return {
                    amazonUsage,
                    messageText,
                    toolInput: toolCalls,
                    status: "EndTurnReached",
                    uuid
                };
            }
            else {
                return {
                    amazonUsage,
                    messageText,
                    status: "EndTurnReached",
                    uuid
                };
            }
        }
        throw new Error("Cannot process this amazon stream");
    });
}
function ingestOpenaiMessageStream(responseStreamP, res, uuid) {
    return __awaiter(this, void 0, void 0, function* () {
        var _a, e_3, _b, _c;
        var _d, _e, _f, _g, _h;
        let status = "Waiting";
        logger_1.logger.info(chalk_1.default.green("Waiting for openai to start responding"));
        const responseStream = yield responseStreamP;
        logger_1.logger.info(chalk_1.default.green("Openai started responding"));
        let messageText = "";
        const toolCalls = [];
        res.write("streamItemStart:" + JSON.stringify({ type: "temp_uuid", uuid }) + ":streamItemEnd");
        // process.stdout.write(chalk.green("-----------------------------\n"));
        res.write("streamItemStart:" +
            JSON.stringify({
                type: "message_start",
                delta: {},
                content_block: {}
            }) +
            ":streamItemEnd");
        try {
            for (var _j = true, responseStream_2 = __asyncValues(responseStream), responseStream_2_1; responseStream_2_1 = yield responseStream_2.next(), _a = responseStream_2_1.done, !_a; _j = true) {
                _c = responseStream_2_1.value;
                _j = false;
                const streamItem = _c;
                // const stringifiedItem = JSON.stringify(streamItem);
                // process.stdout.write('\rReceived from claude: ' + stringifiedItem);
                if (streamItem.choices.length === 0) {
                    logger_1.logger.info("No choices obtained");
                }
                else if (streamItem.choices.length > 1) {
                    logger_1.logger.info("Received multiple choices. Will continue with first");
                }
                // console.log("delta: ", streamItem.choices[0]?.delta);
                const choice = streamItem.choices.find(it => it.index === 0);
                const delta = choice === null || choice === void 0 ? void 0 : choice.delta;
                if (delta && delta.tool_calls) {
                    for (let i = 0; i < delta.tool_calls.length; i++) {
                        const tool_call = delta.tool_calls[i];
                        if (tool_call.index === undefined) {
                            tool_call.index = i;
                        }
                        while (toolCalls.length <= tool_call.index) {
                            toolCalls.push({
                                id: "",
                                tool: "",
                                partialInput: ""
                            });
                        }
                        const targetAccumulator = toolCalls[tool_call.index];
                        if (tool_call.function) {
                            const tool_call_id = tool_call.id || "toolu-" + Date.now();
                            targetAccumulator.id = tool_call_id;
                            res.write("streamItemStart:"
                                + JSON.stringify({
                                    type: "content_block_start",
                                    index: tool_call.index,
                                    content_block: {
                                        type: "tool_use",
                                        id: tool_call_id,
                                        name: (_d = tool_call.function) === null || _d === void 0 ? void 0 : _d.name,
                                        partialInput: ""
                                    }
                                })
                                + ":streamItemEnd");
                        }
                        if ((_e = tool_call.function) === null || _e === void 0 ? void 0 : _e.arguments) {
                            targetAccumulator.partialInput += (_f = tool_call.function) === null || _f === void 0 ? void 0 : _f.arguments;
                            res.write("streamItemStart:"
                                + JSON.stringify({
                                    type: "content_block_delta",
                                    index: tool_call.index,
                                    delta: {
                                        type: "input_json_delta",
                                        partial_json: (_g = tool_call.function) === null || _g === void 0 ? void 0 : _g.arguments
                                    }
                                })
                                + ":streamItemEnd");
                        }
                        if ((_h = tool_call.function) === null || _h === void 0 ? void 0 : _h.name) {
                            targetAccumulator.tool += tool_call.function.name;
                        }
                    }
                }
                else if (delta && typeof delta.content === "string") {
                    messageText += delta.content;
                    res.write("streamItemStart:" +
                        JSON.stringify({
                            type: "content_block_delta",
                            delta: {
                                type: "text_delta",
                                text: delta.content
                            },
                            content_block: {}
                        }) +
                        ":streamItemEnd");
                }
                else if (delta && choice.finish_reason === "tool_calls") {
                    status = "EndTurnReached";
                    console.log("[OPENAI_AGENT] Openai is finished because " + choice.finish_reason);
                }
                else {
                    console.error("[OPENAI_AGENT] Unrecognized delta received", choice);
                }
            }
        }
        catch (e_3_1) { e_3 = { error: e_3_1 }; }
        finally {
            try {
                if (!_j && !_a && (_b = responseStream_2.return)) yield _b.call(responseStream_2);
            }
            finally { if (e_3) throw e_3.error; }
        }
        // process.stdout.write(chalk.green("\n--------------------------------\n"));
        res.write("streamItemStart:" +
            JSON.stringify({
                type: "message_stop",
                delta: {},
                content_block: {}
            }) +
            ":streamItemEnd");
        if (toolCalls.length > 0) {
            return {
                messageText,
                toolInput: toolCalls,
                status: "EndTurnReached",
                uuid
            };
        }
        else {
            return {
                messageText,
                status: "EndTurnReached",
                uuid
            };
        }
    });
}
function ingestClaudeMessages(responseStream) {
    return __awaiter(this, void 0, void 0, function* () {
        var _a, responseStream_3, responseStream_3_1;
        var _b, e_4, _c, _d;
        let status = "Waiting";
        let messageText = "";
        let toolInput = null;
        try {
            for (_a = true, responseStream_3 = __asyncValues(responseStream); responseStream_3_1 = yield responseStream_3.next(), _b = responseStream_3_1.done, !_b; _a = true) {
                _d = responseStream_3_1.value;
                _a = false;
                const streamItem = _d;
                switch (streamItem.type) {
                    case "message_start":
                        {
                            // console.log("Message started");
                        }
                        break;
                    case "message_delta":
                        {
                            // console.log("Claude is done with the chat.", streamItem.delta);
                            if (streamItem.delta.stop_reason === "end_turn") {
                                console.log("Claude is done with the chat.");
                                status = "EndTurnReached";
                            }
                        }
                        break;
                    case "message_stop":
                        {
                            // console.log("stopping: ", streamItem);
                        }
                        break;
                    case "content_block_start":
                        {
                            switch (streamItem.content_block.type) {
                                case "text":
                                    {
                                        messageText = streamItem.content_block.text;
                                    }
                                    break;
                                case "tool_use":
                                    {
                                        toolInput = {
                                            tool: streamItem.content_block.name,
                                            id: streamItem.content_block.id,
                                            partialInput: "",
                                        };
                                    }
                                    break;
                                default:
                                    console.error("unknown delta in content_block");
                            }
                        }
                        break;
                    case "content_block_delta":
                        {
                            switch (streamItem.delta.type) {
                                case "input_json_delta":
                                    {
                                        toolInput.partialInput += streamItem.delta.partial_json;
                                    }
                                    break;
                                case "text_delta":
                                    {
                                        messageText += streamItem.delta.text;
                                    }
                                    break;
                            }
                        }
                        break;
                    case "content_block_stop":
                        {
                            console.log("content block ended: ", streamItem);
                        }
                        break;
                    default:
                        console.error("unknown stream item: ", streamItem);
                }
            }
        }
        catch (e_4_1) { e_4 = { error: e_4_1 }; }
        finally {
            try {
                if (!_a && !_b && (_c = responseStream_3.return)) yield _c.call(responseStream_3);
            }
            finally { if (e_4) throw e_4.error; }
        }
        if (toolInput) {
            return { messageText, toolInput: [toolInput], end_turn: status === "EndTurnReached" };
        }
        else {
            return { messageText, end_turn: status === "EndTurnReached" };
        }
    });
}
