hp-prod-tracker/src/app/api/chat/route.ts

import { NextRequest, NextResponse } from "next/server";
import { serverError } from "@/lib/api-utils";
import { requireAuth } from "@/lib/rbac/require-auth";
import { chat, buildSystemPrompt, getProviderStatus } from "@/lib/chat/provider";
import { executeTool } from "@/lib/chat/tool-executor";
import { prisma } from "@/lib/prisma";
import { MUTATION_TOOLS } from "@/lib/chat/tool-executor";

// ─── Rate Limiter ──────────────────────────────────────
const RATE_LIMIT = 20; // requests per window
const RATE_WINDOW_MS = 60_000; // 1 minute
const rateLimitMap = new Map<string, { count: number; resetAt: number }>();

function checkRateLimit(userId: string): boolean {
  const now = Date.now();
  const entry = rateLimitMap.get(userId);
  if (!entry || now > entry.resetAt) {
    rateLimitMap.set(userId, { count: 1, resetAt: now + RATE_WINDOW_MS });
    return true;
  }
  entry.count++;
  return entry.count <= RATE_LIMIT;
}

// ─── Entity Card Extraction ────────────────────────────
// Extracts navigable entity cards from tool results so the
// chat panel can render clickable cards like the old Smart Search.

export interface EntityCard {
  id: string;
  type: "project" | "deliverable" | "user" | "stage";
  name: string;
  status: string;
  priority?: string;
  code?: string;
  projectName?: string;
  projectId?: string;
  link: string;
  relevanceScore?: number;
}

/** Tools whose results may contain navigable entities */
const ENTITY_PRODUCING_TOOLS = new Set([
  "search_entities",
  "list_projects",
  "list_deliverables",
  "list_overdue",
  "get_blocked_stages",
  "get_project",
  "get_workload",
  "list_users",
  "get_available_artists",
]);

/**
 * Extract entity cards from a tool's result data.
 * Handles various return shapes from different tools.
 */
function extractEntities(toolName: string, data: any): EntityCard[] {
  if (!data || !ENTITY_PRODUCING_TOOLS.has(toolName)) return [];

  // If data is a "no results" message object, skip
  if (data.message && !Array.isArray(data)) return [];

  // ── Special: get_workload returns { users: [...], weeks: [...] }
  if (toolName === "get_workload" && data.users) {
    return extractWorkloadEntities(data);
  }

  // ── Special: list_users / get_available_artists return arrays of user objects
  if (toolName === "list_users" || toolName === "get_available_artists") {
    return extractUserEntities(Array.isArray(data) ? data : [data]);
  }

  const items = Array.isArray(data) ? data : [data];
  const entities: EntityCard[] = [];

  for (const item of items) {
    if (!item || typeof item !== "object" || !item.id) continue;

    // Determine entity type from the item's type field or the tool name
    const itemType =
      item.type ||
      (toolName === "list_projects" || toolName === "get_project"
        ? "project"
        : toolName === "list_deliverables"
          ? "deliverable"
          : null);

    if (itemType === "project") {
      entities.push({
        id: item.id,
        type: "project",
        name: item.name || "Unnamed Project",
        status: item.status || "",
        priority: item.priority,
        code: item.projectCode || item.code,
        link: `/projects/${item.id}`,
        relevanceScore: item._relevanceScore,
      });
    } else if (itemType === "deliverable") {
      const projId = item.projectId || item.project?.id;
      entities.push({
        id: item.id,
        type: "deliverable",
        name: item.name || "Unnamed Deliverable",
        status: item.status || "",
        priority: item.priority,
        code: item.projectCode || item.code || item.project?.projectCode,
        projectName: item.projectName || item.project?.name,
        projectId: projId,
        link: projId
          ? `/projects/${projId}/deliverables/${item.id}`
          : `/projects`,
        relevanceScore: item._relevanceScore,
      });
    } else if (itemType === "user") {
      entities.push({
        id: item.id,
        type: "user",
        name: item.name || "Unnamed User",
        status: item.department || item.role || "",
        link: `/workload`,
        relevanceScore: item._relevanceScore,
      });
    } else if (
      toolName === "list_overdue" ||
      toolName === "get_blocked_stages"
    ) {
      const delivId = item.deliverableId || item.deliverable?.id;
      const projId = item.projectId || item.project?.id;
      if (delivId && projId) {
        entities.push({
          id: delivId,
          type: "deliverable",
          name:
            item.deliverableName ||
            item.deliverable?.name ||
            item.name ||
            "Deliverable",
          status: item.stageStatus || item.status || "",
          priority: item.priority,
          code: item.projectCode || item.project?.projectCode,
          projectName: item.projectName || item.project?.name,
          projectId: projId,
          link: `/projects/${projId}/deliverables/${delivId}`,
        });
      }
    }
  }

  return entities;
}

