sim/app/api/chat/route.ts

import { OpenAI } from 'openai'
import { NextResponse } from 'next/server'
import { z } from 'zod'
import { ChatCompletionMessageParam } from 'openai/resources/chat/completions'

// Validation schemas
const MessageSchema = z.object({
  role: z.enum(['user', 'assistant', 'system']),
  content: z.string()
})

const RequestSchema = z.object({
  messages: z.array(MessageSchema),
  workflowState: z.object({
    blocks: z.record(z.any()),
    edges: z.array(z.any())
  })
})

// Define function schemas with strict typing
const workflowActions = {
  addBlock: {
    description: "Add one new block to the workflow",
    parameters: {
      type: "object",
      required: ["type"],
      properties: {
        type: {
          type: "string",
          enum: ["agent", "api", "condition", "function", "router"],
          description: "The type of block to add"
        },
        name: {
          type: "string",
          description: "Optional custom name for the block. Do not provide a name unless the user has specified it."
        },
        position: {
          type: "object",
          description: "Optional position for the block. Do not provide a position unless the user has specified it.",
          properties: {
            x: { type: "number" },
            y: { type: "number" }
          }
        }
      }
    }
  },
  addEdge: {
    description: "Create a connection between two blocks",
    parameters: {
      type: "object",
      required: ["sourceId", "targetId"],
      properties: {
        sourceId: {
          type: "string",
          description: "ID of the source block"
        },
        targetId: {
          type: "string",
          description: "ID of the target block"
        },
        sourceHandle: { 
          type: "string",
          description: "Optional handle identifier for the source connection point"
        },
        targetHandle: { 
          type: "string",
          description: "Optional handle identifier for the target connection point"
        }
      }
    }
  },
  removeBlock: {
    description: "Remove a block from the workflow",
    parameters: {
      type: "object",
      required: ["id"],
      properties: {
        id: { type: "string", description: "ID of the block to remove" }
      }
    }
  },
  removeEdge: {
    description: "Remove a connection between blocks",
    parameters: {
      type: "object",
      required: ["id"],
      properties: {
        id: { type: "string", description: "ID of the edge to remove" }
      }
    }
  }
}

// System prompt that references workflow state
const getSystemPrompt = (workflowState: any) => {
  const blockCount = Object.keys(workflowState.blocks).length
  const edgeCount = workflowState.edges.length
  
  // Create a summary of existing blocks
  const blockSummary = Object.values(workflowState.blocks)
    .map((block: any) => `- ${block.type} block named "${block.name}" with id ${block.id}`)
    .join('\n')

  // Create a summary of existing edges
  const edgeSummary = workflowState.edges
    .map((edge: any) => `- ${edge.source} -> ${edge.target} with id ${edge.id}`)
    .join('\n')

  return `You are a workflow assistant that helps users modify their workflow by adding/removing blocks and connections.

Current Workflow State:
${blockCount === 0 ? 'The workflow is empty.' : `${blockSummary}

Connections:
${edgeCount === 0 ? 'No connections between blocks.' : edgeSummary}`}

When users request changes:
- Consider existing blocks when suggesting connections
- Provide clear feedback about what actions you've taken

Use the following functions to modify the workflow:
1. Use the addBlock function to create new blocks
2. Use the addEdge function to connect blocks
3. Use the removeBlock function to remove blocks
4. Use the removeEdge function to remove connections

Only use the provided functions and respond naturally to the user's requests.`
}

export async function POST(request: Request) {
  try {
    // Validate API key
    const apiKey = request.headers.get('X-OpenAI-Key')
    if (!apiKey) {
      return NextResponse.json(
        { error: 'OpenAI API key is required' },
        { status: 401 }
      )
    }

    // Parse and validate request body
    const body = await request.json()
    const validatedData = RequestSchema.parse(body)
    const { messages, workflowState } = validatedData

    // Initialize OpenAI client
    const openai = new OpenAI({ apiKey })
    
    // Create message history with workflow context
    const messageHistory = [
      { role: 'system', content: getSystemPrompt(workflowState) },
      ...messages
    ]

    // Make OpenAI API call with workflow context
    const completion = await openai.chat.completions.create({
      model: "gpt-4o",
      messages: messageHistory as ChatCompletionMessageParam[],
      tools: Object.entries(workflowActions).map(([name, config]) => ({
        type: 'function',
        function: {
          name,
          description: config.description,
          parameters: config.parameters
        }
      })),
      tool_choice: "auto"
    })

    const message = completion.choices[0].message

    // Process tool calls if present
    if (message.tool_calls) {
        console.log(message.tool_calls)
      const actions = message.tool_calls.map(call => ({
        name: call.function.name,
        parameters: JSON.parse(call.function.arguments)
      }))

      return NextResponse.json({
        message: message.content || "I've updated the workflow based on your request.",
        actions
      })
    }

    // Return response with no actions
    return NextResponse.json({ 
      message: message.content || "I'm not sure what changes to make to the workflow. Can you please provide more specific instructions?"
    })

  } catch (error) {
    console.error('Chat API error:', error)
    
    // Handle specific error types
    if (error instanceof z.ZodError) {
      return NextResponse.json(
        { error: 'Invalid request format', details: error.errors },
        { status: 400 }
      )
    }

    return NextResponse.json(
      { error: 'Failed to process chat message' },
      { status: 500 }
    )
  }
}