/**
 * @license
 * Copyright 2025 Google LLC
 * SPDX-License-Identifier: Apache-2.0
 */

import {
  Config,
  ToolCallRequestInfo,
  ToolCallResponseInfo,
  ToolConfirmationOutcome,
  Tool,
  ToolCallConfirmationDetails,
  ToolResult,
} from '@gemini-code/core';
import { Part, PartUnion, PartListUnion } from '@google/genai';
import { useCallback, useEffect, useState } from 'react';
import {
  HistoryItemToolGroup,
  IndividualToolCallDisplay,
  ToolCallStatus,
  HistoryItemWithoutId,
} from '../types.js';

type ValidatingToolCall = {
  status: 'validating';
  request: ToolCallRequestInfo;
  tool: Tool;
};

type ScheduledToolCall = {
  status: 'scheduled';
  request: ToolCallRequestInfo;
  tool: Tool;
};

type ErroredToolCall = {
  status: 'error';
  request: ToolCallRequestInfo;
  response: ToolCallResponseInfo;
};

type SuccessfulToolCall = {
  status: 'success';
  request: ToolCallRequestInfo;
  tool: Tool;
  response: ToolCallResponseInfo;
};

export type ExecutingToolCall = {
  status: 'executing';
  request: ToolCallRequestInfo;
  tool: Tool;
  liveOutput?: string;
};

type CancelledToolCall = {
  status: 'cancelled';
  request: ToolCallRequestInfo;
  response: ToolCallResponseInfo;
  tool: Tool;
};

type WaitingToolCall = {
  status: 'awaiting_approval';
  request: ToolCallRequestInfo;
  tool: Tool;
  confirmationDetails: ToolCallConfirmationDetails;
};

export type Status = ToolCall['status'];

export type ToolCall =
  | ValidatingToolCall
  | ScheduledToolCall
  | ErroredToolCall
  | SuccessfulToolCall
  | ExecutingToolCall
  | CancelledToolCall
  | WaitingToolCall;

export type ScheduleFn = (
  request: ToolCallRequestInfo | ToolCallRequestInfo[],
) => void;
export type CancelFn = () => void;
export type CompletedToolCall =
  | SuccessfulToolCall
  | CancelledToolCall
  | ErroredToolCall;

/**
 * Formats a PartListUnion response from a tool into JSON suitable for a Gemini
 * FunctionResponse and additional Parts to include after that response.
 *
 * This is required because FunctionReponse appears to only support JSON
 * and not arbitrary parts. Including parts like inlineData or fileData
 * directly in a FunctionResponse confuses the model resulting in a failure
 * to interpret the multimodal content and context window exceeded errors.
 */

export function formatLlmContentForFunctionResponse(
  llmContent: PartListUnion,
): {
  functionResponseJson: Record<string, string>;
  additionalParts: PartUnion[];
} {
  const additionalParts: PartUnion[] = [];
  let functionResponseJson: Record<string, string>;

  if (Array.isArray(llmContent) && llmContent.length === 1) {
    // Ensure that length 1 arrays are treated as a single Part.
    llmContent = llmContent[0];
  }

  if (typeof llmContent === 'string') {
    functionResponseJson = { output: llmContent };
  } else if (Array.isArray(llmContent)) {
    functionResponseJson = { status: 'Tool execution succeeded.' };
    additionalParts.push(...llmContent);
  } else {
    if (
      llmContent.inlineData !== undefined ||
      llmContent.fileData !== undefined
    ) {
      // For Parts like inlineData or fileData, use the returnDisplay as the textual output for the functionResponse.
      // The actual Part will be added to additionalParts.
      functionResponseJson = {
        status: `Binary content of type ${llmContent.inlineData?.mimeType || llmContent.fileData?.mimeType || 'unknown'} was processed.`,
      };
      additionalParts.push(llmContent);
    } else if (llmContent.text !== undefined) {
      functionResponseJson = { output: llmContent.text };
    } else {
      functionResponseJson = { status: 'Tool execution succeeded.' };
      additionalParts.push(llmContent);
    }
  }

  return {
    functionResponseJson,
    additionalParts,
  };
}

