/**
 * useEntityExecution - Unified hook for executing agents, orchestrators, and workflows
 *
 * Consolidates all streaming execution logic into a single, reusable hook.
 * Uses strategy pattern for entity-specific message handling.
 */

import { useState, useCallback, useRef, useEffect } from "react";
import { apiClient } from "@/services/api";
import type {
  Message,
  StreamEvent,
  RunEntityRequest,
  SessionInfo,
  SessionMessagesResponse,
  ToolApprovalRequest,
  ToolApprovalResponse,
} from "@/types";

function replaceTrailingAssistantOrAppend(
  messages: Message[],
  nextMessage: Message
): Message[] {
  const lastMessage = messages[messages.length - 1];
  if (lastMessage?.role === "assistant") {
    return [...messages.slice(0, -1), nextMessage];
  }
  return [...messages, nextMessage];
}

function calculateUsageTotals(
  messages: Message[]
): { tokens_input: number; tokens_output: number; cost_estimate: number } {
  return messages.reduce(
    (acc, msg) => {
      if (msg.role === "assistant" && (msg as any).usage) {
        const usage = (msg as any).usage;
        return {
          tokens_input: acc.tokens_input + (usage.tokens_input || 0),
          tokens_output: acc.tokens_output + (usage.tokens_output || 0),
          cost_estimate:
            (acc.cost_estimate || 0) + (usage.cost_estimate || 0),
        };
      }
      return acc;
    },
    { tokens_input: 0, tokens_output: 0, cost_estimate: 0 }
  );
}

function mergeUsageFromCurrentMessages(
  refreshedMessages: Message[],
  currentMessages: Message[]
): Message[] {
  return refreshedMessages.map((message, index) => {
    if (message.role !== "assistant" || (message as any).usage) {
      return message;
    }

    const fallbackMatch =
      currentMessages[index] &&
      currentMessages[index].role === "assistant" &&
      currentMessages[index].source === message.source &&
      currentMessages[index].content === message.content
        ? currentMessages[index]
        : currentMessages.find(
            (candidate) =>
              candidate.role === "assistant" &&
              candidate.source === message.source &&
              candidate.content === message.content &&
              (candidate as any).usage
          );

    if (fallbackMatch && (fallbackMatch as any).usage) {
      return {
        ...message,
        usage: (fallbackMatch as any).usage,
      };
    }

    return message;
  });
}

export type EntityType = "agent" | "orchestrator" | "workflow";

export interface MessageHandler {
  /**
   * Process a stream event and return updated messages array.
   * Handler maintains internal state for streaming logic.
   */
  handleEvent(
    event: StreamEvent,
    currentMessages: Message[],
    entityName: string
  ): Message[];

  /**
   * Reset handler state for new execution
   */
  reset(): void;
}

export interface UseEntityExecutionOptions {
  entityId: string;
  entityType: EntityType;
  entityName: string;
  userId?: string;
  currentSession?: SessionInfo;
  onDebugEvent: (event: StreamEvent) => void;
  onSessionChange?: (session: SessionInfo) => void;
  messageHandler: MessageHandler;
  supportsToolApproval?: boolean;
  supportsTokenStreaming?: boolean;
}

export interface UseEntityExecutionReturn {
  messages: Message[];
  isStreaming: boolean;
  sessionTotalUsage: {
    tokens_input: number;
    tokens_output: number;
    cost_estimate: number;
  };
  pendingApproval: ToolApprovalRequest | null;
  currentAgentSpeaking: string | null;
  finalResult: string | null;
  stopReason: string | null;

  handleSendMessage: (
    newMessages: Message[],
    approvalResponses?: ToolApprovalResponse[]
  ) => Promise<void>;
  handleStop: () => void;
  handleClearMessages: () => void;
  handleApprove: (response: ToolApprovalResponse) => void;
  handleReject: (response: ToolApprovalResponse) => void;
}

