// src/run.ts
import './instrumentation';
import { ObservabilityCallbackHandler } from '@illuma-ai/observability-langchain';
import { Command, isGraphInterrupt } from '@langchain/langgraph';
import { PromptTemplate } from '@langchain/core/prompts';
import { RunnableLambda } from '@langchain/core/runnables';
import { AzureChatOpenAI, ChatOpenAI } from '@langchain/openai';
import { BaseCallbackHandler } from '@langchain/core/callbacks/base';
import type {
  MessageContentComplex,
  BaseMessage,
} from '@langchain/core/messages';
import type { StringPromptValue } from '@langchain/core/prompt_values';
import type { RunnableConfig } from '@langchain/core/runnables';
import type * as t from '@/types';
import {
  createCompletionTitleRunnable,
  createTitleRunnable,
} from '@/utils/title';
import { createTokenCounter, encodingForModel } from '@/utils/tokens';
import { GraphEvents, Callback, TitleMethod } from '@/common';
import { MultiAgentGraph } from '@/graphs/MultiAgentGraph';
import { StandardGraph } from '@/graphs/Graph';
import { HandlerRegistry } from '@/events';
import { executeHooks } from '@/hooks';
import type { HookRegistry } from '@/hooks';
import { ToolOutputReferenceRegistry } from '@/tools/toolOutputReferences';
import { isOpenAILike } from '@/utils/llm';
import { isPresent } from '@/utils/misc';

function findLastMessageOfType(
  messages: BaseMessage[],
  type: string
): BaseMessage | undefined {
  for (let i = messages.length - 1; i >= 0; i--) {
    if (messages[i].getType() === type) {
      return messages[i];
    }
  }
  return undefined;
}

function extractPromptText(message: BaseMessage): string {
  const content = message.content;
  if (typeof content === 'string') {
    return content;
  }
  if (!Array.isArray(content)) {
    return String(content);
  }
  const parts: string[] = [];
  for (const block of content) {
    if (
      typeof block === 'object' &&
      'type' in block &&
      block.type === 'text' &&
      'text' in block &&
      typeof block.text === 'string'
    ) {
      parts.push(block.text);
    }
  }
  return parts.join('\n');
}

export const defaultOmitOptions = new Set([
  'stream',
  'thinking',
  'streaming',
  'maxTokens',
  'clientOptions',
  'thinkingConfig',
  'thinkingBudget',
  'includeThoughts',
  'maxOutputTokens',
  'additionalModelRequestFields',
]);

export class Run<_T extends t.BaseGraphState> {
  id: string;
  private tokenCounter?: t.TokenCounter;
  private handlerRegistry?: HandlerRegistry;
  private hookRegistry?: HookRegistry;
  private toolOutputRegistry?: ToolOutputReferenceRegistry;
  private indexTokenCountMap?: Record<string, number>;
  graphRunnable?: t.CompiledStateWorkflow;
  Graph: StandardGraph | MultiAgentGraph | undefined;
  returnContent: boolean = false;
  private skipCleanup: boolean = false;
  private _originalCallbacksSnapshot?: t.ProvidedCallbacks;

  private constructor(config: Partial<t.RunConfig>) {
    const runId = config.runId ?? '';
    if (!runId) {
      throw new Error('Run ID not provided');
    }

    this.id = runId;
    this.tokenCounter = config.tokenCounter;
    this.indexTokenCountMap = config.indexTokenCountMap;

    const handlerRegistry = new HandlerRegistry();

    if (config.customHandlers) {
      for (const [eventType, handler] of Object.entries(
        config.customHandlers
      )) {
        handlerRegistry.register(eventType, handler);
      }
    }

    this.handlerRegistry = handlerRegistry;
    /**
     * Capture hookRegistry BEFORE graph compilation so it propagates into
     * any ToolNode the graph builds during createWorkflow(). createLegacyGraph
     * threads `this.hookRegistry` onto the graph instance via Graph.hookRegistry
     * before compile, which Graph.createToolNodeForAgent reads when constructing
     * each ToolNode.
     */
    this.hookRegistry = config.hooks;
    /**
     * Same lifecycle for the tool-output reference registry: build it
     * here (when enabled) and let createLegacyGraph / createMultiAgentGraph
     * thread it onto the graph before compile so every ToolNode sees the
     * same instance.
     */
    if (config.toolOutputReferences?.enabled === true) {
      this.toolOutputRegistry = new ToolOutputReferenceRegistry({
        maxOutputSize: config.toolOutputReferences.maxOutputSize,
        maxTotalSize: config.toolOutputReferences.maxTotalSize,
      });
    }

    if (!config.graphConfig) {
      throw new Error('Graph config not provided');
    }

    /** Handle different graph types */
    if (config.graphConfig.type === 'multi-agent') {
      this.graphRunnable = this.createMultiAgentGraph(config.graphConfig);
      if (this.Graph) {
        this.Graph.handlerRegistry = handlerRegistry;
      }
    } else {
      /** Default to legacy graph for 'standard' or undefined type */
      this.graphRunnable = this.createLegacyGraph(config.graphConfig);
      if (this.Graph) {
        this.Graph.compileOptions =
          config.graphConfig.compileOptions ?? this.Graph.compileOptions;
        this.Graph.handlerRegistry = handlerRegistry;
      }
    }

    this.returnContent = config.returnContent ?? false;
    this.skipCleanup = config.skipCleanup ?? false;
  }

