import Anthropic from '@anthropic-ai/sdk';
import type { Stream } from '@anthropic-ai/sdk/streaming';

import { CitationItem, ModelTokensUsage } from '@/types/message';

import { ChatStreamCallbacks } from '../../types';
import {
  StreamContext,
  StreamProtocolChunk,
  StreamProtocolToolCallChunk,
  StreamToolCallChunkData,
  convertIterableToStream,
  createCallbacksTransformer,
  createSSEProtocolTransformer,
  createTokenSpeedCalculator,
} from './protocol';

export const transformAnthropicStream = (
  chunk: Anthropic.MessageStreamEvent,
  context: StreamContext,
): StreamProtocolChunk | StreamProtocolChunk[] => {
  // maybe need another structure to add support for multiple choices
  switch (chunk.type) {
    case 'message_start': {
      context.id = chunk.message.id;
      context.returnedCitationArray = [];
      let totalInputTokens = chunk.message.usage?.input_tokens;

      if (
        chunk.message.usage?.cache_creation_input_tokens ||
        chunk.message.usage?.cache_read_input_tokens
      ) {
        totalInputTokens =
          chunk.message.usage?.input_tokens +
          (chunk.message.usage.cache_creation_input_tokens || 0) +
          (chunk.message.usage.cache_read_input_tokens || 0);
      }

      context.usage = {
        inputCacheMissTokens: chunk.message.usage?.input_tokens,
        inputCachedTokens: chunk.message.usage?.cache_read_input_tokens || undefined,
        inputWriteCacheTokens: chunk.message.usage?.cache_creation_input_tokens || undefined,
        totalInputTokens,
        totalOutputTokens: chunk.message.usage?.output_tokens,
      };

      return { data: chunk.message, id: chunk.message.id, type: 'data' };
    }
    case 'content_block_start': {
      switch (chunk.content_block.type) {
        case 'redacted_thinking': {
          return {
            data: chunk.content_block.data,
            id: context.id,
            type: 'flagged_reasoning_signature',
          };
        }

        case 'text': {
          return { data: chunk.content_block.text, id: context.id, type: 'data' };
        }

        case 'server_tool_use':
        case 'tool_use': {
          const toolChunk = chunk.content_block;

          // if toolIndex is not defined, set it to 0
          if (typeof context.toolIndex === 'undefined') {
            context.toolIndex = 0;
          }
          // if toolIndex is defined, increment it
          else {
            context.toolIndex += 1;
          }

          const toolCall: StreamToolCallChunkData = {
            function: {
              arguments: '',
              name: toolChunk.name,
            },
            id: toolChunk.id,
            index: context.toolIndex,
            type: 'function',
          };

          context.tool = { id: toolChunk.id, index: context.toolIndex, name: toolChunk.name };

          return { data: [toolCall], id: context.id, type: 'tool_calls' };
        }

        /*
        case 'web_search_tool_result': {
          const citations = chunk.content_block.content;

          return [
            {
              data: {
                citations: (citations as any[]).map(
                  (item) =>
                    ({
                      title: item.title,
                      url: item.url,
                    }) as CitationItem,
                ),
              },
              id: context.id,
              type: 'grounding',
            },
          ];
        }
        */

        case 'thinking': {
          const thinkingChunk = chunk.content_block;

          // if there is signature in the thinking block, return both thinking and signature
          if (!!thinkingChunk.signature) {
            return [
              { data: thinkingChunk.thinking, id: context.id, type: 'reasoning' },
              { data: thinkingChunk.signature, id: context.id, type: 'reasoning_signature' },
            ];
          }

          if (typeof thinkingChunk.thinking === 'string')
            return { data: thinkingChunk.thinking, id: context.id, type: 'reasoning' };

          return { data: thinkingChunk, id: context.id, type: 'data' };
        }

        default: {
          break;
        }
      }

      return { data: chunk, id: context.id, type: 'data' };
    }

    case 'content_block_delta': {
      switch (chunk.delta.type) {
        case 'text_delta': {
          return { data: chunk.delta.text, id: context.id, type: 'text' };
        }

        case 'input_json_delta': {
          const delta = chunk.delta.partial_json;

          const toolCall: StreamToolCallChunkData = {
            function: { arguments: delta },
            index: context.toolIndex || 0,
            type: 'function',
          };

          return {
            data: [toolCall],
            id: context.id,
            type: 'tool_calls',
          } as StreamProtocolToolCallChunk;
        }

        case 'signature_delta': {
          return {
            data: chunk.delta.signature,
            id: context.id,
            type: 'reasoning_signature',
          };
        }

        case 'thinking_delta': {
          return {
            data: chunk.delta.thinking,
            id: context.id,
            type: 'reasoning',
          };
        }

        case 'citations_delta': {
          const citations = (chunk as any).delta.citation;

          if (context.returnedCitationArray) {
            context.returnedCitationArray.push({
              title: citations.title,
              url: citations.url,
            } as CitationItem);
          }

          return { data: null, id: context.id, type: 'text' };
        }

        default: {
          break;
        }
      }
      return { data: chunk, id: context.id, type: 'data' };
    }

    case 'message_delta': {
      const totalOutputTokens =
        chunk.usage?.output_tokens + (context.usage?.totalOutputTokens || 0);
      const totalInputTokens = context.usage?.totalInputTokens || 0;
      const totalTokens = totalInputTokens + totalOutputTokens;

      if (totalTokens > 0) {
        return [
          { data: chunk.delta.stop_reason, id: context.id, type: 'stop' },
          {
            data: {
              ...context.usage,
              totalInputTokens,
              totalOutputTokens,
              totalTokens,
            } as ModelTokensUsage,
            id: context.id,
            type: 'usage',
          },
        ];
      }
      return { data: chunk.delta.stop_reason, id: context.id, type: 'stop' };
    }

    case 'message_stop': {
      return [
        ...(context.returnedCitationArray?.length
          ? [
              {
                data: { citations: context.returnedCitationArray },
                id: context.id,
                type: 'grounding',
              },
            ]
          : []),
        { data: 'message_stop', id: context.id, type: 'stop' },
      ] as any;
    }

    default: {
      return { data: chunk, id: context.id, type: 'data' };
    }
  }
};

export interface AnthropicStreamOptions {
  callbacks?: ChatStreamCallbacks;
  inputStartAt?: number;
}

export const AnthropicStream = (
  stream: Stream<Anthropic.MessageStreamEvent> | ReadableStream,
  { callbacks, inputStartAt }: AnthropicStreamOptions = {},
) => {
  const streamStack: StreamContext = { id: '' };

  const readableStream =
    stream instanceof ReadableStream ? stream : convertIterableToStream(stream);

  return readableStream
    .pipeThrough(
      createTokenSpeedCalculator(transformAnthropicStream, { inputStartAt, streamStack }),
    )
    .pipeThrough(createSSEProtocolTransformer((c) => c, streamStack))
    .pipeThrough(createCallbacksTransformer(callbacks));
};