export function useEntityExecution(
  options: UseEntityExecutionOptions
): UseEntityExecutionReturn {
  const {
    entityId,
    entityType,
    entityName,
    userId,
    currentSession,
    onDebugEvent,
    onSessionChange,
    messageHandler,
    supportsToolApproval = false,
    supportsTokenStreaming = false,
  } = options;

  const [messages, setMessages] = useState<Message[]>([]);
  const [isStreaming, setIsStreaming] = useState(false);
  const [sessionTotalUsage, setSessionTotalUsage] = useState<{
    tokens_input: number;
    tokens_output: number;
    cost_estimate: number;
  }>({ tokens_input: 0, tokens_output: 0, cost_estimate: 0 });
  const [pendingApproval, setPendingApproval] =
    useState<ToolApprovalRequest | null>(null);
  const [pendingApprovalResponses, setPendingApprovalResponses] = useState<
    ToolApprovalResponse[]
  >([]);
  const [currentAgentSpeaking, setCurrentAgentSpeaking] = useState<
    string | null
  >(null);
  const [finalResult, setFinalResult] = useState<string | null>(null);
  const [stopReason, setStopReason] = useState<string | null>(null);

  const abortControllerRef = useRef<AbortController | null>(null);

  const applySessionMessagesResponse = useCallback(
    (
      response: SessionMessagesResponse,
      currentMessagesForMerge?: Message[]
    ) => {
      if (currentMessagesForMerge) {
        setMessages(() => {
          const mergedMessages = mergeUsageFromCurrentMessages(
            response.messages,
            currentMessagesForMerge
          );
          setSessionTotalUsage(calculateUsageTotals(mergedMessages));
          return mergedMessages;
        });
      } else {
        setMessages(response.messages);
        setSessionTotalUsage(calculateUsageTotals(response.messages));
      }
      setFinalResult(response.final_result ?? null);
      setStopReason(response.stop_reason ?? null);
    },
    []
  );

  // Load session messages when currentSession changes
  useEffect(() => {
    const loadSessionMessages = async () => {
      if (currentSession) {
        try {
          const response = await apiClient.getSessionMessages(currentSession.id);
          applySessionMessagesResponse(response);
        } catch (error) {
          console.error("Failed to load session messages:", error);
          setMessages([]);
          setSessionTotalUsage({
            tokens_input: 0,
            tokens_output: 0,
            cost_estimate: 0,
          });
          setFinalResult(null);
          setStopReason(null);
        }
      } else {
        // No session - clear messages and usage
        setMessages([]);
        setSessionTotalUsage({
          tokens_input: 0,
          tokens_output: 0,
          cost_estimate: 0,
        });
        setFinalResult(null);
        setStopReason(null);
      }
      setPendingApproval(null);
      setPendingApprovalResponses([]);
    };

    loadSessionMessages();
  }, [currentSession?.id, applySessionMessagesResponse]);

  const handleSendMessage = useCallback(
    async (
      newMessages: Message[],
      approvalResponses?: ToolApprovalResponse[]
    ) => {
      // Add new messages to state
      setMessages((prev) => [...prev, ...newMessages]);
      setIsStreaming(true);
      setCurrentAgentSpeaking(null);
      setFinalResult(null);
      setStopReason(null);

      // Create new AbortController for this request
      abortControllerRef.current = new AbortController();

      // Use provided approval responses or fallback to state
      const approvalsToSend =
        approvalResponses ||
        (pendingApprovalResponses.length > 0
          ? pendingApprovalResponses
          : undefined);

      let activeSessionId = currentSession?.id;
      let didComplete = false;

      try {
        const request: RunEntityRequest = {
          messages: newMessages, // Send only NEW messages
          session_id: currentSession?.id, // Backend will append to session
          user_id: userId,
          stream_tokens: supportsTokenStreaming,
          approval_responses: approvalsToSend,
        };

        // Clear pending approvals after sending
        if (approvalsToSend) {
          setPendingApprovalResponses([]);
        }

        // Reset message handler for new execution
        messageHandler.reset();

        // Stream execution
        for await (const event of apiClient.streamEntityExecution(
          entityId,
          request,
          abortControllerRef.current.signal
        )) {
          onDebugEvent(event);

          // Capture session_id from first event and update parent
          if (!currentSession && event.session_id && onSessionChange) {
            activeSessionId = event.session_id;
            const newSession: SessionInfo = {
              id: event.session_id,
              entity_id: entityId,
              entity_type: entityType,
              created_at: new Date().toISOString(),
              last_activity: new Date().toISOString(),
              message_count: 0,
              user_id: userId,
            };
            onSessionChange(newSession);
          }
          if (event.session_id) {
            activeSessionId = event.session_id;
          }

          // Track current agent speaking (for orchestrators)
          if (event.type === "agent_selection" && event.data?.selected_agent) {
            setCurrentAgentSpeaking(event.data.selected_agent);
          } else if (event.type === "agent_execution_complete") {
            setCurrentAgentSpeaking(null);
          }

          if (event.type === "orchestration_complete") {
            // Use the streamed terminal event for immediate UX, then let the
            // post-run session refresh replace it with the persisted value.
            setFinalResult(event.data?.result ?? null);
            setStopReason(event.data?.stop_reason ?? null);
          }

          // Handle tool approval requests
          if (
            supportsToolApproval &&
            event.type === "tool_approval" &&
            event.data?.approval_request
          ) {
            setPendingApproval(event.data.approval_request);
          }

          // Let message handler process the event
          setMessages((prevMessages) =>
            messageHandler.handleEvent(event, prevMessages, entityName)
          );

          // Handle completion - extract usage
          if (event.type === "complete" && event.data?.usage) {
            const usage = event.data.usage;
            setSessionTotalUsage((prev) => ({
              tokens_input: prev.tokens_input + (usage.tokens_input || 0),
              tokens_output: prev.tokens_output + (usage.tokens_output || 0),
              cost_estimate:
                (prev.cost_estimate || 0) + (usage.cost_estimate || 0),
            }));
            didComplete = true;
            break;
          }
        }

        if (didComplete && activeSessionId) {
          try {
            const response = await apiClient.getSessionMessages(activeSessionId);
            applySessionMessagesResponse(response, messages);
          } catch (refreshError) {
            console.error("Failed to refresh session messages after completion:", refreshError);
          }
        }
      } catch (error) {
        console.error("Failed to send message:", error);

        // Check if this was an abort (user clicked stop)
        if (error instanceof Error && error.name === "AbortError") {
          const cancelMessage: Message = {
            role: "assistant",
            content: "Cancelled by user",
            source: "system",
          };
          setMessages((prev) =>
            replaceTrailingAssistantOrAppend(prev, cancelMessage)
          );
        } else {
          const errorMessage: Message = {
            role: "assistant",
            content: `Error: ${
              error instanceof Error ? error.message : "Unknown error"
            }`,
            source: "system",
          };
          setMessages((prev) =>
            replaceTrailingAssistantOrAppend(prev, errorMessage)
          );
        }
      } finally {
        setIsStreaming(false);
        setCurrentAgentSpeaking(null);
        abortControllerRef.current = null;
      }
    },
    [
      entityId,
      entityType,
      entityName,
      userId,
      currentSession,
      pendingApprovalResponses,
      onDebugEvent,
      onSessionChange,
      messageHandler,
      supportsToolApproval,
      supportsTokenStreaming,
      applySessionMessagesResponse,
      messages,
    ]
  );

  const handleStop = useCallback(() => {
    if (abortControllerRef.current) {
      console.log(`🛑 Stopping ${entityType} execution`);
      abortControllerRef.current.abort();
    }
  }, [entityType]);

  const handleClearMessages = useCallback(() => {
    setMessages([]);
    setSessionTotalUsage({
      tokens_input: 0,
      tokens_output: 0,
      cost_estimate: 0,
    });
    setCurrentAgentSpeaking(null);
  }, []);

  const handleApprove = useCallback(
    (response: ToolApprovalResponse) => {
      console.log("📝 handleApprove called with:", response);
      setPendingApproval(null);
      handleSendMessage([], [response]);
    },
    [handleSendMessage]
  );

  const handleReject = useCallback(
    (response: ToolApprovalResponse) => {
      console.log("📝 handleReject called with:", response);
      setPendingApproval(null);
      handleSendMessage([], [response]);
    },
    [handleSendMessage]
  );

  return {
    messages,
    isStreaming,
    sessionTotalUsage,
    pendingApproval,
    currentAgentSpeaking,
    finalResult,
    stopReason,
    handleSendMessage,
    handleStop,
    handleClearMessages,
    handleApprove,
    handleReject,
  };
}
