import {
  TOOL_SWAP_AGENT_SHORT_NAME,
  TOOL_SWAP_TO_DEFAULT_AGENT_SHORT_NAME,
  TOOL_TODO_WRITE_SHORT_NAME,
  type ToolStateMcpToolError,
  ToolStateMcpToolErrorSchema,
} from "@archestra/shared";
import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
import { z } from "zod";
import { isAgentTypeAdmin } from "@/auth/agent-type-permissions";
import logger from "@/logging";
import {
  AgentModel,
  ChatOpsChannelBindingModel,
  ChatOpsThreadAgentOverrideModel,
  ConversationModel,
  OrganizationModel,
} from "@/models";
import { resolveConversationLlmSelectionForAgent } from "@/utils/llm-resolution";
import {
  catchError,
  defineArchestraTool,
  defineArchestraTools,
  EmptyToolArgsSchema,
  errorResult,
  structuredSuccessResult,
  structuredToolErrorResult,
} from "./helpers";
import type { ArchestraContext } from "./types";

// === Constants ===

const TodoItemSchema = z
  .object({
    id: z.number().int().describe("Unique identifier for the todo item."),
    content: z
      .string()
      .describe("The content or description of the todo item."),
    status: z
      .enum(["pending", "in_progress", "completed"])
      .describe("The current status of the todo item."),
  })
  .strict();

const TodoWriteOutputSchema = z.object({
  success: z.literal(true).describe("Whether the write succeeded."),
  todoCount: z
    .number()
    .int()
    .nonnegative()
    .describe("How many todo items were written."),
});

const SwapAgentStateCodeSchema = z.enum([
  "no_agent_found",
  "already_using_agent",
  "no_default_agent",
  "default_agent_not_found",
  "already_using_default_agent",
]);

const SwapAgentOutputSchema = z.discriminatedUnion("success", [
  z.object({
    success: z.literal(true).describe("Whether the swap succeeded."),
    agent_id: z.string().describe("The agent ID the conversation now uses."),
    agent_name: z
      .string()
      .describe("The agent name the conversation now uses."),
  }),
  z.object({
    success: z.literal(false).describe("Whether the swap succeeded."),
    code: SwapAgentStateCodeSchema.describe("Why the swap was not applied."),
    message: z.string().describe("Human-readable explanation."),
    archestraError: ToolStateMcpToolErrorSchema,
  }),
]);

type SwapAgentStateCode = z.infer<typeof SwapAgentStateCodeSchema>;

const registry = defineArchestraTools([
  defineArchestraTool({
    shortName: TOOL_TODO_WRITE_SHORT_NAME,
    title: "Write Todos",
    description:
      "Write todos to the current conversation. You have access to this tool to help you manage and plan tasks. Use it VERY frequently to ensure that you are tracking your tasks and giving the user visibility into your progress. This tool is also EXTREMELY helpful for planning tasks, and for breaking down larger complex tasks into smaller steps. If you do not use this tool when planning, you may forget to do important tasks - and that is unacceptable. It is critical that you mark todos as completed as soon as you are done with a task. Do not batch up multiple tasks before marking them as completed.",
    schema: z
      .object({
        todos: z
          .array(TodoItemSchema)
          .describe("Array of todo items to write to the conversation."),
      })
      .strict(),
    outputSchema: TodoWriteOutputSchema,
    async handler({ args, context }) {
      const { agent: contextAgent } = context;

      logger.info(
        { agentId: contextAgent.id, todoArgs: args },
        "todo_write tool called",
      );

      try {
        return structuredSuccessResult(
          { success: true, todoCount: args.todos.length },
          `Successfully wrote ${args.todos.length} todo item(s) to the conversation`,
        );
      } catch (error) {
        return catchError(error, "writing todos");
      }
    },
  }),
  defineArchestraTool({
    shortName: TOOL_SWAP_AGENT_SHORT_NAME,
    title: "Swap Agent",
    description:
      "Switch the current conversation to a different agent. The new agent will automatically continue the conversation. Use this when the user asks to switch to or talk to a different agent.",
    schema: z
      .object({
        agent_name: z
          .string()
          .trim()
          .min(1)
          .describe("The name of the agent to switch to."),
      })
      .strict(),
    outputSchema: SwapAgentOutputSchema,
    async handler({ args, context }) {
      return handleSwapAgent({
        agentName: args.agent_name,
        context,
      });
    },
  }),
  defineArchestraTool({
    shortName: TOOL_SWAP_TO_DEFAULT_AGENT_SHORT_NAME,
    title: "Swap to Default Agent",
    description:
      "Return to the default agent. You MUST call this — without asking the user — when you don't have the right tools to fulfill a request, when you are stuck and cannot help further, when you are done with your task, or when the user wants to go back. Always write a brief message before calling this tool summarizing why you are switching back (e.g. what you accomplished, what tool is missing, or why you cannot continue).",
    schema: EmptyToolArgsSchema,
    outputSchema: SwapAgentOutputSchema,
    async handler({ context }) {
      return handleSwapToDefaultAgent({ context });
    },
  }),
] as const);

