import { StringEnum } from "@earendil-works/pi-ai";
import type { ExtensionAPI, ExtensionContext, Theme } from "@earendil-works/pi-coding-agent";
import { Text } from "@earendil-works/pi-tui";
import { Type } from "typebox";
import type { TaskState } from "./state.ts";
import type { AddTaskDraft, Task, TaskPatch, TaskToolDetails } from "./types.ts";
import { TaskValidationError } from "./types.ts";
import { orderedTreeRows, renderToolResult } from "./ui.ts";

const TaskStatusSchema = StringEnum(["pending", "in-progress", "blocked", "done"] as const, {
  description: "Task status. Defaults to pending when omitted during add.",
});

const AddTaskSchema = Type.Cyclic(
  {
    AddTask: Type.Object(
      {
        name: Type.String({ minLength: 1 }),
        description: Type.String({ minLength: 1 }),
        status: Type.Optional(TaskStatusSchema),
        parentId: Type.Optional(
          Type.Integer({ minimum: 1, description: "Existing parent ID. Allowed only on top-level tasks." }),
        ),
        children: Type.Optional(Type.Array(Type.Ref("AddTask"))),
      },
      { additionalProperties: false },
    ),
  },
  "AddTask",
);

export const AddTaskParams = Type.Object(
  {
    tasks: Type.Array(AddTaskSchema, { minItems: 1, description: "Tasks to add in one atomic batch." }),
  },
  { additionalProperties: false },
);

const TaskPatchSchema = Type.Object(
  {
    id: Type.Integer({ minimum: 1 }),
    name: Type.Optional(Type.String({ minLength: 1 })),
    description: Type.Optional(Type.String({ minLength: 1 })),
    status: Type.Optional(TaskStatusSchema),
    parentId: Type.Optional(
      Type.Union([Type.Integer({ minimum: 1 }), Type.Null()], {
        description: "New parent ID, or null to move the task to the root.",
      }),
    ),
    position: Type.Optional(
      Type.Integer({ minimum: 0, description: "Zero-based insertion index among the task's final siblings." }),
    ),
  },
  { additionalProperties: false },
);

export const UpdateTaskParams = Type.Object(
  {
    tasks: Type.Array(TaskPatchSchema, { minItems: 1, description: "Task patches applied atomically in array order." }),
  },
  { additionalProperties: false },
);

export const IdsParams = Type.Object(
  {
    ids: Type.Array(Type.Integer({ minimum: 1 }), {
      minItems: 1,
      description: "Task IDs to process in one call. Duplicate IDs are deduplicated.",
    }),
  },
  { additionalProperties: false },
);

export const ListTaskParams = Type.Object({}, { additionalProperties: false });

interface AddTaskInput {
  tasks: AddTaskDraft[];
}

interface UpdateTaskInput {
  tasks: TaskPatch[];
}

interface IdsInput {
  ids: number[];
}

type ChangeHandler = (ctx: ExtensionContext) => void;

function taskIds(tasks: Task[]): string {
  return tasks.map((task) => task.id).join(", ");
}

function compactTree(tasks: Task[]): string {
  if (tasks.length === 0) return "No tasks.";
  return orderedTreeRows(tasks)
    .map(({ task, prefix }) => `${prefix}${task.id} [${task.status}] ${task.name}`)
    .join("\n");
}

function details(action: TaskToolDetails["action"], state: TaskState, affected: Task[]): TaskToolDetails {
  return { action, snapshot: state.snapshot(), affected };
}

function mutationText(verb: string, affected: Task[], state: TaskState): string {
  const stats = state.stats();
  return `${verb} ${affected.length} task${affected.length === 1 ? "" : "s"}: ${taskIds(affected)}. ${stats.done}/${stats.total} done.`;
}

function throwToolError(label: string, error: unknown): never {
  if (error instanceof TaskValidationError) {
    const count = error.errors.length;
    throw new Error(
      `${label} failed · ${count} validation error${count === 1 ? "" : "s"}\n${error.errors.map((item) => `- ${item}`).join("\n")}`,
    );
  }
  throw error;
}

function renderCall(title: string, count: number | undefined, theme: Theme): Text {
  const suffix = count === undefined ? "" : theme.fg("dim", ` (${count})`);
  return new Text(theme.fg("toolTitle", theme.bold(title)) + suffix, 0, 0);
}

function notifyChange(onChange: ChangeHandler, ctx: ExtensionContext): void {
  try {
    onChange(ctx);
  } catch (error) {
    console.error("[pi-todo-herdr] failed to refresh task UI:", error);
  }
}

