import { useMemo } from 'react';

import { SYSTEM_MSG_TAG } from '@/lib/shared/constants/system-message-tags';

import type {
  DocumentWriteApproval,
  HumanInputRequest,
  IntegrationApproval,
  KnowledgeWriteApproval,
  LocationRequest,
  WorkflowCreationApproval,
  WorkflowRunApproval,
  WorkflowUpdateApproval,
} from './queries';
import type { ChatMessage } from './use-message-processing';

export type ChatItem =
  | { type: 'message'; data: ChatMessage }
  | { type: 'approval'; data: IntegrationApproval }
  | { type: 'workflow_approval'; data: WorkflowCreationApproval }
  | { type: 'workflow_update_approval'; data: WorkflowUpdateApproval }
  | { type: 'workflow_run_approval'; data: WorkflowRunApproval }
  | { type: 'human_input_request'; data: HumanInputRequest }
  | { type: 'document_write_approval'; data: DocumentWriteApproval }
  | { type: 'knowledge_write_approval'; data: KnowledgeWriteApproval }
  | { type: 'location_request'; data: LocationRequest };

type ApprovalChatItem = Exclude<ChatItem, { type: 'message' }>;

interface UseMergedChatItemsParams {
  messages: ChatMessage[];
  integrationApprovals: IntegrationApproval[] | undefined;
  workflowCreationApprovals: WorkflowCreationApproval[] | undefined;
  workflowUpdateApprovals: WorkflowUpdateApproval[] | undefined;
  workflowRunApprovals: WorkflowRunApproval[] | undefined;
  humanInputRequests: HumanInputRequest[] | undefined;
  /** Completed/rejected human-input requests — rendered INLINE at their
   *  chronological slot (after their source message) so an answered request
   *  stays visible and editable in the history. */
  resolvedHumanInputRequests?: HumanInputRequest[];
  locationRequests: LocationRequest[] | undefined;
  documentWriteApprovals: DocumentWriteApproval[] | undefined;
  /** Optional — surfaces only on the main chat path (the `knowledge_write`
   *  tool); arena/automation callers don't pass it. */
  knowledgeWriteApprovals?: KnowledgeWriteApproval[];
}

export interface MergedChatItemsResult {
  messages: ChatItem[];
  activeApproval: ChatItem | null;
  /** True when `activeApproval` already renders INLINE in `messages` (human
   *  input anchored to its source message) — the footer card must then be
   *  skipped, or the request would show twice. Keeping the active card
   *  inline gives it ONE stable DOM slot through pending → executing →
   *  completed, instead of hopping from the footer into the history. */
  activeApprovalInline: boolean;
}

function isActiveStatus(status: string) {
  return status === 'pending' || status === 'executing';
}

/**
 * Splice human-input requests (active AND resolved) into the chronologically
 * sorted message items, each immediately AFTER the message whose id matches
 * its `messageId` (anchor-to-source: the request must appear under the
 * assistant turn that asked it, even when later messages carry earlier
 * client clocks). When at least one card is inserted, the duplicate
 * `[HUMAN_INPUT_RESPONSE]` pill system messages are dropped — the card shows
 * the (latest) response itself. Requests whose source message isn't loaded
 * (pagination) are skipped; they appear once load-more brings the message
 * in. Pure + exported for unit testing.
 */
export function mergeHumanInputItems(
  messageItems: ChatItem[],
  requests: HumanInputRequest[],
  /** Whether the resolved-requests subscription has loaded. Pills are only
   *  suppressed once it has — suppressing while it's still in flight would
   *  briefly hide previously answered Q&As (their cards aren't in yet). */
  suppressPills: boolean,
): ChatItem[] {
  if (requests.length === 0) return messageItems;

  const byMessageId = new Map<string, HumanInputRequest[]>();
  for (const request of requests) {
    if (!request.messageId) continue;
    const bucket = byMessageId.get(request.messageId);
    if (bucket) bucket.push(request);
    else byMessageId.set(request.messageId, [request]);
  }
  for (const bucket of byMessageId.values()) {
    bucket.sort((a, b) => a._creationTime - b._creationTime);
  }

  const merged: ChatItem[] = [];
  let insertedAny = false;
  for (const item of messageItems) {
    merged.push(item);
    if (item.type !== 'message') continue;
    const bucket = byMessageId.get(item.data.id);
    if (!bucket) continue;
    for (const request of bucket) {
      merged.push({ type: 'human_input_request', data: request });
      insertedAny = true;
    }
  }
  if (!insertedAny) return messageItems;
  if (!suppressPills) return merged;

  // The inline card displays the response — the pill would repeat it (and
  // every edit appends ANOTHER "corrected answer supersedes" pill). Only
  // suppressed when a card actually rendered, so the pill survives as a
  // fallback when approval data is missing.
  return merged.filter(
    (item) =>
      item.type !== 'message' ||
      item.data.role !== 'system' ||
      !item.data.content.startsWith(SYSTEM_MSG_TAG.HUMAN_INPUT_RESPONSE),
  );
}

/**
 * Hook to merge messages with approvals.
 * Messages are returned chronologically.
 * The latest active (pending/executing) approval is returned separately;
 * resolved human-input requests are merged INLINE after their source message
 * (other completed/rejected approvals stay hidden).
 */
