import { get } from "@noya-app/observable";
import { Static, Type } from "@sinclair/typebox";
import { createOrCastValue } from "./checkType";
import { diff } from "./diff";
import type { NoyaManager } from "./NoyaManager";
import { OperationManager } from "./stateManager";

const pathSchema = Type.Array(Type.Union([Type.String(), Type.Number()]), {
  description:
    "An exact path into the JSON document, expressed as string and number segments.",
});

const documentOperationSchema = Type.Union([
  Type.Object({
    op: Type.Literal("set"),
    path: pathSchema,
    value: Type.Any(),
  }),
  Type.Object({
    op: Type.Literal("add"),
    path: pathSchema,
    value: Type.Any(),
  }),
  Type.Object({
    op: Type.Literal("remove"),
    path: pathSchema,
  }),
  Type.Object({
    op: Type.Literal("move"),
    from: pathSchema,
    path: pathSchema,
  }),
]);

export const readDocumentParametersSchema = Type.Object({
  path: Type.Optional(pathSchema),
});

export const applyDocumentOperationsParametersSchema = Type.Object({
  operations: Type.Array(documentOperationSchema, { minItems: 1 }),
  expectedHash: Type.Optional(
    Type.String({
      description:
        "The document hash returned by an earlier tool result. The edit is rejected if the document changed.",
    })
  ),
  summary: Type.String({
    description: "A concise description of the requested document changes.",
  }),
});

export const runDocumentScriptParametersSchema = Type.Object({
  sourceCode: Type.String({
    description:
      'JavaScript defining a named function "updateState(state)" which mutates state in place.',
  }),
  expectedHash: Type.Optional(
    Type.String({
      description:
        "The document hash returned by an earlier tool result. The edit is rejected if the document changed.",
    })
  ),
  summary: Type.String({
    description: "A concise description of the requested document changes.",
  }),
});

type ReadDocumentParameters = Static<typeof readDocumentParametersSchema>;
type ApplyDocumentOperationsParameters = Static<
  typeof applyDocumentOperationsParametersSchema
>;
type RunDocumentScriptParameters = Static<
  typeof runDocumentScriptParametersSchema
>;

const MAX_SCRIPT_LENGTH = 20_000;
const MAX_SERIALIZED_STATE_LENGTH = 2_000_000;

const stringifyToolResult = (value: unknown) =>
  JSON.stringify(value, (_key, item) =>
    typeof item === "bigint" ? item.toString() : item
  );

function parseParameters<T>(
  schema: Parameters<typeof createOrCastValue>[0]["schema"],
  value: unknown
): T {
  return createOrCastValue({
    schema,
    value,
    defs: {},
  }) as T;
}

function staleDocumentResult(expectedHash: string, actualHash: string) {
  return stringifyToolResult({
    ok: false,
    error: "The document changed since it was read. Read it again and retry.",
    expectedHash,
    actualHash,
  });
}