export function useToolScheduler(
  onComplete: (tools: CompletedToolCall[]) => void,
  config: Config,
  setPendingHistoryItem: React.Dispatch<
    React.SetStateAction<HistoryItemWithoutId | null>
  >,
): [ToolCall[], ScheduleFn, CancelFn] {
  const [toolRegistry] = useState(() => config.getToolRegistry());
  const [toolCalls, setToolCalls] = useState<ToolCall[]>([]);
  const [abortController, setAbortController] = useState<AbortController>(
    () => new AbortController(),
  );

  const isRunning = toolCalls.some(
    (t) => t.status === 'executing' || t.status === 'awaiting_approval',
  );
  // Note: request array[] typically signal pending tool calls
  const schedule = useCallback(
    async (request: ToolCallRequestInfo | ToolCallRequestInfo[]) => {
      if (isRunning) {
        throw new Error(
          'Cannot schedule tool calls while other tool calls are running',
        );
      }
      const requestsToProcess = Array.isArray(request) ? request : [request];

      // Step 1: Create initial calls with 'validating' status (or 'error' if tool not found)
      // and add them to the state immediately to make the UI busy.
      const initialNewCalls: ToolCall[] = requestsToProcess.map(
        (r): ToolCall => {
          const tool = toolRegistry.getTool(r.name);
          if (!tool) {
            return {
              status: 'error',
              request: r,
              response: toolErrorResponse(
                r,
                new Error(`tool ${r.name} does not exist`),
              ),
            };
          }
          // Set to 'validating' immediately. This will make streamingState 'Responding'.
          return { status: 'validating', request: r, tool };
        },
      );
      setToolCalls((prevCalls) => prevCalls.concat(initialNewCalls));

      // Step 2: Asynchronously check for confirmation and update status for each new call.
      initialNewCalls.forEach(async (initialCall) => {
        // If the call was already marked as an error (tool not found), skip further processing.
        if (initialCall.status !== 'validating') return;

        const { request: r, tool } = initialCall;
        try {
          const userApproval = await tool.shouldConfirmExecute(
            r.args,
            abortController.signal,
          );
          if (userApproval) {
            // Confirmation is needed. Update status to 'awaiting_approval'.
            setToolCalls(
              setStatus(r.callId, 'awaiting_approval', {
                ...userApproval,
                onConfirm: async (outcome) => {
                  // This onConfirm is triggered by user interaction later.
                  await userApproval.onConfirm(outcome);
                  setToolCalls(
                    outcome === ToolConfirmationOutcome.Cancel
                      ? setStatus(
                          r.callId,
                          'cancelled',
                          'User did not allow tool call',
                        )
                      : // If confirmed, it goes to 'scheduled' to be picked up by the execution effect.
                        setStatus(r.callId, 'scheduled'),
                  );
                },
              }),
            );
          } else {
            // No confirmation needed, move to 'scheduled' for execution.
            setToolCalls(setStatus(r.callId, 'scheduled'));
          }
        } catch (e) {
          // Handle errors from tool.shouldConfirmExecute() itself.
          setToolCalls(
            setStatus(
              r.callId,
              'error',
              toolErrorResponse(
                r,
                e instanceof Error ? e : new Error(String(e)),
              ),
            ),
          );
        }
      });
    },
    [isRunning, setToolCalls, toolRegistry, abortController.signal],
  );

  const cancel = useCallback(
    (reason: string = 'unspecified') => {
      abortController.abort();
      setAbortController(new AbortController());
      setToolCalls((tc) =>
        tc.map((c) =>
          c.status !== 'error' && c.status !== 'executing'
            ? {
                ...c,
                status: 'cancelled',
                response: {
                  callId: c.request.callId,
                  responseParts: {
                    functionResponse: {
                      id: c.request.callId,
                      name: c.request.name,
                      response: {
                        error: `[Operation Cancelled] Reason: ${reason}`,
                      },
                    },
                  },
                  resultDisplay: undefined,
                  error: undefined,
                },
              }
            : c,
        ),
      );
    },
    [abortController],
  );

  useEffect(() => {
    // effect for executing scheduled tool calls
    const allToolsConfirmed = toolCalls.every(
      (t) => t.status === 'scheduled' || t.status === 'cancelled',
    );
    if (allToolsConfirmed) {
      const signal = abortController.signal;
      toolCalls
        .filter((t) => t.status === 'scheduled')
        .forEach((t) => {
          const callId = t.request.callId;
          setToolCalls(setStatus(t.request.callId, 'executing'));

          const updateOutput = t.tool.canUpdateOutput
            ? (output: string) => {
                setPendingHistoryItem(
                  (prevItem: HistoryItemWithoutId | null) => {
                    if (prevItem?.type === 'tool_group') {
                      return {
                        ...prevItem,
                        tools: prevItem.tools.map(
                          (toolDisplay: IndividualToolCallDisplay) =>
                            toolDisplay.callId === callId &&
                            toolDisplay.status === ToolCallStatus.Executing
                              ? {
                                  ...toolDisplay,
                                  resultDisplay: output,
                                }
                              : toolDisplay,
                        ),
                      };
                    }
                    return prevItem;
                  },
                );
                // Also update the toolCall itself so that mapToDisplay
                // can pick up the live output if the item is not pending
                // (e.g. if it's being re-rendered from history)
                setToolCalls((prevToolCalls) =>
                  prevToolCalls.map((tc) =>
                    tc.request.callId === callId && tc.status === 'executing'
                      ? { ...tc, liveOutput: output }
                      : tc,
                  ),
                );
              }
            : undefined;

          t.tool
            .execute(t.request.args, signal, updateOutput)
            .then((result: ToolResult) => {
              if (signal.aborted) {
                // TODO(jacobr): avoid stringifying the LLM content.
                setToolCalls(
                  setStatus(callId, 'cancelled', String(result.llmContent)),
                );
                return;
              }
              const { functionResponseJson, additionalParts } =
                formatLlmContentForFunctionResponse(result.llmContent);
              const functionResponse: Part = {
                functionResponse: {
                  name: t.request.name,
                  id: callId,
                  response: functionResponseJson,
                },
              };
              const response: ToolCallResponseInfo = {
                callId,
                responseParts: [functionResponse, ...additionalParts],
                resultDisplay: result.returnDisplay,
                error: undefined,
              };
              setToolCalls(setStatus(callId, 'success', response));
            })
            .catch((e: Error) =>
              setToolCalls(
                setStatus(
                  callId,
                  'error',
                  toolErrorResponse(
                    t.request,
                    e instanceof Error ? e : new Error(String(e)),
                  ),
                ),
              ),
            );
        });
    }
  }, [toolCalls, toolRegistry, abortController.signal, setPendingHistoryItem]);

  useEffect(() => {
    const allDone = toolCalls.every(
      (t) =>
        t.status === 'success' ||
        t.status === 'error' ||
        t.status === 'cancelled',
    );
    if (toolCalls.length && allDone) {
      setToolCalls([]);
      onComplete(toolCalls);
      setAbortController(() => new AbortController());
    }
  }, [toolCalls, onComplete]);

  return [toolCalls, schedule, cancel];
}