/**
 * Extract project entities from workload data.
 * Only extracts when the result is a targeted query (≤3 active users).
 * For org-wide workload dumps, returns nothing — too noisy for cards.
 * The search_entities tool already provides user cards.
 */
function extractWorkloadEntities(data: {
  users: any[];
  weeks: any[];
}): EntityCard[] {
  const activeUsers = (data.users || []).filter(
    (u: any) => (u.totalActiveAssignments || 0) > 0
  );

  // If this is an org-wide query (many active users), skip extraction entirely
  if (activeUsers.length > 3) return [];

  const entities: EntityCard[] = [];

  for (const user of activeUsers) {
    // Add user card
    if (user.userId) {
      entities.push({
        id: user.userId,
        type: "user",
        name: user.userName || "User",
        status: user.department || user.role || "",
        link: `/workload`,
      });
    }

    // Extract unique projects from their assignments
    const seenProjects = new Set<string>();
    for (const week of user.weeks || []) {
      for (const assignment of week.assignments || []) {
        if (assignment.projectId && !seenProjects.has(assignment.projectId)) {
          seenProjects.add(assignment.projectId);
          entities.push({
            id: assignment.projectId,
            type: "project",
            name: assignment.projectName || "Project",
            status: assignment.stageStatus || "",
            priority: assignment.priority,
            link: `/projects/${assignment.projectId}`,
            relevanceScore: 0.9,
          });
        }
      }
    }
  }

  return entities;
}

/**
 * Extract user entities from list_users / get_available_artists.
 * Cap at 6 to avoid flooding the chat with dozens of artist cards.
 */
function extractUserEntities(users: any[]): EntityCard[] {
  return users
    .filter((u) => u && u.id)
    .slice(0, 6)
    .map((u) => ({
      id: u.id,
      type: "user" as const,
      name: u.name || "User",
      status: u.department || u.role || "",
      link: `/workload`,
    }));
}

/** Deduplicate entities by id, keeping the first occurrence */
function deduplicateEntities(entities: EntityCard[]): EntityCard[] {
  const seen = new Set<string>();
  return entities.filter((e) => {
    if (seen.has(e.id)) return false;
    seen.add(e.id);
    return true;
  });
}

/**
 * SSE helper — writes a JSON event to the stream.
 */
function sendEvent(
  controller: ReadableStreamDefaultController,
  encoder: TextEncoder,
  event: string,
  data: any
) {
  controller.enqueue(
    encoder.encode(`event: ${event}\ndata: ${JSON.stringify(data)}\n\n`)
  );
}

