route.ts

import { NextRequest } from 'next/server';
import { appendSessionMessages } from '@/lib/agent-memory';
import { globalEventEmitter } from '@/lib/event-emitter';

// Type definitions for streaming events
interface AgentEventData {
  sessionId: string;
  status?: string;
  message?: string;
  content?: string;
  fullContent?: string;
  tool?: string;
  response?: string;
  sandboxUrl?: string;
  error?: string;
  hasSandbox?: boolean;
  sandboxId?: string;
  isNew?: boolean;
}

// Helper function to create SSE messages
function createSSEMessage(type: string, data: Record<string, unknown>): string {
  return `data: ${JSON.stringify({ type, data })}\n\n`;
}

export async function GET(request: NextRequest) {
  const { searchParams } = new URL(request.url);
  const sessionId = searchParams.get('sessionId');

  if (!sessionId) {
    return new Response('Missing sessionId parameter', { status: 400 });
  }

  // Create a ReadableStream for Server-Sent Events
  const stream = new ReadableStream({
    start(controller) {
      // Send headers to establish SSE connection
      const encoder = new TextEncoder();
      
      let isClosed = false;
      
      const send = (data: string) => {
        if (isClosed) {
          return;
        }
        try {
          controller.enqueue(encoder.encode(data));
        } catch (error) {
          console.error('[SSE Stream] Error sending data:', error);
          isClosed = true;
        }
      };

      // Initial connection message
      send(createSSEMessage('connected', { sessionId }));

      // Event handlers for different types of updates
      const handleAgentUpdate = (data: AgentEventData) => {
        if (data.sessionId === sessionId) {
          send(createSSEMessage('status', {
            sessionId: data.sessionId,
            status: data.status,
            message: data.message,
            hasSandbox: data.hasSandbox,
          }));
        }
      };

      const handlePartialContent = (data: AgentEventData) => {
        if (data.sessionId === sessionId) {
          send(createSSEMessage('partial', {
            sessionId: data.sessionId,
            content: data.content,
            fullContent: data.fullContent,
          }));
        }
      };

      const handleToolUsed = (data: AgentEventData & { args?: Record<string, unknown>; result?: string; status?: string }) => {
        if (data.sessionId === sessionId) {
          send(createSSEMessage('tool', {
            sessionId: data.sessionId,
            tool: data.tool,
            args: data.args,
            result: data.result,
            status: data.status,
          }));
        }
      };

      const handleSandboxStatus = (data: AgentEventData) => {
        if (data.sessionId === sessionId) {
          send(createSSEMessage('sandbox', {
            sessionId: data.sessionId,
            sandboxId: data.sandboxId,
            sandboxUrl: data.sandboxUrl,
            isNew: data.isNew,
          }));
        }
      };

      const handleComplete = (data: AgentEventData) => {
        if (data.sessionId === sessionId) {
          send(createSSEMessage('complete', {
            sessionId: data.sessionId,
            response: data.response,
            sandboxUrl: data.sandboxUrl,
            hasSandbox: data.hasSandbox,
          }));
          if (data.response) {
            appendSessionMessages(sessionId, [{ role: 'ai', content: String(data.response), ts: Date.now() }]);
          }
        }
      };

      const handleError = (data: AgentEventData) => {
        if (data.sessionId === sessionId) {
          send(createSSEMessage('error', {
            sessionId: data.sessionId,
            error: data.error,
          }));
        }
      };

      const handleReasoning = (data: AgentEventData & { reasoning?: string }) => {
        if (data.sessionId === sessionId) {
          send(createSSEMessage('reasoning', {
            sessionId: data.sessionId,
            reasoning: data.reasoning,
          }));
        }
      };

      const handleFileUpdate = (data: AgentEventData & { filePath?: string; content?: string; action?: 'start' | 'update' | 'complete' }) => {
        if (data.sessionId === sessionId) {
          send(createSSEMessage('file_update', {
            sessionId: data.sessionId,
            filePath: data.filePath,
            content: data.content,
            action: data.action,
          }));
        }
      };

      const handleCodePatch = (data: AgentEventData & { filePath?: string; content?: string; action?: 'start' | 'patch' | 'complete' }) => {
        if (data.sessionId === sessionId) {
          send(createSSEMessage('code_patch', {
            sessionId: data.sessionId,
            filePath: data.filePath,
            content: data.content,
            action: data.action,
          }));
        }
      };

      const handleFileTreeSync = (data: AgentEventData & { fileTree?: unknown }) => {
        if (data.sessionId === sessionId) {
          send(createSSEMessage('file_tree_sync', {
            sessionId: data.sessionId,
            fileTree: data.fileTree,
          }));
        }
      };

      // Listen for different types of events
      globalEventEmitter.on('agent:status', handleAgentUpdate);
      globalEventEmitter.on('agent:partial', handlePartialContent);
      globalEventEmitter.on('agent:tool', handleToolUsed);
      globalEventEmitter.on('agent:sandbox', handleSandboxStatus);
      globalEventEmitter.on('agent:complete', handleComplete);
      globalEventEmitter.on('agent:error', handleError);
      globalEventEmitter.on('agent:reasoning', handleReasoning);
      globalEventEmitter.on('agent:fileUpdate', handleFileUpdate);
      globalEventEmitter.on('agent:codePatch', handleCodePatch);
      globalEventEmitter.on('agent:fileTreeSync', handleFileTreeSync);

      // Keep connection alive with periodic heartbeat
      const heartbeat = setInterval(() => {
        send(createSSEMessage('heartbeat', { timestamp: Date.now() }));
      }, 30000);

      // Cleanup when client disconnects
      request.signal.addEventListener('abort', () => {
        isClosed = true;
        clearInterval(heartbeat);
        globalEventEmitter.off('agent:status', handleAgentUpdate);
        globalEventEmitter.off('agent:partial', handlePartialContent);
        globalEventEmitter.off('agent:tool', handleToolUsed);
        globalEventEmitter.off('agent:sandbox', handleSandboxStatus);
        globalEventEmitter.off('agent:complete', handleComplete);
        globalEventEmitter.off('agent:error', handleError);
        globalEventEmitter.off('agent:reasoning', handleReasoning);
        globalEventEmitter.off('agent:fileUpdate', handleFileUpdate);
        globalEventEmitter.off('agent:codePatch', handleCodePatch);
        globalEventEmitter.off('agent:fileTreeSync', handleFileTreeSync);
        try {
          controller.close();
        } catch {
          // Controller already closed
        }
      });
    },
  });

  return new Response(stream, {
    headers: {
      'Content-Type': 'text/event-stream',
      'Cache-Control': 'no-cache',
      'Connection': 'keep-alive',
      'Access-Control-Allow-Origin': '*',
      'Access-Control-Allow-Headers': 'Content-Type',
    },
  });
}