function setStatus(
  targetCallId: string,
  status: 'success',
  response: ToolCallResponseInfo,
): (t: ToolCall[]) => ToolCall[];
function setStatus(
  targetCallId: string,
  status: 'awaiting_approval',
  confirm: ToolCallConfirmationDetails,
): (t: ToolCall[]) => ToolCall[];
function setStatus(
  targetCallId: string,
  status: 'error',
  response: ToolCallResponseInfo,
): (t: ToolCall[]) => ToolCall[];
function setStatus(
  targetCallId: string,
  status: 'cancelled',
  reason: string,
): (t: ToolCall[]) => ToolCall[];
function setStatus(
  targetCallId: string,
  status: 'executing' | 'scheduled' | 'validating',
): (t: ToolCall[]) => ToolCall[];
function setStatus(
  targetCallId: string,
  status: Status,
  auxiliaryData?: unknown,
): (t: ToolCall[]) => ToolCall[] {
  return function (tc: ToolCall[]): ToolCall[] {
    return tc.map((t) => {
      if (t.request.callId !== targetCallId || t.status === 'error') {
        return t;
      }
      switch (status) {
        case 'success': {
          const next: SuccessfulToolCall = {
            ...t,
            status: 'success',
            response: auxiliaryData as ToolCallResponseInfo,
          };
          return next;
        }
        case 'error': {
          const next: ErroredToolCall = {
            ...t,
            status: 'error',
            response: auxiliaryData as ToolCallResponseInfo,
          };
          return next;
        }
        case 'awaiting_approval': {
          const next: WaitingToolCall = {
            ...t,
            status: 'awaiting_approval',
            confirmationDetails: auxiliaryData as ToolCallConfirmationDetails,
          };
          return next;
        }
        case 'scheduled': {
          const next: ScheduledToolCall = {
            ...t,
            status: 'scheduled',
          };
          return next;
        }
        case 'cancelled': {
          const next: CancelledToolCall = {
            ...t,
            status: 'cancelled',
            response: {
              callId: t.request.callId,
              responseParts: {
                functionResponse: {
                  id: t.request.callId,
                  name: t.request.name,
                  response: {
                    error: `[Operation Cancelled] Reason: ${auxiliaryData}`,
                  },
                },
              },
              resultDisplay: undefined,
              error: undefined,
            },
          };
          return next;
        }
        case 'validating': {
          const next: ValidatingToolCall = {
            ...(t as ValidatingToolCall), // Added type assertion for safety
            status: 'validating',
          };
          return next;
        }
        case 'executing': {
          const next: ExecutingToolCall = {
            ...t,
            status: 'executing',
          };
          return next;
        }
        default: {
          // ensures every case is checked for above
          const exhaustiveCheck: never = status;
          return exhaustiveCheck;
        }
      }
    });
  };
}