  private createLegacyGraph(
    config: t.LegacyGraphConfig | t.StandardGraphConfig
  ): t.CompiledStateWorkflow {
    let agentConfig: t.AgentInputs;
    let signal: AbortSignal | undefined;

    /** Check if this is a multi-agent style config (has agents array) */
    if ('agents' in config && Array.isArray(config.agents)) {
      if (config.agents.length === 0) {
        throw new Error('At least one agent must be provided');
      }
      agentConfig = config.agents[0];
      signal = config.signal;
    } else {
      /** Legacy path: build agent config from llmConfig */
      const {
        type: _type,
        llmConfig,
        signal: legacySignal,
        tools = [],
        ...agentInputs
      } = config as t.LegacyGraphConfig;
      const { provider, ...clientOptions } = llmConfig;

      agentConfig = {
        ...agentInputs,
        tools,
        provider,
        clientOptions,
        agentId: 'default',
      };
      signal = legacySignal;
    }

    const standardGraph = new StandardGraph({
      signal,
      runId: this.id,
      agents: [agentConfig],
      tokenCounter: this.tokenCounter,
      indexTokenCountMap: this.indexTokenCountMap,
    });
    /** Propagate compile options from graph config */
    standardGraph.compileOptions = config.compileOptions;
    if (this.hookRegistry) {
      standardGraph.hookRegistry = this.hookRegistry;
    }
    if (this.toolOutputRegistry) {
      standardGraph.toolOutputRegistry = this.toolOutputRegistry;
    }
    this.Graph = standardGraph;
    return standardGraph.createWorkflow();
  }

  private createMultiAgentGraph(
    config: t.MultiAgentGraphConfig
  ): t.CompiledStateWorkflow {
    const { agents, edges, compileOptions, resumeFromAgentId } = config;

    const multiAgentGraph = new MultiAgentGraph({
      runId: this.id,
      agents,
      edges,
      resumeFromAgentId,
      tokenCounter: this.tokenCounter,
      indexTokenCountMap: this.indexTokenCountMap,
    });

    if (compileOptions != null) {
      multiAgentGraph.compileOptions = compileOptions;
    }

    if (this.hookRegistry) {
      multiAgentGraph.hookRegistry = this.hookRegistry;
    }
    if (this.toolOutputRegistry) {
      multiAgentGraph.toolOutputRegistry = this.toolOutputRegistry;
    }
    this.Graph = multiAgentGraph;
    return multiAgentGraph.createWorkflow();
  }

  static async create<T extends t.BaseGraphState>(
    config: t.RunConfig
  ): Promise<Run<T>> {
    /** Create tokenCounter if indexTokenCountMap is provided but tokenCounter is not */
    if (config.indexTokenCountMap && !config.tokenCounter) {
      const gc = config.graphConfig;
      const clientOpts =
        'agents' in gc ? gc.agents[0]?.clientOptions : gc.clientOptions;
      const model = (clientOpts as { model?: string } | undefined)?.model ?? '';
      config.tokenCounter = await createTokenCounter(encodingForModel(model));
    }
    return new Run<T>(config);
  }

