import { z } from 'zod'

// 用 zod 定義任務資料結構，擔任資料驗證的守門員
const TaskSchema = z.object({
  id: z.string(),
  description: z.string(),
  priority: z.enum(['P0', 'P1', 'P2']).optional(),
  assigned_to: z.string().optional(),
})
type Task = z.infer<typeof TaskSchema>

class SuperSimpleTaskManager {
  private tasks = new Map<string, Task>()
  private nextTaskId = 1

  createTask(description: string): Task {
    const id = `TASK-${String(this.nextTaskId).padStart(3, '0')}`
    this.nextTaskId += 1
    const task = TaskSchema.parse({ id, description })
    this.tasks.set(id, task)
    console.log(`DEBUG: Task created - ${id}: ${description}`)
    return task
  }

  updateTask(taskId: string, updates: Partial<Task>): Task | null {
    const task = this.tasks.get(taskId)
    if (!task) {
      console.log(`DEBUG: Task ${taskId} not found for update.`)
      return null
    }
    const updated = TaskSchema.parse({ ...task, ...updates })
    this.tasks.set(taskId, updated)
    console.log(`DEBUG: Task ${taskId} updated with ${JSON.stringify(updates)}`)
    return updated
  }

  listAllTasks(): string {
    if (this.tasks.size === 0) return 'No tasks in the system.'
    return (
      'Current Tasks:\n' +
      [...this.tasks.values()]
        .map(
          (t) =>
            `ID: ${t.id}, Desc: '${t.description}', ` +
            `Priority: ${t.priority ?? 'N/A'}, Assigned To: ${t.assigned_to ?? 'N/A'}`,
        )
        .join('\n')
    )
  }
}

import { DynamicStructuredTool } from '@langchain/core/tools'

const taskManager = new SuperSimpleTaskManager()

const createTaskTool = new DynamicStructuredTool({
  name: 'create_new_task',
  description: 'Use this first to create a new task and get its ID.',
  schema: z.object({ description: z.string() }),
  func: async ({ description }) => {
    const task = taskManager.createTask(description)
    return `Created task ${task.id}: '${task.description}'.`
  },
})

const assignPriorityTool = new DynamicStructuredTool({
  name: 'assign_priority_to_task',
  description: 'Use this to assign a priority to a task after it has been created.',
  schema: z.object({
    task_id: z.string(),
    priority: z.enum(['P0', 'P1', 'P2']),
  }),
  func: async ({ task_id, priority }) => {
    const task = taskManager.updateTask(task_id, { priority })
    return task
      ? `Assigned priority ${priority} to task ${task.id}.`
      : `Task ${task_id} not found.`
  },
})

const assignWorkerTool = new DynamicStructuredTool({
  name: 'assign_task_to_worker',
  description: 'Use this to assign a task to a specific worker after it has been created.',
  schema: z.object({ task_id: z.string(), worker_name: z.string() }),
  func: async ({ task_id, worker_name }) => {
    const task = taskManager.updateTask(task_id, { assigned_to: worker_name })
    return task
      ? `Assigned task ${task.id} to ${worker_name}.`
      : `Task ${task_id} not found.`
  },
})

const listTasksTool = new DynamicStructuredTool({
  name: 'list_all_tasks',
  description: 'Use this to list all current tasks and their status.',
  schema: z.object({}),
  func: async () => taskManager.listAllTasks(),
})

import { ChatOpenAI } from '@langchain/openai'
import { ChatPromptTemplate } from '@langchain/core/prompts'
import { createToolCallingAgent, AgentExecutor } from 'langchain/agents'
import { BufferMemory } from 'langchain/memory'

const llm = new ChatOpenAI({ temperature: 0.5, model: 'gpt-4o-mini' })

const systemPrompt = `
You are a focused Project Manager LLM agent. Your goal is to manage project tasks efficiently.
When you receive a new task request, follow these steps:
1. First, create the task with the given description using the create_new_task tool. You must do this first to get a task_id.
2. Next, analyze the user's request to see if a priority or an assignee is mentioned.
   - If a priority is mentioned (e.g., "urgent", "ASAP", "critical"), map it to P0. Use assign_priority_to_task.
   - If a worker is mentioned, use assign_task_to_worker.
3. If any information (priority, assignee) is missing, you must make a reasonable default assignment (e.g., assign P1 priority and assign to 'Worker A').
4. Once the task is fully processed, use list_all_tasks to show the final state.

Available workers: 'Worker A', 'Worker B', 'Review Team'
Priority levels: P0 (highest), P1 (medium), P2 (lowest)
`

const prompt = ChatPromptTemplate.fromMessages([
  ['system', systemPrompt],
  ['placeholder', '{chat_history}'],
  ['human', '{input}'],
  ['placeholder', '{agent_scratchpad}'],
])

const agent = await createToolCallingAgent({
  llm,
  tools: [createTaskTool, assignPriorityTool, assignWorkerTool, listTasksTool],
  prompt,
})

const executor = new AgentExecutor({
  agent,
  tools: [createTaskTool, assignPriorityTool, assignWorkerTool, listTasksTool],
  memory: new BufferMemory({ returnMessages: true, memoryKey: 'chat_history' }),
  inputKey: 'input',
  outputKey: 'output',
})

// 情境一：緊急任務，有指定負責人
await executor.invoke({
  input: "Create a task to implement a new login system. It's urgent and should be assigned to Worker B.",
})

// 情境二：一般任務，資訊缺漏
await executor.invoke({
  input: 'Manage a new task: Review marketing website content.',
})