interface StreamChoice {
  index?: number;
  delta: {
    content: string;
  };
  logprobs: null;
  finish_reason: string | null;
}

interface StreamResponse {
  type: string;
  data: {
    choices?: StreamChoice[];
    delta?: {
      text?: string;
    };
    type?: string;
    content_block?: {
      type: string;
      text: string;
    };
  };
}

const handleStreamData = (controller: ReadableStreamDefaultController, encoder: TextEncoder) => {
  return (data: StreamResponse, transformFn?: (data: StreamResponse) => StreamResponse) => {
    if (!data?.type || data.type !== 'chat') {
      return;
    }

    let transformedData: StreamResponse;

    if (transformFn) {
      transformedData = transformFn(data);
    } else {
      if (data.data.type === 'content_block_start' && data.data.content_block?.type === 'text') {
        transformedData = {
          type: 'chat',
          data: {
            choices: [
              {
                delta: {
                  content: data.data.content_block.text || '',
                },
                logprobs: null,
                finish_reason: null,
              },
            ],
          },
        };
      } else if (data.data.delta?.text) {
        transformedData = {
          type: 'chat',
          data: {
            choices: [
              {
                delta: {
                  content: data.data.delta.text,
                },
                logprobs: null,
                finish_reason: null,
              },
            ],
          },
        };
      } else if (data.data.choices?.[0]?.delta?.content) {
        transformedData = {
          type: 'chat',
          data: {
            choices: [
              {
                index: data.data.choices[0].index,
                delta: {
                  content: data.data.choices[0].delta.content,
                },
                logprobs: null,
                finish_reason: data.data.choices[0].finish_reason,
              },
            ],
          },
        };
      } else if (data.data.choices) {
        transformedData = data;
      } else {
        return;
      }
    }

    controller.enqueue(encoder.encode(`data: ${JSON.stringify(transformedData)}\n\n`));
  };
};

export default handleStreamData;