  getRunMessages(): BaseMessage[] | undefined {
    if (!this.Graph) {
      throw new Error(
        'Graph not initialized. Make sure to use Run.create() to instantiate the Run.'
      );
    }
    return this.Graph.getRunMessages();
  }

  /**
   * Manually trigger cleanup of heavy state (messages, config, etc.).
   * Call this after all continuations are complete when using skipCleanup=true.
   */
  clearState(): void {
    if (this.Graph) {
      this.Graph.clearHeavyState();
    }
  }

  /**
   * Returns the normalized finish/stop reason from the last LLM invocation.
   * Delegates to the underlying Graph instance.
   */
  getLastFinishReason(): string | undefined {
    if (this.Graph && 'getLastFinishReason' in this.Graph) {
      return (this.Graph as StandardGraph).getLastFinishReason();
    }
    return undefined;
  }

  /**
   * Returns the ID of the last agent that produced output in a multi-agent graph.
   * Used by auto-continuation to determine which agent's context to preserve
   * when a response is truncated after an agent handoff.
   * Returns undefined for single-agent graphs.
   */
  getLastActiveAgentId(): string | undefined {
    if (this.Graph && this.Graph instanceof MultiAgentGraph) {
      // Last entry in the run's participating agent set is the most-recent
      // active agent. Derived from contentData step trail rather than
      // a tracked field — same observation surface, fewer state slots.
      const ids = this.Graph.getActiveAgentIds();
      return ids.length > 0 ? ids[ids.length - 1] : undefined;
    }
    return undefined;
  }

  /**
   * Creates a custom event callback handler that intercepts custom events
   * and processes them through our handler registry instead of EventStreamCallbackHandler
   */
  private createCustomEventCallback() {
    return async (
      eventName: string,
      data: unknown,
      runId: string,
      tags?: string[],
      metadata?: Record<string, unknown>
    ): Promise<void> => {
      const handler = this.handlerRegistry?.getHandler(eventName);
      if (handler && this.Graph) {
        return await handler.handle(
          eventName,
          data as
            | t.StreamEventData
            | t.ModelEndData
            | t.RunStep
            | t.RunStepDeltaEvent
            | t.MessageDeltaEvent
            | t.ReasoningDeltaEvent
            | { result: t.ToolEndEvent },
          metadata,
          this.Graph
        );
      }
    };
  }

