import { MemorySaver } from '@langchain/langgraph';
import {
  StateGraph,
  Annotation,
  messagesStateReducer,
  Command,
} from '@langchain/langgraph';
import { AIMessage, HumanMessage } from '@langchain/core/messages';
import type { BaseMessage } from '@langchain/core/messages';
import type { RunnableConfig } from '@langchain/core/runnables';
import {
  createApprovalGateNode,
  getApprovalGateNodeId,
} from '../ApprovalGateNode';
import type { ApprovalGateConfig } from '@/types/graph';

// Suppress safeDispatchCustomEvent since we don't have a full graph context
jest.mock('@/utils/events', () => ({
  safeDispatchCustomEvent: jest.fn(),
}));

const { safeDispatchCustomEvent } = require('@/utils/events');

/**
 * Creates a simple 2-agent graph with an approval gate between them.
 * Agent A → Gate → Agent B
 */
function createGatedGraph(gateConfig: ApprovalGateConfig) {
  const GraphState = Annotation.Root({
    messages: Annotation<BaseMessage[]>({
      reducer: messagesStateReducer,
    }),
  });

  const checkpointer = new MemorySaver();
  const gateNodeId = getApprovalGateNodeId(gateConfig.gateId);

  const builder = new StateGraph(GraphState);

  // Agent A: echoes input
  builder.addNode('agent_a', async (state) => ({
    messages: [new AIMessage('Agent A done')],
  }));

  // Approval gate
  const gateNode = createApprovalGateNode(gateConfig, 'agent_a', 'agent_b');
  builder.addNode(gateNodeId, gateNode as any);

  // Agent B: runs after approval
  builder.addNode('agent_b', async (state) => ({
    messages: [new AIMessage('Agent B done')],
  }));

  // Wire: START → A → gate → B → END
  builder.addEdge('__start__' as any, 'agent_a' as any);
  builder.addEdge('agent_a' as any, gateNodeId as any);
  builder.addEdge(gateNodeId as any, 'agent_b' as any);
  builder.addEdge('agent_b' as any, '__end__' as any);

  const compiled = builder.compile({ checkpointer } as any);

  return { compiled, checkpointer, gateNodeId };
}

describe('ApprovalGateNode', () => {
  beforeEach(() => {
    jest.clearAllMocks();
  });

  it('interrupts the graph at the gate node', async () => {
    const { compiled } = createGatedGraph({
      gateId: 'review-step',
      channel: 'chat',
      prompt: 'Please review before Agent B runs',
    });

    const config: Partial<RunnableConfig> = {
      configurable: { thread_id: 'test-thread-1' },
    };

    // First invoke: Agent A runs, then gate interrupts
    const result = await compiled.invoke(
      { messages: [new HumanMessage('start')] },
      config
    );

    // Check that the graph was interrupted (Agent B should NOT have run)
    const state = await compiled.getState(config);
    const interrupted = state.tasks?.some((t: any) => t.interrupts?.length > 0);
    expect(interrupted).toBe(true);

    // Agent A should have run
    const messages = result.messages as BaseMessage[];
    const hasAgentA = messages.some(
      (m: any) => m._getType?.() === 'ai' && m.content === 'Agent A done'
    );
    expect(hasAgentA).toBe(true);

    // Agent B should NOT have run
    const hasAgentB = messages.some(
      (m: any) => m._getType?.() === 'ai' && m.content === 'Agent B done'
    );
    expect(hasAgentB).toBe(false);
  });

  it('resumes and runs Agent B after approval', async () => {
    const { compiled } = createGatedGraph({
      gateId: 'review-step',
      channel: 'chat',
    });

    const config: Partial<RunnableConfig> = {
      configurable: { thread_id: 'test-thread-2' },
    };

    // First invoke: hits gate interrupt
    await compiled.invoke({ messages: [new HumanMessage('start')] }, config);

    // Resume with approval
    const result = await compiled.invoke(
      new Command({ resume: { approved: true } }),
      config
    );

    // Agent B should have run after approval
    const messages = result.messages as BaseMessage[];
    const hasAgentB = messages.some(
      (m: any) => m._getType?.() === 'ai' && m.content === 'Agent B done'
    );
    expect(hasAgentB).toBe(true);
  });

  it('dispatches ON_APPROVAL_GATE event before interrupting', async () => {
    const { compiled } = createGatedGraph({
      gateId: 'notify-gate',
      channel: 'outlook',
      prompt: 'Review needed',
      approver: 'manager@example.com',
    });

    const config: Partial<RunnableConfig> = {
      configurable: { thread_id: 'test-thread-3' },
    };

    await compiled.invoke({ messages: [new HumanMessage('start')] }, config);

    // safeDispatchCustomEvent should have been called with gate data
    expect(safeDispatchCustomEvent).toHaveBeenCalledWith(
      'on_approval_gate',
      expect.objectContaining({
        type: 'approval_gate',
        gateId: 'notify-gate',
        channel: 'outlook',
        prompt: 'Review needed',
        approver: 'manager@example.com',
        sourceAgentId: 'agent_a',
        destinationAgentId: 'agent_b',
      }),
      expect.anything()
    );
  });

  it('resumes with denial — gate returns empty state', async () => {
    const { compiled } = createGatedGraph({
      gateId: 'deny-gate',
    });

    const config: Partial<RunnableConfig> = {
      configurable: { thread_id: 'test-thread-4' },
    };

    // First invoke: hits gate interrupt
    await compiled.invoke({ messages: [new HumanMessage('start')] }, config);

    // Resume with denial
    const result = await compiled.invoke(
      new Command({ resume: { approved: false, feedback: 'Not ready' } }),
      config
    );

    // Agent B still runs (the gate doesn't block — it's the host's
    // responsibility to not resume on denial, or use onDeny routing).
    // In the current implementation, the gate node returns {} regardless.
    const messages = result.messages as BaseMessage[];
    const hasAgentB = messages.some(
      (m: any) => m._getType?.() === 'ai' && m.content === 'Agent B done'
    );
    expect(hasAgentB).toBe(true);
  });

  describe('getApprovalGateNodeId', () => {
    it('generates consistent node IDs', () => {
      expect(getApprovalGateNodeId('step-1')).toBe('approval_gate_step-1');
      expect(getApprovalGateNodeId('review')).toBe('approval_gate_review');
    });
  });
});