export function useMergedChatItems({
  messages,
  integrationApprovals,
  workflowCreationApprovals,
  workflowUpdateApprovals,
  workflowRunApprovals,
  humanInputRequests,
  resolvedHumanInputRequests,
  locationRequests,
  documentWriteApprovals,
  knowledgeWriteApprovals,
}: UseMergedChatItemsParams): MergedChatItemsResult {
  return useMemo((): MergedChatItemsResult => {
    // Build message items
    const loadedMessageIds = new Set();
    for (const message of messages || []) {
      loadedMessageIds.add(message.id);
    }

    const messageItems: ChatItem[] = (messages || []).map((message) => ({
      type: 'message' as const,
      data: message,
    }));

    // Sort messages chronologically
    messageItems.sort((a, b) => {
      if (a.type !== 'message' || b.type !== 'message') return 0;
      const aTime = a.data._creationTime ?? a.data.timestamp.getTime();
      const bTime = b.data._creationTime ?? b.data.timestamp.getTime();
      return aTime - bTime;
    });

    // Inline human-input cards — ACTIVE ones too, so the card occupies one
    // stable DOM slot through pending → executing → completed instead of
    // hopping from the footer into the history at completion. Resolved rows
    // whose id is still in the ACTIVE set are skipped (status-flip race
    // between the two subscriptions) so a request never renders two cards.
    const activeHumanInputIds = new Set(
      (humanInputRequests ?? []).map((a) => a._id),
    );
    const humanInputToInline = [
      ...(humanInputRequests ?? []),
      ...(resolvedHumanInputRequests ?? []).filter(
        (a) => !activeHumanInputIds.has(a._id),
      ),
    ];
    const itemsWithHumanInput = mergeHumanInputItems(
      messageItems,
      humanInputToInline,
      resolvedHumanInputRequests !== undefined,
    );

    // Collect active approvals (pending/executing only, linked to loaded messages)
    const activeApprovals: ApprovalChatItem[] = [];

    for (const a of integrationApprovals ?? []) {
      if (
        a.messageId &&
        loadedMessageIds.has(a.messageId) &&
        isActiveStatus(a.status)
      ) {
        activeApprovals.push({ type: 'approval', data: a });
      }
    }
    for (const a of workflowCreationApprovals ?? []) {
      if (
        a.messageId &&
        loadedMessageIds.has(a.messageId) &&
        isActiveStatus(a.status)
      ) {
        activeApprovals.push({ type: 'workflow_approval', data: a });
      }
    }
    for (const a of workflowUpdateApprovals ?? []) {
      if (
        a.messageId &&
        loadedMessageIds.has(a.messageId) &&
        isActiveStatus(a.status)
      ) {
        activeApprovals.push({ type: 'workflow_update_approval', data: a });
      }
    }
    for (const a of workflowRunApprovals ?? []) {
      if (
        a.messageId &&
        loadedMessageIds.has(a.messageId) &&
        isActiveStatus(a.status)
      ) {
        activeApprovals.push({ type: 'workflow_run_approval', data: a });
      }
    }
    for (const a of humanInputRequests ?? []) {
      if (
        a.messageId &&
        loadedMessageIds.has(a.messageId) &&
        isActiveStatus(a.status)
      ) {
        activeApprovals.push({ type: 'human_input_request', data: a });
      }
    }
    for (const a of locationRequests ?? []) {
      if (
        a.messageId &&
        loadedMessageIds.has(a.messageId) &&
        isActiveStatus(a.status)
      ) {
        activeApprovals.push({ type: 'location_request', data: a });
      }
    }
    for (const a of documentWriteApprovals ?? []) {
      if (
        a.messageId &&
        loadedMessageIds.has(a.messageId) &&
        isActiveStatus(a.status)
      ) {
        activeApprovals.push({ type: 'document_write_approval', data: a });
      }
    }
    for (const a of knowledgeWriteApprovals ?? []) {
      if (
        a.messageId &&
        loadedMessageIds.has(a.messageId) &&
        isActiveStatus(a.status)
      ) {
        activeApprovals.push({ type: 'knowledge_write_approval', data: a });
      }
    }

    // Pick the latest active approval by creation time
    let activeApproval: ChatItem | null = null;
    if (activeApprovals.length > 0) {
      activeApprovals.sort(
        (a, b) => b.data._creationTime - a.data._creationTime,
      );
      activeApproval = activeApprovals[0];
    }

    // Human-input requests anchored to a loaded message render inline (see
    // above) — the footer must skip them or the card shows twice.
    const activeApprovalInline =
      activeApproval !== null &&
      activeApproval.type === 'human_input_request' &&
      activeApproval.data.messageId !== undefined &&
      loadedMessageIds.has(activeApproval.data.messageId);

    return {
      messages: itemsWithHumanInput,
      activeApproval,
      activeApprovalInline,
    };
  }, [
    messages,
    integrationApprovals,
    workflowCreationApprovals,
    workflowUpdateApprovals,
    workflowRunApprovals,
    humanInputRequests,
    resolvedHumanInputRequests,
    locationRequests,
    documentWriteApprovals,
    knowledgeWriteApprovals,
  ]);
}