  /**
   * Processes the graph stream for a given input.
   *
   * @param inputs - Either the initial state (IState) for a new run, or a
   *   Command (e.g., `new Command({ resume: ... })`) to resume from an interrupt.
   * @param config - Runnable config with version and optional run_id.
   * @param streamOptions - Optional stream event callbacks and options.
   */
  async processStream(
    inputs: t.IState | Command,
    config: Partial<RunnableConfig> & { version: 'v1' | 'v2'; run_id?: string },
    streamOptions?: t.EventStreamOptions
  ): Promise<MessageContentComplex[] | undefined> {
    if (this.graphRunnable == null) {
      throw new Error(
        'Run not initialized. Make sure to use Run.create() to instantiate the Run.'
      );
    }
    if (!this.Graph) {
      throw new Error(
        'Graph not initialized. Make sure to use Run.create() to instantiate the Run.'
      );
    }

    this.Graph.resetValues(streamOptions?.keepContent);

    /** Custom event callback to intercept and handle custom events */
    const customEventCallback = this.createCustomEventCallback();

    // IMPORTANT: snapshot the ORIGINAL caller-provided callbacks once and
    // reuse that snapshot on every invocation. Previously we read
    // `config.callbacks` on each call and concat'd onto it — but the same
    // `config` object is passed back to processStream() on HITL resume, so
    // the previous run's appended callbacks became the new "base" and we
    // ended up stacking duplicate stream + custom handlers on every cycle.
    // That produced N-times duplicated message_delta events on the SSE
    // stream (character-interleaved text in the UI after tool approval).
    if (!this._originalCallbacksSnapshot) {
      this._originalCallbacksSnapshot =
        (config.callbacks as t.ProvidedCallbacks) ?? [];
    }
    const baseCallbacks = this._originalCallbacksSnapshot;
    const streamCallbacks = streamOptions?.callbacks
      ? this.getCallbacks(streamOptions.callbacks)
      : [];

    const customHandler = BaseCallbackHandler.fromMethods({
      [Callback.CUSTOM_EVENT]: customEventCallback,
    });
    customHandler.awaitHandlers = true;

    config.callbacks = baseCallbacks
      .concat(streamCallbacks)
      .concat(customHandler);

    const illumaSecretKey = process.env.ILLUMA_SECRET_KEY;
    const illumaPublicKey = process.env.ILLUMA_PUBLIC_KEY;
    const illumaBaseUrl = process.env.ILLUMA_BASE_URL;

    if (
      isPresent(illumaSecretKey) &&
      isPresent(illumaPublicKey) &&
      isPresent(illumaBaseUrl)
    ) {
      try {
        const userId = config.configurable?.user_id;
        const sessionId = config.configurable?.thread_id;
        const handler = new ObservabilityCallbackHandler({
          clientOptions: {
            secretKey: illumaSecretKey!,
            publicKey: illumaPublicKey!,
            baseUrl: illumaBaseUrl!,
          },
          userId,
          sessionId,
          metadata: {
            messageId: this.id,
            parentMessageId: config.configurable?.requestBody?.parentMessageId,
          },
        });
        config.callbacks = (
          (config.callbacks as t.ProvidedCallbacks) ?? []
        ).concat([handler]);
      } catch {
        // Gracefully skip if @illuma-ai/observability-node is not installed
      }
    }

    if (!this.id) {
      throw new Error('Run ID not provided');
    }

    config.run_id = this.id;
    config.configurable = Object.assign(config.configurable ?? {}, {
      run_id: this.id,
    });

    const threadId = config.configurable.thread_id as string | undefined;

    if (this.hookRegistry != null) {
      await executeHooks({
        registry: this.hookRegistry,
        input: {
          hook_event_name: 'RunStart',
          runId: this.id,
          threadId,
          agentId: this.Graph.defaultAgentId,
          messages: (inputs as { messages?: BaseMessage[] }).messages ?? [],
        },
        sessionId: this.id,
      });

      const messages = (inputs as { messages?: BaseMessage[] }).messages ?? [];
      const lastHuman = findLastMessageOfType(messages, 'human');
      if (lastHuman != null) {
        const promptResult = await executeHooks({
          registry: this.hookRegistry,
          input: {
            hook_event_name: 'UserPromptSubmit',
            runId: this.id,
            threadId,
            agentId: this.Graph.defaultAgentId,
            prompt: extractPromptText(lastHuman),
          },
          sessionId: this.id,
        });
        if (
          promptResult.decision === 'deny' ||
          promptResult.decision === 'ask'
        ) {
          this.hookRegistry.clearSession(this.id);
          config.callbacks = undefined;
          return undefined;
        }
      }
    }

    const stream = this.graphRunnable.streamEvents(inputs, config, {
      raiseError: true,
      /**
       * Prevent EventStreamCallbackHandler from processing custom events.
       * Custom events are already handled via our createCustomEventCallback()
       * which routes them through the handlerRegistry.
       * Without this flag, EventStreamCallbackHandler throws errors when
       * custom events are dispatched for run IDs not in its internal map
       * (due to timing issues in parallel execution or after run cleanup).
       */
      ignoreCustomEvent: true,
    });

    try {
      for await (const event of stream) {
        const { data, metadata, ...info } = event;

        const eventName: t.EventName = info.event;

        /** Skip custom events as they're handled by our callback */
        if (eventName === GraphEvents.ON_CUSTOM_EVENT) {
          continue;
        }

        const handler = this.handlerRegistry?.getHandler(eventName);
        if (handler) {
          await handler.handle(eventName, data, metadata, this.Graph);
        }
      }

      if (this.hookRegistry?.hasHookFor('Stop', this.id) === true) {
        const inputMessages =
          (inputs as { messages?: BaseMessage[] }).messages ?? [];
        await executeHooks({
          registry: this.hookRegistry,
          input: {
            hook_event_name: 'Stop',
            runId: this.id,
            threadId,
            agentId: this.Graph.defaultAgentId,
            messages: this.Graph.getRunMessages() ?? inputMessages,
            stopHookActive: false,
          },
          sessionId: this.id,
        }).catch(() => {
          /* Stop hook errors must not masquerade as stream failures */
        });
      }
    } catch (e) {
      // GraphInterrupt is expected when interrupt() fires (HITL approval flow).
      // Exit gracefully so the host can check hasInterrupts() and resume.
      if (isGraphInterrupt(e)) {
        return undefined;
      }
      if (this.hookRegistry?.hasHookFor('StopFailure', this.id) === true) {
        const runMessages = this.Graph.getRunMessages() ?? [];
        await executeHooks({
          registry: this.hookRegistry,
          input: {
            hook_event_name: 'StopFailure',
            runId: this.id,
            threadId,
            agentId: this.Graph.defaultAgentId,
            error: e instanceof Error ? e.message : String(e),
            lastAssistantMessage: findLastMessageOfType(runMessages, 'ai'),
          },
          sessionId: this.id,
        }).catch(() => {
          /* swallow hook errors — original error must propagate */
        });
      }
      throw e;
    } finally {
      this.hookRegistry?.clearSession(this.id);
    }

    /**
     * Break the reference chain that keeps heavy data alive via
     * LangGraph's internal `__pregel_scratchpad.currentTaskInput` →
     * `@langchain/core` `RunTree.extra[lc:child_config]` →
     * Node.js `AsyncLocalStorage` context captured by timers/promises.
     *
     * Without this, base64-encoded images/PDFs in message content remain
     * reachable from lingering `Timeout` handles until GC runs.
     */
    if (!this.skipCleanup) {
      if (
        (config.configurable as Record<string, unknown> | undefined) != null
      ) {
        for (const key of Object.getOwnPropertySymbols(config.configurable)) {
          const val = config.configurable[key as unknown as string] as
            | Record<string, unknown>
            | undefined;
          if (
            val != null &&
            typeof val === 'object' &&
            'currentTaskInput' in val
          ) {
            (val as Record<string, unknown>).currentTaskInput = undefined;
          }
          delete config.configurable[key as unknown as string];
        }
        config.configurable = undefined;
      }
      config.callbacks = undefined;
    }

    const result = this.returnContent
      ? this.Graph.getContentParts()
      : undefined;

    if (!this.skipCleanup) {
      this.Graph.clearHeavyState();
    }
    return result;
  }

