import { DynamicStructuredTool } from "@langchain/core/tools";
import { z } from "zod";

interface Task {
  id: number;
  description: string;
  status: "pending" | "in_progress" | "completed";
  createdAt: Date;
  completedAt?: Date;
}

class TaskTracker {
  private tasks: Task[] = [];
  private nextId = 1;

  addTask(description: string): Task {
    const task: Task = {
      id: this.nextId++,
      description,
      status: "pending",
      createdAt: new Date(),
    };
    this.tasks.push(task);
    return task;
  }

  updateTaskStatus(id: number, status: Task["status"]): Task | null {
    const task = this.tasks.find((t) => t.id === id);
    if (task) {
      task.status = status;
      if (status === "completed") {
        task.completedAt = new Date();
      }
      return task;
    }
    return null;
  }

  getTasks(): Task[] {
    return [...this.tasks];
  }

  getActiveTasks(): Task[] {
    return this.tasks.filter((t) => t.status !== "completed");
  }

  clear(): void {
    this.tasks = [];
    this.nextId = 1;
  }

  formatTaskList(): string {
    if (this.tasks.length === 0) {
      return "No tasks in the list.";
    }

    const statusEmoji = {
      pending: "⏳",
      in_progress: "🔄",
      completed: "✅",
    };

    return this.tasks
      .map(
        (t) =>
          `${statusEmoji[t.status]} [${t.id}] ${t.description} (${t.status})`,
      )
      .join("\n");
  }
}

const taskTracker = new TaskTracker();

export const planTasksTool = new DynamicStructuredTool({
  name: "plan_tasks",
  description:
    "Plan and break down a complex task into smaller steps. Use this BEFORE starting any multi-step work to organize your approach.",
  schema: z.object({
    tasks: z
      .array(z.string())
      .min(1)
      .describe(
        "List of task descriptions in order of execution. Keep each task focused and achievable with a single tool call.",
      ),
  }),
  func: async (input: { tasks: string[] }) => {
    taskTracker.clear();
    const createdTasks = input.tasks.map((desc) => taskTracker.addTask(desc));

    return `Task plan created with ${createdTasks.length} tasks:\n${taskTracker.formatTaskList()}\n\nStart with task 1 and mark it as in_progress when you begin.`;
  },
});

export const updateTaskTool = new DynamicStructuredTool({
  name: "update_task",
  description:
    "Update the status of a task. Mark as 'in_progress' when starting, 'completed' when done.",
  schema: z.object({
    taskId: z.number().min(1).describe("The task ID to update"),
    status: z
      .enum(["pending", "in_progress", "completed"])
      .describe("The new status for the task"),
  }),
  func: async (input: { taskId: number; status: Task["status"] }) => {
    const task = taskTracker.updateTaskStatus(input.taskId, input.status);

    if (!task) {
      return `Error: Task ${input.taskId} not found.`;
    }

    const activeTasks = taskTracker.getActiveTasks();
    const nextTask = activeTasks.find((t) => t.status === "pending");

    let response = `Task ${task.id} marked as ${task.status}: "${task.description}"`;

    if (input.status === "completed" && nextTask) {
      response += `\n\nNext task: [${nextTask.id}] ${nextTask.description}`;
    } else if (input.status === "completed" && activeTasks.length === 0) {
      response += "\n\nAll tasks completed! 🎉";
    }

    response += `\n\nCurrent task list:\n${taskTracker.formatTaskList()}`;
    return response;
  },
});

export const viewTasksTool = new DynamicStructuredTool({
  name: "view_tasks",
  description: "View the current task list and their statuses.",
  schema: z.object({}),
  func: async () => {
    return taskTracker.formatTaskList();
  },
});

export const taskTrackerTools = [planTasksTool, updateTaskTool, viewTasksTool];