import { Response as ExpressResponse } from "express";
import { MessageStream } from "@anthropic-ai/sdk/lib/MessageStream";
import { APIPromise } from "openai/core";
import { ChatCompletionChunk } from "openai/resources/chat/completions";
import { Stream } from "openai/streaming";
import { createWriteStream } from "node:fs";
import { RawMessageStreamEvent } from "@anthropic-ai/sdk/resources";
import { logger } from "../utils/logger";
import chalk from "chalk";
import { ConverseStreamCommandOutput } from "@aws-sdk/client-bedrock-runtime";
import { TokenUsage } from "@aws-sdk/client-bedrock-runtime";
type ParserStatus =
  | "Waiting"
  | "CollectingTextDelta"
  | "CollectingToolInputJSONDelta"
  | "EndTurnReached"
  | "MaxTokensUsed"
  | "ToolCallRequested";

export async function ingestClaudeResponseStream(
  responseStream: MessageStream,
  res: ExpressResponse,
  uuid: string
) {
  let status: ParserStatus = "Waiting";
  let messageText = "";
  let toolInput: any = null;

  // process.stdout.write(chalk.green("-----------------------------\n"))
  res.write(
    "streamItemStart:" + JSON.stringify({type: "temp_uuid", uuid}) + ":streamItemEnd"
  );
  for await (const streamItem of responseStream) {
    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);
    }
  }
  // process.stdout.write(chalk.green("\n--------------------------------\n"));

  if (toolInput) {
    return { messageText, toolInput: [toolInput], status, uuid };
  } else {
    return {messageText, status, uuid}
  }
}

export async function ingestAmazonMessageStream(
  response: ConverseStreamCommandOutput, 
  res: ExpressResponse,
  uuid: string
) {
  let blocks: Array<string|{name: string; id: string; input: string;}> = [];
  let amazonUsage: TokenUsage | null = null;
  if (response.stream) {
    res.write(
      "streamItemStart:" + JSON.stringify({type: "temp_uuid", uuid}) + ":streamItemEnd"
    );
    for await (const chunk of response.stream) {
      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 (chunk.contentBlockStart.start?.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 (chunk.contentBlockDelta.delta?.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.error("Failed to process chunk in aws response " + toolpart)
          }
        }

        if (chunk.contentBlockDelta.delta?.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);
      }
    }

    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");
}

export async function ingestOpenaiMessageStream(
  responseStreamP: APIPromise<Stream<ChatCompletionChunk>>, 
  res: ExpressResponse,
  uuid: string
) {
  let status: ParserStatus = "Waiting";
  logger.info(chalk.green("Waiting for openai to start responding"));
  const responseStream = await responseStreamP;
  logger.info(chalk.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"
  );

  for await (const streamItem of responseStream) {
    // const stringifiedItem = JSON.stringify(streamItem);
    // process.stdout.write('\rReceived from claude: ' + stringifiedItem);
    if (streamItem.choices.length === 0) {
      logger.info("No choices obtained");
    } else if (streamItem.choices.length > 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?.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: tool_call.function?.name,
                partialInput: ""
              }
            })
            + ":streamItemEnd"
          );
        }

        if (tool_call.function?.arguments) {
          targetAccumulator.partialInput += tool_call.function?.arguments;
          res.write("streamItemStart:" 
            + JSON.stringify({
              type: "content_block_delta",
              index: tool_call.index,
              delta: {
                type: "input_json_delta",
                partial_json: tool_call.function?.arguments
              }
            }) 
            + ":streamItemEnd"
          )
        } 

        if (tool_call.function?.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);
    }
  }

  // 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
    };
  }
}

export async function ingestClaudeMessages(responseStream: MessageStream) {
  let status: ParserStatus = "Waiting";
  let messageText = "";
  let toolInput: any = null;

  for await (const streamItem of responseStream) {
    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);
    }
  }

  if (toolInput) {
    return { messageText, toolInput: [toolInput], end_turn: status === "EndTurnReached" };
  } else {
    return { messageText, end_turn: status === "EndTurnReached" };
  }
}