  private createSystemCallback<K extends keyof t.ClientCallbacks>(
    clientCallbacks: t.ClientCallbacks,
    key: K
  ): t.SystemCallbacks[K] {
    return ((...args: unknown[]) => {
      const clientCallback = clientCallbacks[key];
      if (clientCallback && this.Graph) {
        (clientCallback as (...args: unknown[]) => void)(this.Graph, ...args);
      }
    }) as t.SystemCallbacks[K];
  }

  /**
   * Checks whether the graph was interrupted (e.g., by HITL tool approval).
   * Call after processStream() returns to determine if the graph is waiting
   * for a resume via Command({ resume }).
   *
   * Requires a checkpointer to be configured — without one, interrupt state
   * is not persisted and this always returns false.
   */
  async hasInterrupts(config: Partial<RunnableConfig>): Promise<boolean> {
    if (!this.graphRunnable) {
      return false;
    }
    try {
      const state = await this.graphRunnable.getState(config);
      return state.tasks.some((task) => task.interrupts.length > 0) ?? false;
    } catch {
      return false;
    }
  }

  /**
   * Returns the interrupt values from the graph state.
   * Each interrupt's `value` contains the data passed to `interrupt()` by the node
   * (e.g., a ToolApprovalRequest for HITL).
   */
  async getInterruptValues(
    config: Partial<RunnableConfig>
  ): Promise<unknown[]> {
    if (!this.graphRunnable) {
      return [];
    }
    try {
      const state = await this.graphRunnable.getState(config);
      return (
        state.tasks.flatMap((task) =>
          (task.interrupts ?? []).map((i) => i.value)
        ) ?? []
      );
    } catch {
      return [];
    }
  }

  getCallbacks(clientCallbacks: t.ClientCallbacks): t.SystemCallbacks {
    return {
      [Callback.TOOL_ERROR]: this.createSystemCallback(
        clientCallbacks,
        Callback.TOOL_ERROR
      ),
      [Callback.TOOL_START]: this.createSystemCallback(
        clientCallbacks,
        Callback.TOOL_START
      ),
      [Callback.TOOL_END]: this.createSystemCallback(
        clientCallbacks,
        Callback.TOOL_END
      ),
    };
  }

