import { toToolId } from "@/utils";
import { AppConfig, ToolExtension } from "@mcpx/shared-model";
import { TargetServerTool } from "@mcpx/shared-model/api";
import { create } from "zustand";
import { devtools } from "zustand/middleware";
import { useShallow } from "zustand/react/shallow";
import { SocketStore, socketStore } from "./socket";
import { createDevtoolsOptions } from "./devtools";

export interface ServerTool {
  description?: string;
  id: string;
  inputSchema?: TargetServerTool["inputSchema"];
  name: string;
  serviceName: string;
}
export interface CustomTool {
  description?: {
    action: "append" | "rewrite";
    text: string;
  };
  name: string;
  originalName?: string;
  originalTool: ServerTool;
  overrideParams: ToolExtension["overrideParams"];
}
export interface ToolsState {
  customTools: CustomTool[];
  tools: ServerTool[];
}
export interface ToolsActions {
  createCustomTool: (tool: CustomTool) => Promise<AppConfig>;
  deleteCustomTool: (tool: CustomTool) => Promise<AppConfig>;
  init: (socketStoreState: SocketStore) => void;
  setTools: (
    update:
      | ToolsState["tools"]
      | ((tools: ToolsState["tools"]) => ToolsState["tools"]),
  ) => void;
  updateCustomTool: (tool: CustomTool) => AppConfig;
}
export type ToolsStore = ToolsState & ToolsActions;

const initialState: ToolsState = { customTools: [], tools: [] };

// Helper to wait for appConfig to be available
async function waitForAppConfig(
  maxAttempts = 20,
  delay = 50,
): Promise<AppConfig | null> {
  for (let i = 0; i < maxAttempts; i++) {
    const { appConfig } = socketStore.getState();
    if (appConfig) {
      return appConfig;
    }
    // Wait using a Promise
    await new Promise((resolve) => setTimeout(resolve, delay));
  }

  // Final attempt
  const finalAppConfig = socketStore.getState().appConfig;
  if (!finalAppConfig) {
    console.error("[Tools] appConfig still not available after max attempts");
  }
  return finalAppConfig;
}

// Extended type for tool extensions with childTools
type ToolExtensionWithChildren = Partial<ToolExtension> & {
  childTools?: ToolExtension[];
};

function createCustomToolImpl(
  payload: CustomTool,
  appConfig: AppConfig,
): AppConfig {
  const newToolExtension: ToolExtension = {
    description: payload.description,
    name: payload.name,
    overrideParams: Object.fromEntries(
      Object.entries(payload.overrideParams).filter(
        ([, value]) => value !== undefined,
      ),
    ) as ToolExtension["overrideParams"],
  };

  // Build the tool extensions structure
  const toolExtensions = {
    ...(appConfig.toolExtensions?.services || {}),
  } as Record<string, Record<string, ToolExtensionWithChildren>>;

  // Ensure the service exists
  if (!toolExtensions[payload.originalTool.serviceName]) {
    toolExtensions[payload.originalTool.serviceName] = {};
  }

  // Ensure the original tool exists
  if (
    !toolExtensions[payload.originalTool.serviceName][payload.originalTool.name]
  ) {
    toolExtensions[payload.originalTool.serviceName][
      payload.originalTool.name
    ] = {
      childTools: [],
    } as ToolExtensionWithChildren;
  }

  // Add the new custom tool to the child tools
  const toolExt =
    toolExtensions[payload.originalTool.serviceName][payload.originalTool.name];
  if (!toolExt.childTools) {
    toolExt.childTools = [];
  }
  toolExt.childTools.push(newToolExtension);

  const updates: AppConfig = {
    ...appConfig,
    toolExtensions: {
      services: toolExtensions as AppConfig["toolExtensions"]["services"],
    },
  };

  return updates;
}

function deleteCustomToolImpl(
  tool: CustomTool,
  appConfig: AppConfig,
  customTools: CustomTool[],
): AppConfig {
  const newCustomTools = customTools.filter(
    (t) =>
      !(t.originalTool.id === tool.originalTool.id && t.name === tool.name),
  );

  type ServiceToolsMap = Record<
    string,
    Record<string, ToolExtensionWithChildren>
  >;

  const services: ServiceToolsMap = Object.fromEntries(
    Object.entries(
      (appConfig.toolExtensions?.services || {}) as ServiceToolsMap,
    )
      .filter(([serviceName]) =>
        newCustomTools.some((t) => t.originalTool.serviceName === serviceName),
      )
      .map(([serviceName, serviceTools]) => [
        serviceName,
        Object.fromEntries(
          Object.entries(
            serviceTools as Record<string, ToolExtensionWithChildren>,
          ).filter(([toolName]) =>
            newCustomTools.some(
              (t) =>
                t.originalTool.name === toolName &&
                t.originalTool.serviceName === serviceName,
            ),
          ),
        ),
      ]),
  );

  const updates: AppConfig = {
    ...appConfig,
    toolExtensions: {
      services: newCustomTools.reduce((acc: ServiceToolsMap, t) => {
        const serviceTools = acc[t.originalTool.serviceName] || {};
        const existingTool: ToolExtensionWithChildren = serviceTools[
          t.originalTool.name
        ] || {
          childTools: [],
        };
        existingTool.childTools = (existingTool.childTools || []).filter(
          (ct: ToolExtension) =>
            !(
              t.originalTool.id === tool.originalTool.id &&
              ct.name === tool.name
            ),
        );

        if (existingTool.childTools.length > 0) {
          serviceTools[t.originalTool.name] = existingTool;
        } else {
          delete serviceTools[t.originalTool.name];
        }

        acc[t.originalTool.serviceName] = serviceTools;

        return acc;
      }, services) as AppConfig["toolExtensions"]["services"],
    },
  };

  return updates;
}

