import { tool } from "@opencode-ai/plugin"
import { rm, writeFile } from "node:fs/promises"
import path from "node:path"
import {
  buildWorkerPrompt,
  createTaskID,
  extractTextFromParts,
  formatTaskLine,
  now,
  projectRootFromContext,
  artifactPathForDisplay,
  artifactRoot,
  addTask,
  readTask,
  updateTask,
  loadRegistry,
  recordTaskError,
  recordTaskCompletion,
  errorMessage,
  sleep,
} from "../lib/team-utils.js"

const DEFAULT_WAIT_TIMEOUT_MS = 5 * 60 * 1000
const POLL_INTERVAL_MS = 1000

async function latestAssistantSummary(client, sessionID, directory) {
  const result = await client.session.messages({
    path: { id: sessionID },
    query: { directory },
  })

  const messages = result.data || []
  for (let index = messages.length - 1; index >= 0; index -= 1) {
    const message = messages[index]
    if (message.info.role !== "assistant") continue
    const text = extractTextFromParts(message.parts)
    if (!text) continue
    return {
      messageID: message.info.id,
      text,
    }
  }

  return {
    messageID: undefined,
    text: "",
  }
}

async function writeArtifact(worktree, taskID, text) {
  if (!text.trim()) return undefined
  const file = path.join(artifactRoot(worktree), `${taskID}.md`)
  await writeFile(file, `${text.trim()}\n`, "utf8")
  return file
}

async function getLiveSessionStatus(client, sessionID, directory) {
  const result = await client.session.status({
    query: { directory },
  })
  return result.data?.[sessionID] || { type: "idle" }
}

async function waitForIdle(client, sessionID, directory, timeoutMs, shouldStop) {
  const deadline = Date.now() + timeoutMs
  let lastStatus = { type: "idle" }

  while (Date.now() < deadline) {
    const stop = await shouldStop?.()
    if (stop) {
      return {
        completed: false,
        cancelled: true,
        status: lastStatus,
      }
    }

    lastStatus = await getLiveSessionStatus(client, sessionID, directory)
    if (!lastStatus || lastStatus.type === "idle") {
      return { completed: true, status: lastStatus || { type: "idle" } }
    }

    const sleepMs = Math.min(POLL_INTERVAL_MS, Math.max(0, deadline - Date.now()))
    if (sleepMs > 0) {
      await sleep(sleepMs)
    }
  }

  return { completed: false, cancelled: false, status: lastStatus }
}

async function dispatchPrompt(client, directory, sessionID, body) {
  await client.session.promptAsync({
    path: { id: sessionID },
    query: { directory },
    body,
  })
}

async function awaitTaskResult(client, context, projectRoot, taskID, timeoutMs) {
  const task = await readTask(projectRoot, taskID)
  if (!task) {
    throw new Error(`Unknown task_id: ${taskID}`)
  }

  if (task.status === "cancelled") {
    return {
      status: "cancelled",
      task,
    }
  }

  if (task.status === "error") {
    return {
      status: "error",
      task,
    }
  }

  if (task.status === "done") {
    return {
      status: "done",
      task,
      latest: {
        messageID: task.lastMessageID,
        text: task.summary || "",
      },
      artifactPath: task.artifactPath,
    }
  }

  const waited = await waitForIdle(
    client,
    task.childSessionID,
    context.directory,
    timeoutMs,
    async () => {
      const currentTask = await readTask(projectRoot, taskID)
      return currentTask?.status === "cancelled"
    },
  )
  if (waited.cancelled) {
    const cancelledTask = await readTask(projectRoot, taskID)
    return {
      status: "cancelled",
      task: cancelledTask || task,
    }
  }

  if (!waited.completed) {
    return {
      status: "running",
      task,
      liveStatus: waited.status.type,
      waitedMs: timeoutMs,
    }
  }

  const latest = await latestAssistantSummary(client, task.childSessionID, context.directory)

  if (!latest.messageID || latest.messageID === task.lastMessageID) {
    return {
      status: "ready",
      task,
    }
  }

  const artifactPath = await writeArtifact(projectRoot, task.id, latest.text)
  const completion = await recordTaskCompletion(projectRoot, task.id, latest, artifactPath)

  if (completion.cancelled) {
    if (artifactPath) {
      await rm(artifactPath, { force: true })
    }
    return {
      status: "cancelled",
      task: completion.task,
    }
  }

  return {
    status: "done",
    task: completion.task,
    latest,
    artifactPath,
  }
}