export async function POST(req: NextRequest) {
  try {
    const { session, error } = await requireAuth("PROJECT_VIEW");
    if (error) return error;

    // Rate limit per user
    if (!checkRateLimit(session.user.id)) {
      return NextResponse.json(
        { error: "Too many requests — please wait a moment" },
        { status: 429 }
      );
    }

    const { messages: rawMessages, context } = await req.json();

    if (!rawMessages || !Array.isArray(rawMessages) || rawMessages.length === 0) {
      return NextResponse.json(
        { error: "messages array is required" },
        { status: 400 }
      );
    }

    // Limit conversation history to prevent exceeding Claude's context window.
    // Keep the last 20 messages (10 user + 10 assistant turns).
    const MAX_HISTORY = 20;
    const messages = rawMessages.length > MAX_HISTORY
      ? rawMessages.slice(-MAX_HISTORY)
      : rawMessages;

    const userId = session.user.id;
    const organizationId = session.user.organizationId;
    const userRole = session.user.role;

    // -- Build system prompt with optional page context --
    let pageContext: string | undefined;

    if (context?.activeProjectId) {
      try {
        const project = await prisma.project.findUnique({
          where: { id: context.activeProjectId },
          select: { name: true, projectCode: true, status: true, id: true },
        });

        if (project) {
          pageContext = `The user is currently viewing project "${project.name}" (code: ${project.projectCode}, ID: ${project.id}, status: ${project.status}).`;

          if (context.activeDeliverableId) {
            const deliverable = await prisma.deliverable.findUnique({
              where: { id: context.activeDeliverableId },
              select: { name: true, status: true, id: true },
            });
            if (deliverable) {
              pageContext += ` They are also viewing deliverable "${deliverable.name}" (ID: ${deliverable.id}, status: ${deliverable.status}).`;
            }
          }

          pageContext += " Use these IDs directly when the user's request clearly refers to what they're looking at — no need to search.";
        }
      } catch {
        // Non-critical — proceed without page context
      }
    }

    const systemPrompt = buildSystemPrompt(pageContext);

    // -- Stream SSE response --
    const encoder = new TextEncoder();

    const stream = new ReadableStream({
      async start(controller) {
        try {
          let response = await chat(messages, systemPrompt);

          // Collect all cache keys to invalidate and entities to display
          const allInvalidateKeys: string[] = [];
          const allEntities: EntityCard[] = [];

          // Tool-calling loop with streaming status updates
          let conversationMessages = [...messages];
          let iterations = 0;
          const MAX_ITERATIONS = 10;

          // Check if this request is a mutation confirmation from the client
          const lastUserMsg = messages[messages.length - 1];
          const isConfirmation = typeof lastUserMsg?.content === "string" &&
            lastUserMsg.content.startsWith("[MUTATION_CONFIRMED]");
          const isCancellation = typeof lastUserMsg?.content === "string" &&
            lastUserMsg.content.startsWith("[MUTATION_CANCELLED]");

          while (response.toolCalls.length > 0 && iterations < MAX_ITERATIONS) {
            iterations++;

            const toolResults: { tool_use_id: string; content: string }[] = [];

            for (const toolCall of response.toolCalls) {
              const isMutationTool = MUTATION_TOOLS.has(toolCall.name);

              // If this is a mutation and NOT a confirmed re-execution, pause for confirmation
              if (isMutationTool && !isConfirmation) {
                sendEvent(controller, encoder, "tool_start", {
                  toolName: toolCall.name,
                  toolCallId: toolCall.id,
                  input: toolCall.input,
                });

                // Send pending event with tool details so client can show confirmation UI
                sendEvent(controller, encoder, "mutation_pending", {
                  toolName: toolCall.name,
                  toolCallId: toolCall.id,
                  input: toolCall.input,
                  // Include Claude's text response so client can display it
                  assistantMessage: response.content || null,
                });

                sendEvent(controller, encoder, "done", {});
                controller.close();
                return;
              }

              // Stream tool status to client
              sendEvent(controller, encoder, "tool_start", {
                toolName: toolCall.name,
                toolCallId: toolCall.id,
                input: toolCall.input,
              });

              const result = await executeTool(toolCall.name, toolCall.input, {
                organizationId,
                userId,
                userRole,
              });

              if (result.invalidateKeys.length > 0) {
                allInvalidateKeys.push(...result.invalidateKeys);
              }

              // Extract navigable entities from tool results
              if (result.success && result.data) {
                const extracted = extractEntities(toolCall.name, result.data);
                allEntities.push(...extracted);
              }

              sendEvent(controller, encoder, "tool_end", {
                toolName: toolCall.name,
                toolCallId: toolCall.id,
                success: result.success,
                error: result.error,
              });

              // Truncate large tool results to avoid exceeding context limits
              let resultJson = JSON.stringify(
                result.success ? result.data : { error: result.error }
              );
              if (resultJson.length > 8000) {
                resultJson = resultJson.slice(0, 8000) + '... [truncated]"}';
              }

              toolResults.push({
                tool_use_id: toolCall.id,
                content: resultJson,
              });
            }

            // Build assistant message with text + tool_use blocks
            const assistantContentBlocks: any[] = [];
            if (response.content) {
              assistantContentBlocks.push({ type: "text", text: response.content });
            }
            for (const tc of response.toolCalls) {
              assistantContentBlocks.push({
                type: "tool_use",
                id: tc.id,
                name: tc.name,
                input: tc.input,
              });
            }

            // Append assistant turn (with tool_use blocks) AND the matching
            // tool_result user turn so the next iteration sees a complete
            // conversation. Claude requires every tool_use to be immediately
            // followed by its tool_result in the next message.
            conversationMessages = [
              ...conversationMessages,
              {
                role: "assistant" as const,
                content: assistantContentBlocks,
              },
              {
                role: "user" as const,
                content: toolResults.map((tr) => ({
                  type: "tool_result",
                  tool_use_id: tr.tool_use_id,
                  content: tr.content,
                })),
              },
            ];

            // Continue conversation — tool results are already in conversationMessages,
            // so don't pass them separately.
            response = await chat(conversationMessages, systemPrompt);
          }

          // Parse suggestions from Claude's response if present
          let finalContent = response.content;
          let suggestions: any[] | undefined;

          const suggestionsMatch = finalContent.match(
            /```suggestions\s*\n([\s\S]*?)\n```/
          );
          if (suggestionsMatch) {
            try {
              suggestions = JSON.parse(suggestionsMatch[1]);
              // Remove the suggestions block from the displayed text
              finalContent = finalContent
                .replace(/```suggestions\s*\n[\s\S]*?\n```/, "")
                .trim();
            } catch {
              // If JSON is malformed, leave content as-is
            }
          }

          // Deduplicate, prioritize by relevance score first, then type (projects/deliverables first, then users), cap at 10
          const TYPE_PRIORITY: Record<string, number> = {
            project: 0,
            deliverable: 1,
            stage: 2,
            user: 3,
          };
          const entities = deduplicateEntities(allEntities)
            .sort((a, b) => {
              const scoreA = a.relevanceScore ?? 0;
              const scoreB = b.relevanceScore ?? 0;
              if (scoreA !== scoreB) {
                return scoreB - scoreA;
              }
              return (TYPE_PRIORITY[a.type] ?? 9) - (TYPE_PRIORITY[b.type] ?? 9);
            })
            .slice(0, 10);

          // Send final response
          sendEvent(controller, encoder, "message", {
            content: finalContent,
            provider: response.provider,
            invalidateKeys: [...new Set(allInvalidateKeys)],
            ...(suggestions && suggestions.length > 0 ? { suggestions } : {}),
            ...(entities.length > 0 ? { entities } : {}),
          });

          sendEvent(controller, encoder, "done", {});
          controller.close();
        } catch (err) {
          sendEvent(controller, encoder, "error", {
            error: err instanceof Error ? err.message : "Internal server error",
          });
          controller.close();
        }
      },
    });

    return new Response(stream, {
      headers: {
        "Content-Type": "text/event-stream",
        "Cache-Control": "no-cache",
        Connection: "keep-alive",
      },
    });
  } catch (error) {
    console.error("[Chat API]", error);
    return serverError(error);
  }
}

/**
 * GET /api/chat — returns provider health status.
 */
export async function GET() {
  try {
    const { session, error } = await requireAuth("PROJECT_VIEW");
    if (error) return error;

    const status = await getProviderStatus();
    return NextResponse.json(status);
  } catch (error) {
    return serverError(error);
  }
}