immersive2/src/react/services/diagramAI.ts

import {ChatMessage, DiagramToolCall, ToolResult} from "../types/chatTypes";
import {createEntity, connectEntities, removeEntity, modifyEntity, listEntities} from "./entityBridge";
import {v4 as uuidv4} from 'uuid';

const SYSTEM_PROMPT = `You are a 3D diagram assistant helping users create and modify diagrams in a virtual reality environment.

Available entity shapes: box, sphere, cylinder, cone, plane, person
Available colors: red, blue, green, yellow, orange, purple, cyan, pink, white, black, brown, gray (or any hex color like #ff5500)

Position coordinates:
- x: left/right (negative = left, positive = right)
- y: up/down (1.5 is eye level, 0 is floor)
- z: forward/backward (positive = toward user, negative = away)

When creating diagrams, think about good spatial layout:
- Spread entities apart to avoid overlap (at least 0.5 units)
- Use y=1.5 for entities at eye level
- Use z=2 to z=4 for comfortable viewing distance

Always use the provided tools to create, modify, or interact with entities. Be concise in your responses.`;

const TOOLS = [
    {
        name: "create_entity",
        description: "Create a 3D shape in the diagram. Use this to add new elements like boxes, spheres, cylinders, etc.",
        input_schema: {
            type: "object",
            properties: {
                shape: {
                    type: "string",
                    enum: ["box", "sphere", "cylinder", "cone", "plane", "person"],
                    description: "The type of 3D shape to create"
                },
                color: {
                    type: "string",
                    description: "Color name (red, blue, green, etc.) or hex code (#ff0000)"
                },
                label: {
                    type: "string",
                    description: "Text label to display on or near the entity"
                },
                position: {
                    type: "object",
                    properties: {
                        x: {type: "number", description: "Left/right position"},
                        y: {type: "number", description: "Up/down position (1.5 = eye level)"},
                        z: {type: "number", description: "Forward/backward position"}
                    },
                    description: "3D position. If not specified, defaults to (0, 1.5, 2)"
                }
            },
            required: ["shape"]
        }
    },
    {
        name: "connect_entities",
        description: "Draw a connection line between two entities. Use entity IDs or labels to identify them.",
        input_schema: {
            type: "object",
            properties: {
                from: {
                    type: "string",
                    description: "ID or label of the source entity"
                },
                to: {
                    type: "string",
                    description: "ID or label of the target entity"
                },
                color: {
                    type: "string",
                    description: "Color of the connection line"
                }
            },
            required: ["from", "to"]
        }
    },
    {
        name: "list_entities",
        description: "List all entities currently in the diagram. Use this to see what exists before connecting or modifying.",
        input_schema: {
            type: "object",
            properties: {}
        }
    },
    {
        name: "remove_entity",
        description: "Remove an entity from the diagram by its ID or label.",
        input_schema: {
            type: "object",
            properties: {
                target: {
                    type: "string",
                    description: "ID or label of the entity to remove"
                }
            },
            required: ["target"]
        }
    },
    {
        name: "modify_entity",
        description: "Modify an existing entity's properties like color, label, or position.",
        input_schema: {
            type: "object",
            properties: {
                target: {
                    type: "string",
                    description: "ID or label of the entity to modify"
                },
                color: {
                    type: "string",
                    description: "New color for the entity"
                },
                label: {
                    type: "string",
                    description: "New label text"
                },
                position: {
                    type: "object",
                    properties: {
                        x: {type: "number"},
                        y: {type: "number"},
                        z: {type: "number"}
                    }
                }
            },
            required: ["target"]
        }
    }
];

interface ClaudeMessage {
    role: 'user' | 'assistant';
    content: string | ClaudeContentBlock[];
}

interface ClaudeContentBlock {
    type: 'text' | 'tool_use' | 'tool_result';
    text?: string;
    id?: string;
    name?: string;
    input?: Record<string, unknown>;
    tool_use_id?: string;
    content?: string;
}

interface ClaudeResponse {
    content: ClaudeContentBlock[];
    stop_reason: 'end_turn' | 'tool_use' | 'max_tokens';
}

async function executeToolCall(toolCall: DiagramToolCall): Promise<ToolResult> {
    switch (toolCall.name) {
        case 'create_entity':
            return createEntity(toolCall.input);
        case 'connect_entities':
            return connectEntities(toolCall.input);
        case 'remove_entity':
            return removeEntity(toolCall.input);
        case 'modify_entity':
            return modifyEntity(toolCall.input);
        case 'list_entities':
            return await listEntities();
        default:
            return {
                toolName: 'unknown',
                success: false,
                message: 'Unknown tool'
            };
    }
}

export async function sendMessage(
    userMessage: string,
    conversationHistory: ChatMessage[],
    onToolResult?: (result: ToolResult) => void
): Promise<{response: string; toolResults: ToolResult[]}> {
    const messages: ClaudeMessage[] = conversationHistory
        .filter(m => !m.isLoading)
        .map(m => ({
            role: m.role,
            content: m.content
        }));

    messages.push({role: 'user', content: userMessage});

    const allToolResults: ToolResult[] = [];
    let finalResponse = '';
    let continueLoop = true;

    while (continueLoop) {
        const response = await fetch('/api/claude/v1/messages', {
            method: 'POST',
            headers: {
                'Content-Type': 'application/json',
            },
            body: JSON.stringify({
                model: 'claude-sonnet-4-20250514',
                max_tokens: 1024,
                system: SYSTEM_PROMPT,
                tools: TOOLS,
                messages
            })
        });

        if (!response.ok) {
            const errorText = await response.text();
            throw new Error(`API error: ${response.status} - ${errorText}`);
        }

        const data: ClaudeResponse = await response.json();

        const textBlocks = data.content.filter(b => b.type === 'text');
        const toolBlocks = data.content.filter(b => b.type === 'tool_use');

        if (textBlocks.length > 0) {
            finalResponse = textBlocks.map(b => b.text).join('\n');
        }

        if (data.stop_reason === 'tool_use' && toolBlocks.length > 0) {
            messages.push({
                role: 'assistant',
                content: data.content
            });

            const toolResults: ClaudeContentBlock[] = [];

            for (const toolBlock of toolBlocks) {
                const toolCall: DiagramToolCall = {
                    name: toolBlock.name as DiagramToolCall['name'],
                    input: toolBlock.input as DiagramToolCall['input']
                };

                const result = await executeToolCall(toolCall);
                allToolResults.push(result);
                onToolResult?.(result);

                toolResults.push({
                    type: 'tool_result',
                    tool_use_id: toolBlock.id,
                    content: result.message
                });
            }

            messages.push({
                role: 'user',
                content: toolResults
            });
        } else {
            continueLoop = false;
        }
    }

    return {response: finalResponse, toolResults: allToolResults};
}

export function createUserMessage(content: string): ChatMessage {
    return {
        id: uuidv4(),
        role: 'user',
        content,
        timestamp: new Date()
    };
}

export function createAssistantMessage(content: string, toolResults?: ToolResult[]): ChatMessage {
    return {
        id: uuidv4(),
        role: 'assistant',
        content,
        timestamp: new Date(),
        toolResults
    };
}

export function createLoadingMessage(): ChatMessage {
    return {
        id: uuidv4(),
        role: 'assistant',
        content: '',
        timestamp: new Date(),
        isLoading: true
    };
}