const toolErrorResponse = (
  request: ToolCallRequestInfo,
  error: Error,
): ToolCallResponseInfo => ({
  callId: request.callId,
  error,
  responseParts: {
    functionResponse: {
      id: request.callId,
      name: request.name,
      response: { error: error.message },
    },
  },
  resultDisplay: error.message,
});

function mapStatus(status: Status): ToolCallStatus {
  switch (status) {
    case 'validating':
      return ToolCallStatus.Executing;
    case 'awaiting_approval':
      return ToolCallStatus.Confirming;
    case 'executing':
      return ToolCallStatus.Executing;
    case 'success':
      return ToolCallStatus.Success;
    case 'cancelled':
      return ToolCallStatus.Canceled;
    case 'error':
      return ToolCallStatus.Error;
    case 'scheduled':
      return ToolCallStatus.Pending;
    default: {
      // ensures every case is checked for above
      const exhaustiveCheck: never = status;
      return exhaustiveCheck;
    }
  }
}

// convenient function for callers to map ToolCall back to a HistoryItem
export function mapToDisplay(
  tool: ToolCall[] | ToolCall,
): HistoryItemToolGroup {
  const tools = Array.isArray(tool) ? tool : [tool];
  const toolsDisplays = tools.map((t): IndividualToolCallDisplay => {
    switch (t.status) {
      case 'success':
        return {
          callId: t.request.callId,
          name: t.tool.displayName,
          description: t.tool.getDescription(t.request.args),
          resultDisplay: t.response.resultDisplay,
          status: mapStatus(t.status),
          confirmationDetails: undefined,
          renderOutputAsMarkdown: t.tool.isOutputMarkdown,
        };
      case 'error':
        return {
          callId: t.request.callId,
          name: t.request.name, // Use request.name as tool might be undefined
          description: '', // No description available if tool is undefined
          resultDisplay: t.response.resultDisplay,
          status: mapStatus(t.status),
          confirmationDetails: undefined,
          renderOutputAsMarkdown: false,
        };
      case 'cancelled':
        return {
          callId: t.request.callId,
          name: t.tool.displayName,
          description: t.tool.getDescription(t.request.args),
          resultDisplay: t.response.resultDisplay,
          status: mapStatus(t.status),
          confirmationDetails: undefined,
          renderOutputAsMarkdown: t.tool.isOutputMarkdown,
        };
      case 'awaiting_approval':
        return {
          callId: t.request.callId,
          name: t.tool.displayName,
          description: t.tool.getDescription(t.request.args),
          resultDisplay: undefined,
          status: mapStatus(t.status),
          confirmationDetails: t.confirmationDetails,
          renderOutputAsMarkdown: t.tool.isOutputMarkdown,
        };
      case 'executing':
        return {
          callId: t.request.callId,
          name: t.tool.displayName,
          description: t.tool.getDescription(t.request.args),
          resultDisplay: t.liveOutput ?? undefined,
          status: mapStatus(t.status),
          confirmationDetails: undefined,
          renderOutputAsMarkdown: t.tool.isOutputMarkdown,
        };
      case 'validating': // Add this case
        return {
          callId: t.request.callId,
          name: t.tool.displayName,
          description: t.tool.getDescription(t.request.args),
          resultDisplay: undefined,
          status: mapStatus(t.status),
          confirmationDetails: undefined,
          renderOutputAsMarkdown: t.tool.isOutputMarkdown,
        };
      case 'scheduled':
        return {
          callId: t.request.callId,
          name: t.tool.displayName,
          description: t.tool.getDescription(t.request.args),
          resultDisplay: undefined,
          status: mapStatus(t.status),
          confirmationDetails: undefined,
          renderOutputAsMarkdown: t.tool.isOutputMarkdown,
        };
      default: {
        // ensures every case is checked for above
        const exhaustiveCheck: never = t;
        return exhaustiveCheck;
      }
    }
  });
  const historyItem: HistoryItemToolGroup = {
    type: 'tool_group',
    tools: toolsDisplays,
  };
  return historyItem;
}