/* eslint-disable no-console */
// src/scripts/cli.ts
import { config } from 'dotenv';
config();
import { HumanMessage, BaseMessage } from '@langchain/core/messages';
import type * as t from '@/types';
import { ChatModelStreamHandler, createContentAggregator } from '@/stream';
import { ToolEndHandler, ModelEndHandler } from '@/events';
import { createSearchTool } from '@/tools/search';

import { getArgs } from '@/scripts/args';
import { Run } from '@/run';
import { GraphEvents, Callback, Providers } from '@/common';
import { getLLMConfig } from '@/utils/llmConfig';

const conversationHistory: BaseMessage[] = [];
async function testStandardStreaming(): Promise<void> {
  const { userName, location, provider, currentDate } = await getArgs();
  const { contentParts, aggregateContent } = createContentAggregator();
  const customHandlers = {
    [GraphEvents.TOOL_END]: new ToolEndHandler(),
    [GraphEvents.CHAT_MODEL_END]: new ModelEndHandler(),
    [GraphEvents.CHAT_MODEL_STREAM]: new ChatModelStreamHandler(),
    [GraphEvents.ON_RUN_STEP_COMPLETED]: {
      handle: (
        event: GraphEvents.ON_RUN_STEP_COMPLETED,
        data: t.StreamEventData
      ): void => {
        console.log('====== ON_RUN_STEP_COMPLETED ======');
        // console.dir(data, { depth: null });
        aggregateContent({
          event,
          data: data as unknown as { result: t.ToolEndEvent },
        });
      },
    },
    [GraphEvents.ON_RUN_STEP]: {
      handle: (
        event: GraphEvents.ON_RUN_STEP,
        data: t.StreamEventData
      ): void => {
        aggregateContent({ event, data: data as t.RunStep });
      },
    },
    [GraphEvents.ON_RUN_STEP_DELTA]: {
      handle: (
        event: GraphEvents.ON_RUN_STEP_DELTA,
        data: t.StreamEventData
      ): void => {
        aggregateContent({ event, data: data as t.RunStepDeltaEvent });
      },
    },
    [GraphEvents.ON_MESSAGE_DELTA]: {
      handle: (
        event: GraphEvents.ON_MESSAGE_DELTA,
        data: t.StreamEventData
      ): void => {
        aggregateContent({ event, data: data as t.MessageDeltaEvent });
      },
    },
    [GraphEvents.TOOL_START]: {
      handle: (
        _event: string,
        data: t.StreamEventData,
        metadata?: Record<string, unknown>
      ): void => {
        console.log('====== TOOL_START ======');
        // console.dir(data, { depth: null });
      },
    },
  };

  const llmConfig = getLLMConfig(provider);

  if (llmConfig.provider === Providers.BEDROCK) {
    (llmConfig as t.BedrockAnthropicInput).promptCache = true;
  }

  const run = await Run.create<t.IState>({
    runId: 'test-run-id',
    graphConfig: {
      type: 'standard',
      llmConfig,
      tools: [
        createSearchTool({
          scraperProvider: 'serper',
        }),
      ],
      instructions:
        'You are a friendly AI assistant. Always address the user by their name.',
      // additional_instructions: `The user's name is ${userName} and they are located in ${location}.`,
    },
    returnContent: true,
    skipCleanup: true,
    customHandlers,
  });

  const config = {
    configurable: {
      provider,
      thread_id: 'conversation-num-1',
    },
    streamMode: 'values',
    version: 'v2' as const,
  };

  console.log('Test 1: Search query (search tool test)');

  // const userMessage = `
  // Make a search for the weather in ${location} today, which is ${currentDate}.
  // Before making the search, please let me know what you're about to do, then immediately start searching without hesitation.
  // Make sure to always refer to me by name, which is ${userName}.
  // After giving me a thorough summary, tell me a joke about the weather forecast we went over.
  // `;
  // const userMessage = 'Are massage guns good?';
  // const userMessage = 'What is functional programming?';
  const userMessage = "Get me today's trending news.";
  // const userMessage = "search recent italy earthquake volcano activity";
  // const userMessage =
  //   "use 'Trump' as the exact search query and tell me what you find.";

  conversationHistory.push(new HumanMessage(userMessage));

  const inputs = {
    messages: conversationHistory,
  };
  const finalContentParts = await run.processStream(inputs, config);
  const finalMessages = run.getRunMessages();
  if (finalMessages) {
    conversationHistory.push(...finalMessages);
    console.dir(conversationHistory, { depth: null });
  }
  // console.dir(finalContentParts, { depth: null });
  console.log('\n\n====================\n\n');
  // console.dir(contentParts, { depth: null });
}

process.on('unhandledRejection', (reason, promise) => {
  console.warn('Unhandled Rejection (non-fatal):', reason);
});

testStandardStreaming().catch((err) => {
  console.error(err);
  console.log('Conversation history:');
  console.dir(conversationHistory, { depth: null });
  process.exit(1);
});