codeg/src/components/chat/agent-plan-overlay.tsx

"use client"

import { memo, useMemo, useState } from "react"
import { Badge } from "@/components/ui/badge"
import { Button } from "@/components/ui/button"
import type { LiveMessage } from "@/contexts/acp-connections-context"
import type { PlanEntryInfo } from "@/lib/types"
import { cn } from "@/lib/utils"
import {
  CheckCircle2Icon,
  ChevronDownIcon,
  ChevronUpIcon,
  CircleDashedIcon,
  ListTodoIcon,
  Loader2Icon,
} from "lucide-react"

interface AgentPlanOverlayProps {
  message?: LiveMessage | null
  entries?: PlanEntryInfo[] | null
  planKey?: string | null
  visible?: boolean
  defaultExpanded?: boolean
}

function getLatestPlanEntries(message: LiveMessage | null): PlanEntryInfo[] {
  if (!message) return []

  for (let i = message.content.length - 1; i >= 0; i -= 1) {
    const block = message.content[i]
    if (block.type === "plan") {
      return block.entries
    }
  }

  return []
}

function getStatusLabel(status: string): string {
  switch (status) {
    case "completed":
      return "Completed"
    case "in_progress":
      return "In Progress"
    case "pending":
      return "Pending"
    default:
      return "Unknown"
  }
}

function getPriorityLabel(priority: string): string {
  switch (priority) {
    case "high":
      return "High"
    case "medium":
      return "Medium"
    case "low":
      return "Low"
    default:
      return "Unknown"
  }
}

function getPriorityClassName(priority: string): string {
  switch (priority) {
    case "high":
      return "text-red-700 bg-red-500/10 border-red-500/20 dark:text-red-300"
    case "medium":
      return "text-amber-700 bg-amber-500/10 border-amber-500/20 dark:text-amber-300"
    case "low":
      return "text-slate-700 bg-slate-500/10 border-slate-500/20 dark:text-slate-300"
    default:
      return "text-muted-foreground"
  }
}

function StatusIcon({ status }: { status: string }) {
  if (status === "completed") {
    return <CheckCircle2Icon className="h-3.5 w-3.5 text-emerald-500" />
  }

  if (status === "in_progress") {
    return <Loader2Icon className="h-3.5 w-3.5 text-blue-500 animate-spin" />
  }

  return <CircleDashedIcon className="h-3.5 w-3.5 text-muted-foreground" />
}

export const AgentPlanOverlay = memo(function AgentPlanOverlay({
  message,
  entries,
  planKey,
  visible = true,
  defaultExpanded = true,
}: AgentPlanOverlayProps) {
  const liveEntries = useMemo(
    () => getLatestPlanEntries(message ?? null),
    [message]
  )
  const resolvedEntries = useMemo(
    () => (liveEntries.length > 0 ? liveEntries : (entries ?? [])),
    [liveEntries, entries]
  )
  const hasPlan = visible && resolvedEntries.length > 0
  const fallbackPlanKey = useMemo(() => {
    if (resolvedEntries.length === 0) return null
    return resolvedEntries
      .map((entry) => `${entry.status}:${entry.priority}:${entry.content}`)
      .join("|")
  }, [resolvedEntries])
  const currentPlanKey = planKey ?? message?.id ?? fallbackPlanKey

  const completedCount = useMemo(
    () =>
      resolvedEntries.filter((entry) => entry.status === "completed").length,
    [resolvedEntries]
  )
  const hasIncompleteEntries = completedCount < resolvedEntries.length
  const resolvedDefaultExpanded = defaultExpanded && hasIncompleteEntries
  const currentPlanStateKey = currentPlanKey ?? "__plan__default__"
  const [collapsedByPlanKey, setCollapsedByPlanKey] = useState<
    Record<string, boolean>
  >({})
  const isExpanded = !(
    collapsedByPlanKey[currentPlanStateKey] ?? !resolvedDefaultExpanded
  )

  if (!hasPlan) {
    return null
  }

  if (!isExpanded) {
    return (
      <div className="pointer-events-none absolute right-8 top-4 z-20 flex">
        <Button
          type="button"
          variant="secondary"
          size="sm"
          className="cursor-pointer pointer-events-auto shadow-md bg-secondary/70 hover:bg-secondary"
          onClick={() =>
            setCollapsedByPlanKey((prev) => ({
              ...prev,
              [currentPlanStateKey]: false,
            }))
          }
        >
          <ListTodoIcon className="h-4 w-4" />
          Plan {completedCount}/{resolvedEntries.length}
          <ChevronUpIcon className="h-4 w-4" />
        </Button>
      </div>
    )
  }

  return (
    <div
      className="pointer-events-none absolute right-8 top-4 z-20 flex max-w-[min(22rem,calc(100%-2rem))]"
      data-plan-key={currentPlanKey ?? undefined}
    >
      <div className="pointer-events-auto w-72 max-w-full rounded-xl border bg-card/60 hover:bg-card/95 shadow-lg backdrop-blur transition-colors supports-[backdrop-filter]:bg-card/50 supports-[backdrop-filter]:hover:bg-card/85">
        <div className="flex items-center justify-between border-b px-3 py-2">
          <div className="flex items-center gap-2 min-w-0">
            <ListTodoIcon className="h-4 w-4 text-muted-foreground" />
            <span className="text-sm font-medium truncate">Agent Plan</span>
            <Badge variant="secondary" className="h-5">
              {completedCount}/{resolvedEntries.length}
            </Badge>
          </div>
          <Button
            type="button"
            variant="ghost"
            size="icon-xs"
            aria-label="Collapse plan"
            onClick={() =>
              setCollapsedByPlanKey((prev) => ({
                ...prev,
                [currentPlanStateKey]: true,
              }))
            }
          >
            <ChevronDownIcon className="h-4 w-4" />
          </Button>
        </div>

        <div className="max-h-96 overflow-y-auto p-3 space-y-2">
          {resolvedEntries.map((entry, index) => (
            <div
              key={`${entry.content}-${index}`}
              className="rounded-lg border bg-transparent px-2.5 py-2"
            >
              <div className="flex items-start gap-2">
                <StatusIcon status={entry.status} />
                <p
                  className={cn(
                    "min-w-0 flex-1 text-sm leading-5 break-words [overflow-wrap:anywhere]",
                    entry.status === "completed"
                      ? "text-muted-foreground line-through"
                      : "text-foreground"
                  )}
                >
                  {entry.content}
                </p>
              </div>
              <div className="mt-2 flex items-center gap-1.5 pl-5">
                <Badge variant="outline" className="h-5 text-[10px] uppercase">
                  {getStatusLabel(entry.status)}
                </Badge>
                <Badge
                  variant="outline"
                  className={cn(
                    "h-5 text-[10px] uppercase",
                    getPriorityClassName(entry.priority)
                  )}
                >
                  {getPriorityLabel(entry.priority)}
                </Badge>
              </div>
            </div>
          ))}
        </div>
      </div>
    </div>
  )
})