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

export function createLongcatTransformer<TOOLS extends ToolSet>(): (options: {
    tools: TOOLS;
    stopStream: () => void;
}) => TransformStream<TextStreamPart<TOOLS>, TextStreamPart<TOOLS>> {
    let buffer = '';
    let hasToolCallInStep = false;
    let lastChunkId: string | undefined;
    let lastChunkType: 'text' | 'reasoning' | undefined;
    let step = 0;

    return (_opts) => {
        return new TransformStream<TextStreamPart<TOOLS>, TextStreamPart<TOOLS>>({
            transform(chunk, controller) {
                if (chunk.type === 'finish-step' || chunk.type === 'finish') {
                    step++;

                    if (hasToolCallInStep) {
                        // We clone the chunk and overwrite the finishReason.
                        // This tricks the SDK into thinking the model requested a tool natively.
                        const modifiedChunk = {
                            ...chunk,
                            finishReason: 'tool-calls' as const,
                        };

                        // Reset for the next potential step
                        if (chunk.type === 'finish-step') {
                            hasToolCallInStep = false;
                        }

                        controller.enqueue(modifiedChunk);
                        return;
                    }
                }

                if (chunk.type === 'text-start' || chunk.type === 'reasoning-start') {
                    lastChunkId = chunk.id;
                    lastChunkType = chunk.type.split('-')[1] as 'text' | 'reasoning';
                }

                // We only care about text chunks
                if (chunk.type !== 'text-delta' && chunk.type !== 'reasoning-delta') {
                    controller.enqueue(chunk);
                    return;
                }

                buffer += chunk.text;

                // Check if we have a full tool call in the buffer
                const pattern = /<longcat_tool_call>([\s\S]*?)<\/longcat_tool_call>/g;
                let lastIndex = 0;
                let match;

                while ((match = pattern.exec(buffer)) !== null) {
                    console.log("longcat tool call found at index", match.index);

                    // 1. Enqueue any text that appeared BEFORE the tool call
                    const textBefore = buffer.substring(lastIndex, match.index);
                    if (textBefore) {
                        controller.enqueue({ type: chunk.type, text: textBefore, id: lastChunkId ?? chunk.type.includes('reasoning') ? `reasoning-${step}` : `text-${step}` });
                    }

                    // 2. Parse the XML content
                    const content = match[1]!.trim();
                    const toolNameMatch = content.match(/^([^\s<]+)/);

                    if (toolNameMatch) {
                        hasToolCallInStep = true;

                        const toolName = toolNameMatch[1];
                        const args: Record<string, any> = {};
                        const argRegex = /<longcat_arg_key>(.*?)<\/longcat_arg_key>\s*<longcat_arg_value>(.*?)<\/longcat_arg_value>/gs;

                        let argMatch;
                        while ((argMatch = argRegex.exec(content)) !== null) {
                            args[argMatch[1]!.trim()] = argMatch[2]!.trim();
                        }

                        // 3. EMIT A TOOL CALL PART
                        // This is the "magic" - the SDK will see this and act as if the LLM 
                        // called a native tool.
                        const toolCallId = `lc-${Date.now()}-${Math.random().toString(36).substr(2, 5)}`;
                        controller.enqueue({
                            type: 'tool-call',
                            // @ts-ignore
                            id: toolCallId,
                            toolCallId,
                            toolName,
                            input: args,
                            dynamic: true,
                        });
                    }

                    lastIndex = pattern.lastIndex;
                }

                // Keep the remaining buffer (unclosed tags) for the next chunk
                buffer = buffer.substring(lastIndex);

                // If there's no open tag starting, we can flush the buffer as text
                if (!buffer.includes('<longcat_tool_call>')) {
                    if (buffer) {
                        controller.enqueue({ type: chunk.type, text: buffer, id: lastChunkId ?? chunk.type.includes('reasoning') ? `reasoning-${step}` : `text-${step}` });
                        buffer = '';
                    }
                }
            },
            flush(controller) {
                if (buffer && lastChunkId && lastChunkType) {
                    controller.enqueue({ type: `${lastChunkType}-delta`, text: buffer, id: lastChunkId });
                }
            }
        });
    };
}