export const toolEntries = registry.toolEntries;

// === Exports ===

export const tools = registry.tools;

function swapAgentStateResult(params: {
  code: SwapAgentStateCode;
  message: string;
  toolName: string;
}): CallToolResult {
  const archestraError: ToolStateMcpToolError = {
    type: "tool_state",
    code: params.code,
    message: params.message,
    toolName: params.toolName,
  };

  return structuredToolErrorResult({
    error: archestraError,
    text: JSON.stringify({
      success: false,
      code: params.code,
      message: params.message,
      archestraError,
    }),
    structuredContent: {
      success: false,
      code: params.code,
      message: params.message,
    },
    // These are expected chat-routing states, not MCP transport failures. Keep
    // isError false so the agent can respond normally instead of surfacing a
    // global chat error.
    isError: false,
  });
}

async function handleSwapAgent(params: {
  agentName: string;
  context: ArchestraContext;
}): Promise<CallToolResult> {
  const { agentName, context } = params;
  const { agent: contextAgent } = context;
  logger.info(
    {
      agentId: contextAgent.id,
      agentName,
      chatOpsBindingId: context.chatOpsBindingId ?? null,
      conversationId: context.conversationId ?? null,
    },
    "swap_agent tool called",
  );

  try {
    if (!context.userId || !context.organizationId) {
      return errorResult(
        "This tool requires user and organization context. It can only be used within an authenticated chat session.",
      );
    }

    const hasConversationContext = Boolean(context.conversationId);
    const hasChatOpsContext = Boolean(context.chatOpsBindingId);
    if (!hasConversationContext && !hasChatOpsContext) {
      return errorResult(
        "This tool requires conversation context. It can only be used within an active chat conversation or chatops channel.",
      );
    }

    // Look up agent by name
    const isAdmin =
      context.userId && context.organizationId
        ? await isAgentTypeAdmin({
            userId: context.userId,
            organizationId: context.organizationId,
            agentType: "agent",
          })
        : false;

    const results = await AgentModel.findAllPaginated(
      { limit: 5, offset: 0 },
      undefined,
      {
        name: agentName,
        agentType: "agent",
        // Hide other users' personal agents. swap_agent is the primary
        // Archestra MCP use-case and requires only the caller's own personal
        // agents to be visible, even though admins can see all personal
        // agents in the UI.
        excludeOtherPersonalAgents: true,
      },
      context.userId,
      isAdmin,
    );

    if (results.data.length === 0) {
      return swapAgentStateResult({
        code: "no_agent_found",
        message: `No agent found matching "${agentName}".`,
        toolName: TOOL_SWAP_AGENT_SHORT_NAME,
      });
    }

    // Pick exact name match if available, otherwise first result
    const targetAgent =
      results.data.find(
        (a) => a.name.toLowerCase() === agentName.toLowerCase(),
      ) ?? results.data[0];

    // Prevent swapping to the same agent
    if (targetAgent.id === contextAgent.id) {
      return swapAgentStateResult({
        code: "already_using_agent",
        message: `Already using agent "${targetAgent.name}". Choose a different agent.`,
        toolName: TOOL_SWAP_AGENT_SHORT_NAME,
      });
    }

    // In chatops-triggered A2A runs we can have both:
    // - chatOpsBindingId: real channel binding context
    // - conversationId: synthetic isolation key for tool/session caching
    // Prefer the chatops binding whenever available.
    if (context.chatOpsBindingId) {
      if (!context.chatOpsThreadId) {
        return errorResult(
          "This tool requires thread context in chatops. Cannot determine which thread to swap.",
        );
      }

      // Validate binding exists and user has permission
      const binding = await ChatOpsChannelBindingModel.findById(
        context.chatOpsBindingId,
      );
      if (!binding || binding.organizationId !== context.organizationId) {
        return errorResult("Failed to update chatops channel agent.");
      }

      // Personal agent scope check
      if (targetAgent.scope === "personal") {
        if (!binding.isDm) {
          return errorResult(
            "Personal agents cannot be assigned to channels. Use an org-scoped or team-scoped agent instead.",
          );
        }
        if (targetAgent.authorId !== context.userId) {
          return errorResult(
            "You can only assign your own personal agents to your DM.",
          );
        }
      }

      // Write thread-scoped override instead of mutating channel binding
      const override = await ChatOpsThreadAgentOverrideModel.upsert(
        context.chatOpsBindingId,
        context.chatOpsThreadId,
        targetAgent.id,
      );
      if (!override) {
        return errorResult("Failed to update chatops thread agent.");
      }
    } else if (context.conversationId) {
      const llmSelection = await resolveConversationLlmSelectionForAgent({
        agent: {
          llmApiKeyId: targetAgent.llmApiKeyId ?? null,
          modelId: targetAgent.modelId ?? null,
        },
        organizationId: context.organizationId,
        userId: context.userId,
      });

      // Update the conversation's agent and LLM selection together so the
      // follow-up response uses the new agent's model/key immediately.
      const updated = await ConversationModel.update(
        context.conversationId,
        context.userId,
        context.organizationId,
        {
          agentId: targetAgent.id,
          chatApiKeyId: llmSelection.chatApiKeyId,
          modelId: llmSelection.modelId,
        },
      );
      if (!updated) {
        logger.warn(
          {
            conversationId: context.conversationId,
            userId: context.userId,
            organizationId: context.organizationId,
            chatOpsBindingId: context.chatOpsBindingId ?? null,
          },
          "swap_agent: conversation update failed, possible missing chatOpsBindingId",
        );
        return errorResult("Failed to update conversation agent.");
      }
    }

    return structuredSuccessResult(
      {
        success: true,
        agent_id: targetAgent.id,
        agent_name: targetAgent.name,
      },
      `Successfully swapped to agent "${targetAgent.name}" (ID: ${targetAgent.id}).`,
    );
  } catch (error) {
    return catchError(error, "swapping agent");
  }
}

