Task worker pool

import { tasklist } from '@crux/core/tasks'
import { taskWorker } from '@crux/core/tasks'
import { prompt } from '@crux/core'
import { agent } from '@crux/core/agent'
import { generate } from '@crux/ai'
import { openai } from '@ai-sdk/openai'
import { z } from 'zod'

// 1. Build the task list and seed it
export async function startBatch(documents: { id: string; text: string }[]) {
  const list = await tasklist({})
  for (const doc of documents) {
    await list.addTask({
      id: doc.id, // user-provided meaningful id
      label: `Summarize ${doc.id}`,
      description: doc.text.slice(0, 200),
    })
  }
  return list.id
}

// 2. The summarizer prompt — the worker context provides the assignment
const summarizePrompt = prompt({
  id: 'summarize-task',
  input: z.object({ text: z.string() }),
  system: `You are a summarization worker. Use your task tools to drive the lifecycle:
     1. Call startTask when you begin.
     2. Call completeTask with a summary when finished.
     3. Call failTask if you cannot complete the task.`,
  prompt: ({ input }) => `Summarize the following text:\n\n${input.text}`,
})

// 3. Drain the queue — one agent per todo task, N in parallel
export async function drainList(taskListId: string, documents: Map<string, string>, concurrency = 5) {
  const handle = createHandle(taskListId)
  const tasks = await handle.getTasks()
  const todos = tasks.filter((t) => t.status === 'todo')

  // Simple parallel drainer
  const queue = [...todos]
  const inflight = new Set<Promise<unknown>>()

  while (queue.length || inflight.size) {
    while (inflight.size < concurrency && queue.length) {
      const task = queue.shift()!
      const work = runOneTask(taskListId, task.id, documents.get(task.id) ?? '').finally(() => inflight.delete(work))
      inflight.add(work)
    }
    if (inflight.size) await Promise.race(inflight)
  }
}

async function runOneTask(taskListId: string, taskId: string, text: string) {
  const worker = taskWorker(taskListId, taskId)

  const reviewerAgent = agent({
    id: `worker-${taskId}`,
    prompt: summarizePrompt,
    tools: worker.asTools(), // startTask, reportProgress, completeTask, failTask
  })

  await generate(agent.prompt, {
    model: openai('gpt-4o-mini'),
    input: { text },
    use: [worker.asContext()], // injects the assignment + guidelines
    stopWhen: { stepCountIs: 5 },
  })
}

'use client'
import { useTaskList, useTasks } from '@crux/react'

export function BatchProgress({ taskListId }: { taskListId: string }) {
  const list = useTaskList(taskListId)
  const tasks = useTasks(taskListId)

  if (!list) return null

  const counts = {
    todo: tasks.filter((t) => t.status === 'todo').length,
    in_progress: tasks.filter((t) => t.status === 'in_progress').length,
    completed: tasks.filter((t) => t.status === 'completed').length,
    failed: tasks.filter((t) => t.status === 'failed').length,
  }

  return (
    <div>
      <p>
        {counts.completed} / {tasks.length} done • {counts.in_progress} running • {counts.failed} failed
      </p>
      <ul>
        {tasks.map((t) => (
          <li key={t.id} className={`task task-${t.status}`}>
            {t.label} — <code>{t.status}</code>
            {t.progress && <p className="muted">{t.progress}</p>}
            {t.status === 'failed' && t.error && <pre>{t.error}</pre>}
          </li>
        ))}
      </ul>
    </div>
  )
}

const task = await handle.getTasks().then((ts) => ts.find((t) => t.id === taskId))
const agent = task?.assignee?.agent === 'classifier' ? classifierAgent : summarizerAgent