import type { TextStreamPart, ToolSet } from 'ai';

export function transformCerebrasReasoningStream<TOOLS extends ToolSet>(): (options: {
    tools: TOOLS;
    stopStream: () => void;
}) => TransformStream<TextStreamPart<TOOLS>, TextStreamPart<TOOLS>> {

    return (_opts) => {
        let isThinking = false;
        let currentReasoningId: string | null = null;

        let bufferedTextStart: TextStreamPart<TOOLS> | null = null;
        let hasEmittedTextStart = false;

        return new TransformStream<TextStreamPart<TOOLS>, TextStreamPart<TOOLS>>({
            transform(chunk, controller) {
                if (chunk.type === 'text-start') {
                    bufferedTextStart = chunk;
                    return;
                }

                if (chunk.type === 'text-delta') {
                    let text = chunk.text;

                    if (text.includes('<think>')) {
                        isThinking = true;
                        currentReasoningId = crypto.randomUUID();

                        const [before, after] = text.split('<think>');

                        if (before && before.trim().length > 0) {
                            if (bufferedTextStart && !hasEmittedTextStart) {
                                controller.enqueue(bufferedTextStart);
                                hasEmittedTextStart = true;
                            }
                            controller.enqueue({ type: 'text-delta', text: before, id: chunk.id });
                        }

                        controller.enqueue({ type: 'reasoning-start', id: currentReasoningId });

                        if (after) {
                            controller.enqueue({ type: 'reasoning-delta', text: after, id: currentReasoningId });
                        }
                        return;
                    }

                    if (text.includes('</think>')) {
                        isThinking = false;
                        const [before, after] = text.split('</think>');

                        if (before && currentReasoningId !== null) {
                            controller.enqueue({ type: 'reasoning-delta', text: before, id: currentReasoningId });
                        }

                        if (currentReasoningId !== null) {
                            controller.enqueue({ type: 'reasoning-end', id: currentReasoningId });
                        }

                        if (after && after.length > 0) {
                            if (bufferedTextStart && !hasEmittedTextStart) {
                                controller.enqueue(bufferedTextStart);
                                hasEmittedTextStart = true;
                            }
                            controller.enqueue({ type: 'text-delta', text: after, id: chunk.id });
                        }
                        return;
                    }

                    if (isThinking && currentReasoningId !== null) {
                        controller.enqueue({ type: 'reasoning-delta', text: text, id: currentReasoningId });
                    } else {
                        if (bufferedTextStart && !hasEmittedTextStart) {
                            controller.enqueue(bufferedTextStart);
                            hasEmittedTextStart = true;
                        }
                        controller.enqueue(chunk);
                    }
                } else {
                    controller.enqueue(chunk);
                }
            },
        });
    }
}