export function registerDocumentAITools(
  noyaManager: NoyaManager<any, any, any, any, any>
) {
  const { aiManager, multiplayerStateManager } = noyaManager;
  const previousSystemMessage = aiManager.systemMessage$.get();
  const schema = multiplayerStateManager.schema;
  const systemMessage = [
    "You are editing the user's current Noya document.",
    "Inspect current values with readDocument before changing them.",
    "Prefer applyDocumentOperations for precise changes. Use runDocumentScript only when path operations would be impractical.",
    "Never invent paths or document structure. Use the exact JSON path segments returned by readDocument.",
    schema &&
      `The document conforms to this schema:\n\`\`\`json\n${JSON.stringify(schema, null, 2)}\n\`\`\``,
  ]
    .filter(Boolean)
    .join("\n\n");

  aiManager.systemMessage$.set(systemMessage);

  const unregisterTools = [
    aiManager.registerTool({
      functionName: "readDocument",
      description:
        "Read the current JSON document or an exact value within it. Omit path to read the whole document.",
      parameters: readDocumentParametersSchema,
      onCall: async (parameters) => {
        try {
          const parsed = parseParameters<ReadDocumentParameters>(
            readDocumentParametersSchema,
            parameters
          );
          const state = multiplayerStateManager.getOptimisticState();
          const value =
            parsed.path === undefined ? state : get(state, parsed.path);

          return stringifyToolResult({
            ok: true,
            path: parsed.path ?? [],
            value,
            hash: multiplayerStateManager.getOptimisticHash(),
          });
        } catch (error) {
          return stringifyToolResult({
            ok: false,
            error: error instanceof Error ? error.message : String(error),
          });
        }
      },
    }),
    aiManager.registerTool({
      functionName: "applyDocumentOperations",
      description:
        "Atomically set, add, remove, or move values at exact paths in the current JSON document.",
      parameters: applyDocumentOperationsParametersSchema,
      onCall: async (parameters) => {
        try {
          const parsed = parseParameters<ApplyDocumentOperationsParameters>(
            applyDocumentOperationsParametersSchema,
            parameters
          );
          const actualHash = multiplayerStateManager.getOptimisticHash();

          if (parsed.expectedHash && parsed.expectedHash !== actualHash) {
            return staleDocumentResult(parsed.expectedHash, actualHash);
          }

          const operations = new OperationManager<any>(
            multiplayerStateManager.getOptimisticState()
          );

          for (const operation of parsed.operations) {
            switch (operation.op) {
              case "set":
                operations.set(operation.path, operation.value);
                break;
              case "add":
                operations.add(operation.path, operation.value as never);
                break;
              case "remove":
                operations.remove(operation.path);
                break;
              case "move":
                operations.move(operation.from, operation.path);
                break;
            }
          }

          multiplayerStateManager.applyPatch(
            { name: parsed.summary },
            {
              patches: operations.patches,
              inversePatches: operations.inversePatches,
            }
          );

          return stringifyToolResult({
            ok: true,
            summary: parsed.summary,
            operationCount: parsed.operations.length,
            hash: multiplayerStateManager.getOptimisticHash(),
          });
        } catch (error) {
          return stringifyToolResult({
            ok: false,
            error: error instanceof Error ? error.message : String(error),
          });
        }
      },
    }),
    aiManager.registerTool({
      functionName: "runDocumentScript",
      description:
        'Run isolated JavaScript that defines updateState(state) and mutates the current JSON document. Use "crypto.randomUUID()" when new IDs are required.',
      parameters: runDocumentScriptParametersSchema,
      onCall: async (parameters) => {
        try {
          const parsed = parseParameters<RunDocumentScriptParameters>(
            runDocumentScriptParametersSchema,
            parameters
          );
          const actualHash = multiplayerStateManager.getOptimisticHash();

          if (parsed.expectedHash && parsed.expectedHash !== actualHash) {
            return staleDocumentResult(parsed.expectedHash, actualHash);
          }
          if (parsed.sourceCode.length > MAX_SCRIPT_LENGTH) {
            throw new Error("The document script is too large.");
          }

          const state = multiplayerStateManager.getOptimisticState();
          const serializedState = JSON.stringify(state);
          if (serializedState.length > MAX_SERIALIZED_STATE_LENGTH) {
            throw new Error("The document is too large for script execution.");
          }

          const code = `const __state = ${serializedState};
${parsed.sourceCode}
if (typeof updateState !== "function") {
  throw new Error("The script must define updateState(state)");
}
updateState(__state);
__state;`;
          const nextState = await noyaManager.evalManager.evaluate(code);
          const [patches, inversePatches] = diff(state, nextState);

          multiplayerStateManager.applyPatch(
            { name: parsed.summary },
            { patches, inversePatches }
          );

          return stringifyToolResult({
            ok: true,
            summary: parsed.summary,
            patchCount: patches.length,
            hash: multiplayerStateManager.getOptimisticHash(),
          });
        } catch (error) {
          return stringifyToolResult({
            ok: false,
            error: error instanceof Error ? error.message : String(error),
          });
        }
      },
    }),
  ];

  return () => {
    unregisterTools.forEach((unregister) => unregister());
    if (aiManager.systemMessage$.get() === systemMessage) {
      aiManager.systemMessage$.set(previousSystemMessage);
    }
  };
}