async function handleSwapToDefaultAgent(params: {
  context: ArchestraContext;
}): Promise<CallToolResult> {
  const { context } = params;
  const { agent: contextAgent } = context;

  logger.info(
    {
      agentId: contextAgent.id,
      chatOpsBindingId: context.chatOpsBindingId ?? null,
      conversationId: context.conversationId ?? null,
    },
    "swap_to_default_agent tool called",
  );

  try {
    if (!context.userId || !context.organizationId) {
      return errorResult(
        "This tool requires user and organization context. It can only be used within an authenticated chat session.",
      );
    }

    const hasConversationContext = Boolean(context.conversationId);
    const hasChatOpsContext = Boolean(context.chatOpsBindingId);
    if (!hasConversationContext && !hasChatOpsContext) {
      return errorResult(
        "This tool requires conversation context. It can only be used within an active chat conversation or chatops channel.",
      );
    }

    const org = await OrganizationModel.getById(context.organizationId);
    const defaultAgentId = org?.defaultAgentId ?? null;

    if (!defaultAgentId) {
      return swapAgentStateResult({
        code: "no_default_agent",
        message: "No default agent is configured for this organization.",
        toolName: TOOL_SWAP_TO_DEFAULT_AGENT_SHORT_NAME,
      });
    }

    const targetAgent = await AgentModel.findById(defaultAgentId);
    if (!targetAgent) {
      return swapAgentStateResult({
        code: "default_agent_not_found",
        message: "Default agent not found.",
        toolName: TOOL_SWAP_TO_DEFAULT_AGENT_SHORT_NAME,
      });
    }

    if (targetAgent.id === contextAgent.id) {
      return swapAgentStateResult({
        code: "already_using_default_agent",
        message: `Already using the default agent "${targetAgent.name}".`,
        toolName: TOOL_SWAP_TO_DEFAULT_AGENT_SHORT_NAME,
      });
    }

    // In chatops-triggered A2A runs we can have both:
    // - chatOpsBindingId: real channel binding context
    // - conversationId: synthetic isolation key for tool/session caching
    // Prefer the chatops binding whenever available.
    if (context.chatOpsBindingId) {
      if (!context.chatOpsThreadId) {
        return errorResult(
          "This tool requires thread context in chatops. Cannot determine which thread to swap.",
        );
      }

      // Validate binding exists and user has permission
      const binding = await ChatOpsChannelBindingModel.findById(
        context.chatOpsBindingId,
      );
      if (!binding || binding.organizationId !== context.organizationId) {
        return errorResult("Failed to update chatops channel agent.");
      }

      // Personal agent scope check
      if (targetAgent.scope === "personal") {
        if (!binding.isDm) {
          return errorResult(
            "Personal agents cannot be assigned to channels. Use an org-scoped or team-scoped agent instead.",
          );
        }
        if (targetAgent.authorId !== context.userId) {
          return errorResult(
            "You can only assign your own personal agents to your DM.",
          );
        }
      }

      // Write thread-scoped override instead of mutating channel binding
      const override = await ChatOpsThreadAgentOverrideModel.upsert(
        context.chatOpsBindingId,
        context.chatOpsThreadId,
        targetAgent.id,
      );
      if (!override) {
        return errorResult("Failed to update chatops thread agent.");
      }
    } else if (context.conversationId) {
      const llmSelection = await resolveConversationLlmSelectionForAgent({
        agent: {
          llmApiKeyId: targetAgent.llmApiKeyId ?? null,
          modelId: targetAgent.modelId ?? null,
        },
        organizationId: context.organizationId,
        userId: context.userId,
      });

      const updated = await ConversationModel.update(
        context.conversationId,
        context.userId,
        context.organizationId,
        {
          agentId: defaultAgentId,
          chatApiKeyId: llmSelection.chatApiKeyId,
          modelId: llmSelection.modelId,
        },
      );
      if (!updated) {
        logger.warn(
          {
            conversationId: context.conversationId,
            userId: context.userId,
            organizationId: context.organizationId,
            chatOpsBindingId: context.chatOpsBindingId ?? null,
          },
          "swap_to_default_agent: conversation update failed, possible missing chatOpsBindingId",
        );
        return errorResult("Failed to update conversation agent.");
      }
    }

    return structuredSuccessResult(
      {
        success: true,
        agent_id: targetAgent.id,
        agent_name: targetAgent.name,
      },
      `Successfully swapped to default agent "${targetAgent.name}" (ID: ${targetAgent.id}).`,
    );
  } catch (error) {
    return catchError(error, "swapping to default agent");
  }
}