export function registerTaskTools(pi: ExtensionAPI, state: TaskState, onChange: ChangeHandler): void {
  pi.registerTool({
    name: "add_task",
    label: "Add Tasks",
    description:
      "Atomically add one or more tasks. Use nested children for new descendants and parentId on top-level items for an existing parent.",
    promptGuidelines: [
      "Use add_task for complex multi-step work; use update_task immediately when tasks start, block, resume, or finish; skip task tracking for simple one-step work.",
    ],
    parameters: AddTaskParams,
    async execute(_toolCallId, params: AddTaskInput, _signal, _onUpdate, ctx) {
      try {
        const affected = state.add(params.tasks);
        notifyChange(onChange, ctx);
        return {
          content: [{ type: "text" as const, text: mutationText("Added", affected, state) }],
          details: details("add", state, affected),
        };
      } catch (error) {
        throwToolError("Add", error);
      }
    },
    renderCall(args: AddTaskInput, theme: Theme) {
      return renderCall("Add Tasks", args.tasks?.length, theme);
    },
    renderResult(result, { expanded }, theme) {
      return renderToolResult(result, expanded, theme);
    },
  });

  pi.registerTool({
    name: "update_task",
    label: "Update Tasks",
    description:
      "Atomically patch one or more tasks. Fields are optional except id; parentId null moves to root and position is a zero-based sibling index.",
    parameters: UpdateTaskParams,
    async execute(_toolCallId, params: UpdateTaskInput, _signal, _onUpdate, ctx) {
      try {
        const affected = state.update(params.tasks);
        notifyChange(onChange, ctx);
        return {
          content: [{ type: "text" as const, text: mutationText("Updated", affected, state) }],
          details: details("update", state, affected),
        };
      } catch (error) {
        throwToolError("Update", error);
      }
    },
    renderCall(args: UpdateTaskInput, theme: Theme) {
      return renderCall("Update Tasks", args.tasks?.length, theme);
    },
    renderResult(result, { expanded }, theme) {
      return renderToolResult(result, expanded, theme);
    },
  });

  pi.registerTool({
    name: "rm_task",
    label: "Remove Tasks",
    description: "Atomically remove multiple task IDs. Removing a parent also removes its entire descendant subtree.",
    parameters: IdsParams,
    async execute(_toolCallId, params: IdsInput, _signal, _onUpdate, ctx) {
      try {
        const affected = state.remove(params.ids);
        notifyChange(onChange, ctx);
        return {
          content: [{ type: "text" as const, text: mutationText("Removed", affected, state) }],
          details: details("remove", state, affected),
        };
      } catch (error) {
        throwToolError("Remove", error);
      }
    },
    renderCall(args: IdsInput, theme: Theme) {
      return renderCall("Remove Tasks", args.ids?.length, theme);
    },
    renderResult(result, { expanded }, theme) {
      return renderToolResult(result, expanded, theme);
    },
  });

  pi.registerTool({
    name: "list_task",
    label: "List Tasks",
    description: "List the full task tree compactly with id, name, and status.",
    parameters: ListTaskParams,
    async execute() {
      const affected = state.list();
      return {
        content: [{ type: "text" as const, text: compactTree(affected) }],
        details: details("list", state, affected),
      };
    },
    renderCall(_args, theme: Theme) {
      return renderCall("List Tasks", undefined, theme);
    },
    renderResult(result, { expanded }, theme) {
      return renderToolResult(result, expanded, theme);
    },
  });

  pi.registerTool({
    name: "get_task",
    label: "Get Tasks",
    description: "Get complete details for one or more task IDs, including parent and child IDs.",
    parameters: IdsParams,
    async execute(_toolCallId, params: IdsInput) {
      const uniqueIds = [...new Set(params.ids)];
      const affected = state.get(uniqueIds);
      const found = new Set(affected.map((task) => task.id));
      const missing = uniqueIds.filter((id) => !found.has(id));
      if (missing.length > 0) {
        throwToolError("Get", new TaskValidationError([`task ids do not exist: ${missing.join(", ")}`]));
      }
      return {
        content: [{ type: "text" as const, text: JSON.stringify(affected) }],
        details: details("get", state, affected),
      };
    },
    renderCall(args: IdsInput, theme: Theme) {
      return renderCall("Get Tasks", args.ids?.length, theme);
    },
    renderResult(result, { expanded }, theme) {
      return renderToolResult(result, expanded, theme);
    },
  });
}