const toolsStore = create<ToolsStore>()(
  devtools(
    (set, get) => ({
      ...initialState,
      createCustomTool: async (payload) => {
        const appConfig = await waitForAppConfig();

        if (!appConfig) {
          throw new Error("App config is not available.");
        }

        return createCustomToolImpl(payload, appConfig);
      },

      deleteCustomTool: async (tool) => {
        const appConfig = await waitForAppConfig();

        if (!appConfig) {
          throw new Error("App config is not available.");
        }

        const { customTools } = get();
        return deleteCustomToolImpl(tool, appConfig, customTools);
      },
      init: (socketStoreState: SocketStore) => {
        if (!socketStoreState.systemState || !socketStoreState.appConfig) {
          return;
        }

        const { systemState, appConfig } = socketStoreState;

        const customTools: CustomTool[] = [];

        for (const [serviceName, serviceTools] of Object.entries(
          appConfig.toolExtensions?.services || {},
        )) {
          // TODO: Maybe populate custom tools regardless of the service tools?
          for (const [originalToolName, { childTools }] of Object.entries(
            serviceTools,
          )) {
            const originalTool = systemState.targetServers
              .find((server) => server.name === serviceName)
              ?.originalTools.find((tool) => tool.name === originalToolName);

            if (!originalTool) {
              continue;
            }

            const toolsForService = childTools.map(
              (toolExtension: ToolExtension) => {
                const parameterDescriptions = Object.fromEntries(
                  Object.entries(toolExtension.overrideParams || {})
                    .map(([name, param]) => {
                      const descObject = param?.description;
                      if (!descObject || typeof descObject.text !== "string") {
                        return [name, undefined];
                      }

                      const text = descObject.text.trim();
                      if (!text) {
                        return [name, undefined];
                      }

                      return [name, text];
                    })
                    .filter(([, text]) => text !== undefined),
                );

                return {
                  description: toolExtension.description,
                  id: toToolId(serviceName, toolExtension.name),
                  name: toolExtension.name,
                  originalTool: {
                    description: originalTool.description || "",
                    id: toToolId(serviceName, originalToolName),
                    inputSchema: originalTool.inputSchema,
                    name: originalToolName,
                    serviceName,
                  },
                  overrideParams: toolExtension.overrideParams,
                  parameterDescriptions,
                };
              },
            );

            customTools.push(...toolsForService);
          }
        }

        const tools: ToolsState["tools"] = [];

        for (const targetServer of systemState.targetServers) {
          const serviceName = targetServer.name;
          for (const {
            description,
            name,
            inputSchema,
          } of targetServer.originalTools) {
            const id = toToolId(serviceName, name);
            tools.push({
              description,
              id,
              inputSchema,
              name,
              serviceName,
            });
          }
        }

        set({ customTools, tools });
      },
      setTools: (update) => {
        set((state) => ({
          tools: typeof update === "function" ? update(state.tools) : update,
        }));
      },
      updateCustomTool: (tool) => {
        const { appConfig } = socketStore.getState();

        if (!appConfig) {
          throw new Error("App config is not available.");
        }

        const toolExtensions = {
          ...(appConfig.toolExtensions?.services || {}),
        };

        // Ensure the service exists
        if (!toolExtensions[tool.originalTool.serviceName]) {
          toolExtensions[tool.originalTool.serviceName] = {};
        }

        // Ensure the original tool exists
        if (
          !toolExtensions[tool.originalTool.serviceName][tool.originalTool.name]
        ) {
          toolExtensions[tool.originalTool.serviceName][
            tool.originalTool.name
          ] = {
            childTools: [],
          };
        }

        // Find and update the existing custom tool
        const childTools =
          toolExtensions[tool.originalTool.serviceName][tool.originalTool.name]
            .childTools;

        // For edit mode, we need to find the tool by the original name, not the new name
        // because the user might be changing the name
        const lookupName = tool.originalName || tool.name;
        const toolIndex = childTools.findIndex((ct: ToolExtension) => {
          // Use originalName if available (for edit mode), otherwise use current name
          return ct.name === lookupName;
        });

        if (toolIndex >= 0) {
          const updatedTool: ToolExtension = {
            description: tool.description,
            name: tool.name,
            overrideParams: Object.fromEntries(
              Object.entries(tool.overrideParams).filter(
                ([, value]) => value !== undefined,
              ),
            ),
          };

          childTools[toolIndex] = updatedTool;
        } else {
          // Tool not found, this shouldn't happen in edit mode
        }

        const updates: AppConfig = {
          ...appConfig,
          toolExtensions: {
            services: toolExtensions,
          },
        };

        return updates;
      },
    }),
    createDevtoolsOptions("tools"),
  ),
);

// Subscribe to socket store updates to keep tools state
// in sync with the app configuration and system state.
socketStore.subscribe((state) => {
  // Only update if there are no pending operations (to prevent overwriting optimistic updates)
  toolsStore.getState().init(state);
});

export const useToolsStore = <T>(selector: (state: ToolsStore) => T) =>
  toolsStore(useShallow(selector));

export { toolsStore };

export const initToolsStore = () => {
  toolsStore.getState().init(socketStore.getState());
};