  async generateTitle({
    provider,
    inputText,
    contentParts,
    titlePrompt,
    clientOptions,
    chainOptions,
    skipLanguage,
    titleMethod = TitleMethod.COMPLETION,
    titlePromptTemplate,
  }: t.RunTitleOptions): Promise<{ language?: string; title?: string }> {
    const titleSecretKey = process.env.ILLUMA_SECRET_KEY;
    const titlePublicKey = process.env.ILLUMA_PUBLIC_KEY;
    const titleBaseUrl = process.env.ILLUMA_BASE_URL;

    if (
      chainOptions != null &&
      isPresent(titleSecretKey) &&
      isPresent(titlePublicKey) &&
      isPresent(titleBaseUrl)
    ) {
      try {
        const userId = chainOptions.configurable?.user_id;
        const sessionId = chainOptions.configurable?.thread_id;
        const handler = new ObservabilityCallbackHandler({
          clientOptions: {
            secretKey: titleSecretKey!,
            publicKey: titlePublicKey!,
            baseUrl: titleBaseUrl!,
          },
          userId,
          sessionId,
          metadata: {
            messageId: 'title-' + this.id,
          },
        });
        chainOptions.callbacks = (
          (chainOptions.callbacks as t.ProvidedCallbacks) ?? []
        ).concat([handler]);
      } catch {
        // Gracefully skip if @illuma-ai/observability-node is not installed
      }
    }

    const convoTemplate = PromptTemplate.fromTemplate(
      titlePromptTemplate ?? 'User: {input}\nAI: {output}'
    );

    const response = contentParts
      .map((part) => {
        if (part?.type === 'text') return part.text;
        return '';
      })
      .join('\n');

    const model = this.Graph?.getNewModel({
      provider,
      clientOptions,
    });
    if (!model) {
      return { language: '', title: '' };
    }
    if (
      isOpenAILike(provider) &&
      (model instanceof ChatOpenAI || model instanceof AzureChatOpenAI)
    ) {
      model.temperature = (clientOptions as t.OpenAIClientOptions | undefined)
        ?.temperature as number;
      model.topP = (clientOptions as t.OpenAIClientOptions | undefined)
        ?.topP as number;
      model.frequencyPenalty = (
        clientOptions as t.OpenAIClientOptions | undefined
      )?.frequencyPenalty as number;
      model.presencePenalty = (
        clientOptions as t.OpenAIClientOptions | undefined
      )?.presencePenalty as number;
      model.n = (clientOptions as t.OpenAIClientOptions | undefined)
        ?.n as number;
    }

    const convoToTitleInput = new RunnableLambda({
      func: (
        promptValue: StringPromptValue
      ): { convo: string; inputText: string; skipLanguage?: boolean } => ({
        convo: promptValue.value,
        inputText,
        skipLanguage,
      }),
    }).withConfig({ runName: 'ConvoTransform' });

    const titleChain =
      titleMethod === TitleMethod.COMPLETION
        ? await createCompletionTitleRunnable(model, titlePrompt)
        : await createTitleRunnable(model, titlePrompt);

    /** Pipes `convoTemplate` -> `transformer` -> `titleChain` */
    const fullChain = convoTemplate
      .withConfig({ runName: 'ConvoTemplate' })
      .pipe(convoToTitleInput)
      .pipe(titleChain)
      .withConfig({ runName: 'TitleChain' });

    const invokeConfig = Object.assign({}, chainOptions, {
      run_id: this.id,
      runId: this.id,
    });

    try {
      return await fullChain.invoke(
        { input: inputText, output: response },
        invokeConfig
      );
    } catch (_e) {
      // Fallback: strip callbacks to avoid EventStream tracer errors in certain environments
      // But preserve observability handler if it exists
      const observabilityHandler = (
        invokeConfig.callbacks as t.ProvidedCallbacks
      )?.find((cb) => cb instanceof ObservabilityCallbackHandler);
      const { callbacks: _cb, ...rest } = invokeConfig;
      const safeConfig = Object.assign({}, rest, {
        callbacks: observabilityHandler ? [observabilityHandler] : [],
      });
      return await fullChain.invoke(
        { input: inputText, output: response },
        safeConfig as Partial<RunnableConfig>
      );
    }
  }
}