export const TeamOrchestratorPlugin = async ({ client }) => {
  return {
    tool: {
      team_delegate: tool({
        description: "Spawn a delegated worker session for a subtask and optionally wait for the result",
        args: {
          title: tool.schema.string().describe("Short task title"),
          agent: tool.schema.string().describe("Subagent to run, such as explore, general, plan, review, or verify"),
          prompt: tool.schema.string().describe("The delegated instruction for the worker"),
          goal: tool.schema.string().optional().describe("Optional concrete success criteria"),
          files: tool.schema.array(tool.schema.string()).optional().describe("Relevant files or directories"),
          wait: tool.schema.boolean().optional().describe("Wait for the worker to finish before returning"),
          timeout_ms: tool.schema.number().int().positive().optional().describe("Wait timeout in milliseconds when wait is true"),
        },
        async execute(args, context) {
          const projectRoot = projectRootFromContext(context)
          const wait = args.wait ?? false
          const timeoutMs = args.timeout_ms ?? DEFAULT_WAIT_TIMEOUT_MS
          const files = args.files ?? []
          const taskID = createTaskID()
          const createdAt = now()
          const created = await client.session.create({
            body: {
              parentID: context.sessionID,
              title: args.title,
            },
            query: {
              directory: context.directory,
            },
          })

          if (!created.data) {
            throw new Error(`Failed to create child session for \"${args.title}\"`)
          }

          const childSessionID = created.data.id
          const task = {
            id: taskID,
            title: args.title,
            agent: args.agent,
            goal: args.goal,
            files,
            prompt: args.prompt,
            parentSessionID: context.sessionID,
            childSessionID,
            status: "running",
            createdAt,
            updatedAt: createdAt,
            artifactPath: undefined,
            lastMessageID: undefined,
            summary: undefined,
            error: undefined,
          }
          await addTask(projectRoot, task)

          context.metadata({
            title: `Delegated ${args.agent}: ${args.title}`,
            metadata: {
              taskID,
              childSessionID,
              agent: args.agent,
            },
          })

          const handoffPrompt = buildWorkerPrompt({
            taskID,
            title: args.title,
            goal: args.goal,
            files,
            prompt: args.prompt,
          })

          try {
            await dispatchPrompt(client, context.directory, childSessionID, {
              agent: args.agent,
              parts: [{ type: "text", text: handoffPrompt }],
            })
          } catch (error) {
            await recordTaskError(projectRoot, taskID, error)
            throw error
          }

          if (!wait) {
            return [
              `▸ Delegated ${args.agent}: ${args.title}`,
              `  task: ${taskID}`,
              `  session: ${childSessionID}`,
              "",
              "Use `team_wait` to collect the result or `team_message` to continue.",
            ].join("\n")
          }

          const outcome = await awaitTaskResult(client, context, projectRoot, taskID, timeoutMs)
          if (outcome.status === "running") {
            return [
              `▸ ${args.title} (still running)`,
              `  task: ${taskID}`,
              `  agent: ${args.agent}`,
              `  session: ${childSessionID}`,
              `  waited: ${timeoutMs}ms`,
              "",
              "Use `team_wait` to collect the result later.",
            ].join("\n")
          }

          if (outcome.status === "cancelled") {
            return [
              `○ ${args.title} (cancelled)`,
              `  task: ${taskID}`,
              `  session: ${childSessionID}`,
            ].join("\n")
          }

          if (outcome.status === "error") {
            return [
              `✗ ${args.title} (error)`,
              `  task: ${taskID}`,
              `  session: ${childSessionID}`,
              `  error: ${outcome.task.error || "unknown error"}`,
            ].join("\n")
          }

          if (outcome.status === "ready") {
            return [
              `◦ ${args.title} (idle, no new result)`,
              `  task: ${taskID}`,
              `  session: ${childSessionID}`,
              "",
              "Use `team_wait` to collect the result or `team_message` to continue.",
            ].join("\n")
          }

          const artifactDisplay = artifactPathForDisplay(projectRoot, outcome.artifactPath)
          return [
            `✓ ${args.title}`,
            `  task: ${taskID}`,
            `  session: ${childSessionID}`,
            artifactDisplay ? `  artifact: ${artifactDisplay}` : "",
            "",
            outcome.latest.text || "The worker finished without a text summary.",
          ].filter(Boolean).join("\n")
        },
      }),
      team_message: tool({
        description: "Send a follow-up message to an existing delegated worker session",
        args: {
          task_id: tool.schema.string().describe("Task ID from team_delegate"),
          message: tool.schema.string().describe("Follow-up instruction for the existing worker"),
          wait: tool.schema.boolean().optional().describe("Wait for the worker to finish this follow-up"),
          timeout_ms: tool.schema.number().int().positive().optional().describe("Wait timeout in milliseconds when wait is true"),
        },
        async execute(args, context) {
          const projectRoot = projectRootFromContext(context)
          const wait = args.wait ?? false
          const timeoutMs = args.timeout_ms ?? DEFAULT_WAIT_TIMEOUT_MS
          const task = await readTask(projectRoot, args.task_id)
          if (!task) {
            throw new Error(`Unknown task_id: ${args.task_id}`)
          }
          if (task.parentSessionID !== context.sessionID) {
            throw new Error(`Task ${args.task_id} belongs to a different coordinator session`)
          }
          if (task.status === "cancelled") {
            throw new Error(`Cannot continue cancelled task: ${args.task_id}`)
          }

          context.metadata({
            title: `Message worker ${args.task_id}`,
            metadata: {
              taskID: args.task_id,
              childSessionID: task.childSessionID,
            },
          })

          await updateTask(projectRoot, args.task_id, currentTask => {
            if (currentTask.status === "cancelled") {
              throw new Error(`Cannot continue cancelled task: ${args.task_id}`)
            }
            currentTask.status = "running"
            currentTask.updatedAt = now()
            currentTask.summary = undefined
            currentTask.artifactPath = undefined
            currentTask.lastMessageID = undefined
            delete currentTask.error
            return { ...currentTask }
          })

          try {
            await dispatchPrompt(client, context.directory, task.childSessionID, {
              agent: task.agent,
              parts: [{ type: "text", text: args.message }],
            })
          } catch (error) {
            await recordTaskError(projectRoot, args.task_id, error)
            throw error
          }

          if (!wait) {
            return [
              `▸ Follow-up sent to ${args.task_id}`,
              `  session: ${task.childSessionID}`,
              `  agent: ${task.agent}`,
              "",
              "Use `team_wait` to collect the updated result.",
            ].join("\n")
          }

          const outcome = await awaitTaskResult(client, context, projectRoot, args.task_id, timeoutMs)
          if (outcome.status === "running") {
            return [
              `▸ ${args.task_id} (still running)`,
              `  session: ${task.childSessionID}`,
              `  agent: ${task.agent}`,
              `  waited: ${timeoutMs}ms`,
              "",
              "Use `team_wait` to collect the result later.",
            ].join("\n")
          }

          if (outcome.status === "cancelled") {
            return [
              `○ ${args.task_id} (cancelled)`,
              `  session: ${task.childSessionID}`,
            ].join("\n")
          }

          if (outcome.status === "error") {
            return [
              `✗ ${args.task_id} (error)`,
              `  session: ${task.childSessionID}`,
              `  error: ${outcome.task.error || "unknown error"}`,
            ].join("\n")
          }

          if (outcome.status === "ready") {
            return [
              `◦ ${args.task_id} (idle, no new result)`,
              `  session: ${task.childSessionID}`,
              "",
              "Use `team_message` to send a follow-up or `team_wait` again later.",
            ].join("\n")
          }

          const artifactDisplay = artifactPathForDisplay(projectRoot, outcome.artifactPath)
          return [
            `✓ ${args.task_id} responded`,
            `  session: ${task.childSessionID}`,
            artifactDisplay ? `  artifact: ${artifactDisplay}` : "",
            "",
            outcome.latest.text || "The worker finished without a text summary.",
          ].filter(Boolean).join("\n")
        },
      }),
      team_wait: tool({
        description: "Wait for a delegated worker session to become idle and fetch its latest summary",
        args: {
          task_id: tool.schema.string().describe("Task ID from team_delegate"),
          timeout_ms: tool.schema.number().int().positive().optional().describe("Maximum time to wait in milliseconds"),
        },
        async execute(args, context) {
          const projectRoot = projectRootFromContext(context)
          const timeoutMs = args.timeout_ms ?? DEFAULT_WAIT_TIMEOUT_MS
          const task = await readTask(projectRoot, args.task_id)
          if (!task) {
            throw new Error(`Unknown task_id: ${args.task_id}`)
          }
          if (task.parentSessionID !== context.sessionID) {
            throw new Error(`Task ${args.task_id} belongs to a different coordinator session`)
          }

          context.metadata({
            title: `Wait for ${args.task_id}`,
            metadata: {
              taskID: args.task_id,
              childSessionID: task.childSessionID,
            },
          })

          const outcome = await awaitTaskResult(client, context, projectRoot, args.task_id, timeoutMs)
          if (outcome.status === "running") {
            return [
              `▸ ${args.task_id} is running`,
              `  session: ${task.childSessionID}`,
              `  status: ${outcome.liveStatus}`,
              `  waited: ${timeoutMs}ms`,
            ].join("\n")
          }

          if (outcome.status === "cancelled") {
            return [
              `○ ${args.task_id} was cancelled`,
              `  session: ${task.childSessionID}`,
            ].join("\n")
          }

          if (outcome.status === "error") {
            return [
              `✗ ${args.task_id} is in an error state`,
              `  session: ${task.childSessionID}`,
              `  error: ${outcome.task.error || "unknown error"}`,
            ].join("\n")
          }

          if (outcome.status === "ready") {
            return [
              `◦ ${args.task_id} is idle (no new result)`,
              `  session: ${task.childSessionID}`,
              "",
              "Use `team_message` to send a follow-up or `team_wait` again later.",
            ].join("\n")
          }

          const artifactDisplay = artifactPathForDisplay(projectRoot, outcome.artifactPath)
          return [
            `✓ ${args.task_id} is idle`,
            `  session: ${task.childSessionID}`,
            artifactDisplay ? `  artifact: ${artifactDisplay}` : "",
            "",
            outcome.latest.text || "The worker is idle, but no assistant summary was found yet.",
          ].filter(Boolean).join("\n")
        },
      }),
      team_list: tool({
        description: "List delegated worker tasks for the current coordinator session",
        args: {
          status: tool.schema.enum(["running", "ready", "done", "cancelled", "error"]).optional().describe("Optional status filter"),
        },
        async execute(args, context) {
          const projectRoot = projectRootFromContext(context)
          const registry = await loadRegistry(projectRoot)
          const scopedTasks = registry.tasks.filter(task => task.parentSessionID === context.sessionID)

          const lines = []
          const statusCounts = {}
          for (const task of scopedTasks) {
            let liveStatus = task.status
            if (task.status === "running") {
              const sessionStatus = await getLiveSessionStatus(client, task.childSessionID, context.directory)
              if (sessionStatus.type === "idle") {
                liveStatus = task.summary ? "done" : "ready"
              } else {
                liveStatus = "running"
              }
            }

            statusCounts[liveStatus] = (statusCounts[liveStatus] || 0) + 1

            if (args.status && liveStatus !== args.status) {
              continue
            }

            lines.push(formatTaskLine(task, liveStatus))
          }

          if (lines.length === 0) {
            return args.status
              ? `No tasks matched status \"${args.status}\" for this coordinator session.`
              : "No delegated tasks found for this coordinator session."
          }

          const summaryParts = Object.entries(statusCounts)
            .sort(([,a], [,b]) => b - a)
            .map(([s, c]) => `${c} ${s}`)
            .join(", ")

          const shown = args.status ? `${lines.length} of ${scopedTasks.length}` : `${lines.length}`
          const summaryLabel = args.status ? "All tasks:" : ""
          const summaryLine = summaryLabel ? `${summaryLabel} ${summaryParts}` : summaryParts
          return [`Delegated tasks (${shown}):`, ...lines, "", summaryLine].join("\n")
        },
      }),
      team_cancel: tool({
        description: "Abort a delegated worker session and mark the task as cancelled",
        args: {
          task_id: tool.schema.string().describe("Task ID from team_delegate"),
        },
        async execute(args, context) {
          const projectRoot = projectRootFromContext(context)
          const task = await readTask(projectRoot, args.task_id)
          if (!task) {
            throw new Error(`Unknown task_id: ${args.task_id}`)
          }
          if (task.parentSessionID !== context.sessionID) {
            throw new Error(`Task ${args.task_id} belongs to a different coordinator session`)
          }
          if (task.status === "done" || task.status === "cancelled") {
            const icon = task.status === "done" ? "✓" : "○"
            return `${icon} ${args.task_id} is already ${task.status}; not cancelling.`
          }

          context.metadata({
            title: `Cancel ${args.task_id}`,
            metadata: {
              taskID: args.task_id,
              childSessionID: task.childSessionID,
            },
          })

          const liveStatus = await getLiveSessionStatus(client, task.childSessionID, context.directory)
          if (liveStatus.type === "idle" && task.summary) {
            return `✓ ${args.task_id} is already idle with a completed result; not cancelling.`
          }

          await client.session.abort({
            path: { id: task.childSessionID },
            query: { directory: context.directory },
          })

          const cancelResult = await updateTask(projectRoot, args.task_id, currentTask => {
            if (currentTask.status === "done" || currentTask.status === "cancelled") {
              return { alreadyTerminal: true, task: currentTask }
            }
            currentTask.status = "cancelled"
            currentTask.updatedAt = now()
            delete currentTask.error
            return { alreadyTerminal: false, task: { ...currentTask } }
          })

          if (cancelResult.alreadyTerminal) {
            const icon = cancelResult.task.status === "done" ? "✓" : "○"
            return `${icon} ${args.task_id} is already ${cancelResult.task.status}; no action taken.`
          }

          return [
            `○ Cancelled ${args.task_id}`,
            `  session: ${task.childSessionID}`,
            `  title: ${task.title}`,
          ].join("\n")
        },
      }),
    },
  }
}