// SPDX-FileCopyrightText: 2025 LiveKit, Inc.
//
// SPDX-License-Identifier: Apache-2.0
import type { AudioFrame } from '@livekit/rtc-node';
import { AsyncLocalStorage } from 'node:async_hooks';
import { ReadableStream } from 'node:stream/web';
import {
  LLM as InferenceLLM,
  STT as InferenceSTT,
  TTS as InferenceTTS,
  type LLMModels,
  type STTModelString,
  type TTSModelString,
} from '../inference/index.js';
import { type Instructions, ReadonlyChatContext } from '../llm/chat_context.js';
import type { ChatMessage, FunctionCall } from '../llm/index.js';
import {
  type ChatChunk,
  ChatContext,
  LLM,
  RealtimeModel,
  type ToolChoice,
  ToolContext,
  type ToolContextLike,
  toToolContext,
} from '../llm/index.js';
import { log } from '../log.js';
import type { STT, SpeechEvent } from '../stt/index.js';
import { StreamAdapter as STTStreamAdapter } from '../stt/index.js';
import { SentenceTokenizer as BasicSentenceTokenizer } from '../tokenize/basic/index.js';
import type { TTS } from '../tts/index.js';
import { SynthesizeStream, StreamAdapter as TTSStreamAdapter } from '../tts/index.js';
import { type FlushSentinel, USERDATA_TIMED_TRANSCRIPT } from '../types.js';
import { Future, Task, toStream } from '../utils.js';
import type { VAD } from '../vad.js';
import { type AgentActivity, agentActivityStorage } from './agent_activity.js';
import type { AgentSession, TurnDetectionMode } from './agent_session.js';
import {
  type AgentCreateOptions,
  type AgentTaskCreateOptions,
  createAgentTaskV2,
  createAgentV2,
} from './agent_v2.js';
import type { UserTurnExceededEvent } from './events.js';
import type { TimedString } from './io.js';
import type { SpeechHandle } from './speech_handle.js';
import type { ToolHandlingOptions } from './tool_executor.js';
import type { TurnHandlingOptions } from './turn_config/turn_handling.js';
import { migrateTurnHandling } from './turn_config/utils.js';

export type {
  AgentContext,
  AgentCreateOptions,
  AgentHookNodeResult,
  AgentHooks,
  AgentTaskContext,
  AgentTaskCreateOptions,
} from './agent_v2.js';

// speechHandle identifies which SpeechHandle owns the current tool call, enabling
// SpeechHandle.waitForPlayout() to distinguish self-wait (deadlock) from waiting
// on a different handle scheduled inside the tool.
export const functionCallStorage = new AsyncLocalStorage<{
  functionCall?: FunctionCall;
  speechHandle?: SpeechHandle;
}>();
export const speechHandleStorage = new AsyncLocalStorage<SpeechHandle>();
const activityTaskInfoStorage = new WeakMap<Task<any>, _ActivityTaskInfo>();

type _ActivityTaskInfo = {
  functionCall: FunctionCall | null;
  speechHandle: SpeechHandle | null;
  inlineTask: boolean;
};

/** @internal */
export function _setActivityTaskInfo<T>(
  task: Task<T>,
  options: {
    functionCall?: FunctionCall | null;
    speechHandle?: SpeechHandle | null;
    inlineTask?: boolean;
  },
): void {
  const info = activityTaskInfoStorage.get(task) ?? {
    functionCall: null,
    speechHandle: null,
    inlineTask: false,
  };

  if (Object.hasOwn(options, 'functionCall')) {
    info.functionCall = options.functionCall ?? null;
  }
  if (Object.hasOwn(options, 'speechHandle')) {
    info.speechHandle = options.speechHandle ?? null;
  }
  if (Object.hasOwn(options, 'inlineTask')) {
    info.inlineTask = options.inlineTask ?? false;
  }

  activityTaskInfoStorage.set(task, info);
}

/** @internal */
export function _getActivityTaskInfo<T>(task: Task<T>): _ActivityTaskInfo | undefined {
  return activityTaskInfoStorage.get(task);
}
export const STOP_RESPONSE_SYMBOL = Symbol('StopResponse');

export class StopResponse extends Error {
  constructor() {
    super();
    this.name = 'StopResponse';

    Object.defineProperty(this, STOP_RESPONSE_SYMBOL, {
      value: true,
    });
  }
}

export function isStopResponse(value: unknown): value is StopResponse {
  return (
    value !== undefined &&
    value !== null &&
    typeof value === 'object' &&
    STOP_RESPONSE_SYMBOL in value
  );
}

export interface ModelSettings {
  /** The tool choice to use when calling the LLM. */
  toolChoice?: ToolChoice;
}

export interface AgentOptions<UserData> {
  id?: string;
  instructions: string | Instructions;
  chatCtx?: ChatContext;
  tools?: ToolContextLike<UserData>;
  stt?: STT | STTModelString;
  vad?: VAD;
  llm?: LLM | RealtimeModel | LLMModels;
  tts?: TTS | TTSModelString;
  turnHandling?: TurnHandlingOptions;
  toolHandling?: ToolHandlingOptions;
  minConsecutiveSpeechDelay?: number;
  useTtsAlignedTranscript?: boolean;
  /** @deprecated use turnHandling.turnDetection instead */
  turnDetection?: TurnDetectionMode;
  /** @deprecated use turnHandling.interruption.enabled instead */
  allowInterruptions?: boolean;
}

export class Agent<UserData = any> {
  private _id: string;
  private _stt?: STT;
  private _vad?: VAD;
  private _llm?: LLM | RealtimeModel;
  private _tts?: TTS;
  private _turnHandling?: Partial<TurnHandlingOptions>;

  private _minConsecutiveSpeechDelay?: number;
  private _useTtsAlignedTranscript?: boolean;

  /** @internal */
  _agentActivity?: AgentActivity;

  /** @internal */
  _chatCtx: ChatContext;

  /** @internal */
  _instructions: string | Instructions;

  /** @internal */
  _toolCtx: ToolContext<UserData>;

  /** @internal */
  _asyncToolOptions?: ToolHandlingOptions['asyncOptions'];

  static create<UserData = unknown>(options: AgentCreateOptions<UserData>): Agent<UserData> {
    return createAgentV2(Agent, options);
  }

  constructor({
    id,
    instructions,
    chatCtx,
    tools,
    turnDetection,
    stt,
    vad,
    llm,
    tts,
    allowInterruptions,
    turnHandling,
    toolHandling,
    minConsecutiveSpeechDelay,
    useTtsAlignedTranscript,
  }: AgentOptions<UserData>) {
    if (id) {
      this._id = id;
    } else {
      const className = this.constructor.name;
      if (className === 'Agent') {
        this._id = 'default_agent';
      } else {
        this._id = className
          .replace(/([A-Z])/g, '_$1')
          .toLowerCase()
          .replace(/^_/, '');
      }
    }

    this._instructions = instructions;
    this._toolCtx = toToolContext(tools) ?? ToolContext.empty<UserData>();
    this._chatCtx = chatCtx
      ? chatCtx.copy({
          toolCtx: this._toolCtx,
        })
      : ChatContext.empty();

    const resolvedTurnHandling = migrateTurnHandling({
      turnDetection,
      allowInterruptions,
      turnHandling,
    });
    this._turnHandling =
      Object.keys(resolvedTurnHandling).length > 0 ? resolvedTurnHandling : undefined;
    this._asyncToolOptions = toolHandling?.asyncOptions;

    this._vad = vad;

    if (typeof stt === 'string') {
      this._stt = InferenceSTT.fromModelString(stt);
    } else {
      this._stt = stt;
    }

    if (typeof llm === 'string') {
      this._llm = InferenceLLM.fromModelString(llm);
    } else {
      this._llm = llm;
    }

    if (typeof tts === 'string') {
      this._tts = InferenceTTS.fromModelString(tts);
    } else {
      this._tts = tts;
    }

    this._minConsecutiveSpeechDelay = minConsecutiveSpeechDelay;
    this._useTtsAlignedTranscript = useTtsAlignedTranscript;

    this._agentActivity = undefined;
  }

  get vad(): VAD | undefined {
    return this._vad;
  }

  get stt(): STT | undefined {
    return this._stt;
  }

  get llm(): LLM | RealtimeModel | undefined {
    return this._llm;
  }

  get tts(): TTS | undefined {
    return this._tts;
  }

  get useTtsAlignedTranscript(): boolean | undefined {
    return this._useTtsAlignedTranscript;
  }

  get chatCtx(): ReadonlyChatContext {
    return new ReadonlyChatContext(this._chatCtx.items);
  }

  get id(): string {
    return this._id;
  }

  get instructions(): string | Instructions {
    return this._instructions;
  }

  get toolCtx(): ToolContext<UserData> {
    return this._toolCtx.copy();
  }

  get session(): AgentSession<UserData> {
    return this.getActivityOrThrow().agentSession as AgentSession<UserData>;
  }

  get turnHandling(): Partial<TurnHandlingOptions> | undefined {
    return this._turnHandling;
  }

  get minConsecutiveSpeechDelay(): number | undefined {
    return this._minConsecutiveSpeechDelay;
  }

  async onEnter(): Promise<void> {}

  async onExit(): Promise<void> {}

  async transcriptionNode(
    text: ReadableStream<string | TimedString> | AsyncIterable<string | TimedString>,
    modelSettings: ModelSettings,
  ): Promise<ReadableStream<string | TimedString> | null> {
    return Agent.default.transcriptionNode(this, text, modelSettings);
  }

  async onUserTurnCompleted(_chatCtx: ChatContext, _newMessage: ChatMessage): Promise<void> {}

  async onUserTurnExceeded(ev: UserTurnExceededEvent): Promise<void> {
    await this.session.generateReply({
      userInput: ev.transcript,
      instructions:
        'The user has been speaking too long without giving a chance to reply. ' +
        'Politely cut in with a short reply or notice. Keep it short since the user cannot interrupt it.',
      allowInterruptions: false,
      toolChoice: 'none',
      inputModality: 'audio',
    });
  }

  async sttNode(
    audio: ReadableStream<AudioFrame> | AsyncIterable<AudioFrame>,
    modelSettings: ModelSettings,
  ): Promise<ReadableStream<SpeechEvent | string> | null> {
    return Agent.default.sttNode(this, audio, modelSettings);
  }

  async llmNode(
    chatCtx: ChatContext,
    toolCtx: ToolContext,
    modelSettings: ModelSettings,
  ): Promise<ReadableStream<ChatChunk | string | FlushSentinel> | null> {
    return Agent.default.llmNode(this, chatCtx, toolCtx, modelSettings);
  }

  async ttsNode(
    text: ReadableStream<string> | AsyncIterable<string>,
    modelSettings: ModelSettings,
  ): Promise<ReadableStream<AudioFrame> | null> {
    return Agent.default.ttsNode(this, text, modelSettings);
  }

  async realtimeAudioOutputNode(
    audio: ReadableStream<AudioFrame> | AsyncIterable<AudioFrame>,
    modelSettings: ModelSettings,
  ): Promise<ReadableStream<AudioFrame> | null> {
    return Agent.default.realtimeAudioOutputNode(this, audio, modelSettings);
  }

  // realtime_audio_output_node

  getActivityOrThrow(): AgentActivity {
    if (!this._agentActivity) {
      throw new Error('Agent activity not found');
    }
    return this._agentActivity;
  }

  async updateChatCtx(chatCtx: ChatContext): Promise<void> {
    if (!this._agentActivity) {
      this._chatCtx = chatCtx.copy({ toolCtx: this.toolCtx });
      return;
    }

    this._agentActivity.updateChatCtx(chatCtx);
  }

  async updateInstructions(instructions: string | Instructions): Promise<void> {
    if (!this._agentActivity) {
      this._instructions = instructions;
      return;
    }

    await this._agentActivity.updateInstructions(instructions);
  }

  // TODO(parity): Add when AgentConfigUpdate is ported to ChatContext.
  async updateTools(tools: ToolContextLike<UserData>): Promise<void> {
    if (!this._agentActivity) {
      this._toolCtx = toToolContext(tools);
      this._chatCtx = this._chatCtx.copy({ toolCtx: this._toolCtx });
      return;
    }

    await this._agentActivity.updateTools(tools);
  }

  static default = {
    async sttNode(
      agent: Agent,
      audio: ReadableStream<AudioFrame> | AsyncIterable<AudioFrame>,
      _modelSettings: ModelSettings,
    ): Promise<ReadableStream<SpeechEvent | string> | null> {
      const input = audio instanceof ReadableStream ? audio : toStream(audio);
      const activity = agent.getActivityOrThrow();
      if (!activity.stt) {
        throw new Error('sttNode called but no STT node is available');
      }

      let wrappedStt = activity.stt;

      if (!wrappedStt.capabilities.streaming) {
        const vad = agent.vad || activity.vad;
        if (!vad) {
          throw new Error(
            'STT does not support streaming, add a VAD to the AgentTask/VoiceAgent to enable streaming',
          );
        }
        wrappedStt = new STTStreamAdapter(wrappedStt, vad);
      }

      const connOptions = activity.agentSession.connOptions.sttConnOptions;
      const stream = wrappedStt.stream({ connOptions });

      // Set startTimeOffset to provide linear timestamps across reconnections
      const audioInputStartedAt =
        activity.inputStartedAt ?? // Use input started at proxied from AudioRecognition if available
        activity.agentSession._recorderIO?.recordingStartedAt ?? // Fallback to recording start time if available
        activity.agentSession._startedAt ?? // Fallback to session start time
        Date.now(); // Fallback to current time

      stream.startTimeOffset = (Date.now() - audioInputStartedAt) / 1000;

      stream.updateInputStream(input);

      let cleaned = false;
      const cleanup = () => {
        if (cleaned) return;
        cleaned = true;
        stream.detachInputStream();
        stream.close();
      };

      return new ReadableStream({
        async start(controller) {
          try {
            for await (const event of stream) {
              controller.enqueue(event);
            }
            controller.close();
          } finally {
            // Always clean up the STT stream, whether it ends naturally or is cancelled
            cleanup();
          }
        },
        cancel() {
          cleanup();
        },
      });
    },

    async llmNode(
      agent: Agent,
      chatCtx: ChatContext,
      toolCtx: ToolContext,
      modelSettings: ModelSettings,
    ): Promise<ReadableStream<ChatChunk | string> | null> {
      const activity = agent.getActivityOrThrow();
      if (!activity.llm) {
        throw new Error('llmNode called but no LLM node is available');
      }

      if (!(activity.llm instanceof LLM)) {
        throw new Error(
          'llmNode should only be used with LLM (non-multimodal/realtime APIs) nodes',
        );
      }

      const { toolChoice } = modelSettings;
      const connOptions = activity.agentSession.connOptions.llmConnOptions;

      // parallelToolCalls is not passed here - it will use the value from LLM's modelOptions
      // This allows users to configure it via: new inference.LLM({ modelOptions: { parallel_tool_calls: false } })
      const stream = activity.llm.chat({
        chatCtx,
        toolCtx,
        toolChoice,
        connOptions,
      });

      let cleaned = false;
      const cleanup = () => {
        if (cleaned) return;
        cleaned = true;
        stream.close();
      };

      return new ReadableStream({
        async start(controller) {
          try {
            for await (const chunk of stream) {
              controller.enqueue(chunk);
            }
            controller.close();
          } finally {
            cleanup();
          }
        },
        cancel() {
          cleanup();
        },
      });
    },

    async ttsNode(
      agent: Agent,
      text: ReadableStream<string> | AsyncIterable<string>,
      _modelSettings: ModelSettings,
    ): Promise<ReadableStream<AudioFrame> | null> {
      const input = text instanceof ReadableStream ? text : toStream(text);
      const activity = agent.getActivityOrThrow();
      if (!activity.tts) {
        throw new Error('ttsNode called but no TTS node is available');
      }

      let wrappedTts = activity.tts;

      if (!activity.tts.capabilities.streaming) {
        wrappedTts = new TTSStreamAdapter(wrappedTts, new BasicSentenceTokenizer());
      }

      const connOptions = activity.agentSession.connOptions.ttsConnOptions;
      const stream = wrappedTts.stream({ connOptions });
      stream.updateInputStream(input);

      let cleaned = false;
      const cleanup = async () => {
        if (cleaned) return;
        cleaned = true;
        stream.close();
        await input.cancel('tts node cleanup').catch(() => {});
        if (wrappedTts !== activity.tts) {
          await wrappedTts.close();
        }
      };

      return new ReadableStream({
        async start(controller) {
          try {
            for await (const chunk of stream) {
              if (chunk === SynthesizeStream.END_OF_STREAM) {
                break;
              }
              // Attach timed transcripts to frame.userdata
              if (chunk.timedTranscripts && chunk.timedTranscripts.length > 0) {
                chunk.frame.userdata[USERDATA_TIMED_TRANSCRIPT] = chunk.timedTranscripts;
              }
              controller.enqueue(chunk.frame);
            }
            controller.close();
          } finally {
            await cleanup();
          }
        },
        cancel() {
          return cleanup();
        },
      });
    },

    async transcriptionNode(
      agent: Agent,
      text: ReadableStream<string | TimedString> | AsyncIterable<string | TimedString>,
      _modelSettings: ModelSettings,
    ): Promise<ReadableStream<string | TimedString> | null> {
      return text instanceof ReadableStream ? text : toStream(text);
    },

    async realtimeAudioOutputNode(
      _agent: Agent,
      audio: ReadableStream<AudioFrame> | AsyncIterable<AudioFrame>,
      _modelSettings: ModelSettings,
    ): Promise<ReadableStream<AudioFrame> | null> {
      return audio instanceof ReadableStream ? audio : toStream(audio);
    },
  };
}

export interface AgentTaskOptions<UserData = any> extends AgentOptions<UserData> {
  preserveFunctionCallHistory?: boolean;
}

export class AgentTask<ResultT = unknown, UserData = any> extends Agent<UserData> {
  private started = false;
  private future = new Future<ResultT>();
  private _preserveFunctionCallHistory: boolean;

  #logger = log();

  static create<ResultT = unknown, UserData = unknown>(
    options: AgentTaskCreateOptions<ResultT, UserData>,
  ): AgentTask<ResultT, UserData> {
    return createAgentTaskV2<ResultT, UserData>(AgentTask, options);
  }

  constructor(options: AgentTaskOptions<UserData>) {
    const { preserveFunctionCallHistory = false, ...rest } = options;
    super(rest);
    this._preserveFunctionCallHistory = preserveFunctionCallHistory;
  }

  get done(): boolean {
    return this.future.done;
  }

  complete(result: ResultT | Error): void {
    if (this.future.done) {
      throw new Error(`${this.constructor.name} is already done`);
    }

    if (result instanceof Error) {
      this.future.reject(result);
    } else {
      this.future.resolve(result);
    }

    const speechHandle = speechHandleStorage.getStore();
    if (speechHandle) {
      speechHandle._maybeRunFinalOutput = result;
    }
  }

  async run(): Promise<ResultT> {
    if (this.started) {
      throw new Error(
        `Task ${this.constructor.name} has already started and cannot be awaited multiple times`,
      );
    }
    this.started = true;

    const currentTask = Task.current();
    if (!currentTask) {
      throw new Error(`${this.constructor.name} must be executed inside a Task context`);
    }

    const taskInfo = _getActivityTaskInfo(currentTask);
    if (!taskInfo || !taskInfo.inlineTask) {
      throw new Error(
        `${this.constructor.name} should only be awaited inside function tools or the onEnter/onExit methods of an Agent`,
      );
    }

    const speechHandle = speechHandleStorage.getStore();
    const oldActivity = agentActivityStorage.getStore();
    if (!oldActivity) {
      throw new Error(`${this.constructor.name} must be executed inside an AgentActivity context`);
    }

    // A non-blocking tool (one that called ctx.update) detaches from its speech
    // task, so that task completes before a later ctx.foreground/AgentTask does.
    // Binding the guard below to the already-finished task would fire it
    // immediately and tear the AgentTask down; the tool's still-running promise
    // and the foreground hold keep it alive instead.
    const ownerIsNonBlocking =
      taskInfo.functionCall?.extra.__livekit_agents_tool_non_blocking === true;
    if (!ownerIsNonBlocking) {
      currentTask.addDoneCallback(() => {
        if (this.future.done) return;

        // If the Task finished before the AgentTask was completed, complete the AgentTask with an error.
        this.#logger.error(`The Task finished before ${this.constructor.name} was completed.`);
        this.complete(
          new Error(`The Task finished before ${this.constructor.name} was completed.`),
        );
      });
    }

    const oldAgent = oldActivity.agent;
    const session = oldActivity.agentSession;

    const blockedTasks: Task<any>[] = [currentTask];
    const onEnterTask = oldActivity._onEnterTask;

    if (onEnterTask && !onEnterTask.done && onEnterTask !== currentTask) {
      blockedTasks.push(onEnterTask);
    }

    const activeRunState = session._globalRunState;
    if (activeRunState && !activeRunState.done()) {
      for (const task of blockedTasks) {
        activeRunState._watchHandle(task);
      }
    }

    if (
      taskInfo.functionCall &&
      oldActivity.llm instanceof RealtimeModel &&
      !oldActivity.llm.capabilities.manualFunctionCalls
    ) {
      this.#logger.error(
        `Realtime model does not support resuming function calls from chat context, ` +
          `using AgentTask inside a function tool may have unexpected behavior.`,
      );
    }

    const suspendedHandles: Array<SpeechHandle | Task<void>> = [];

    await session._updateActivity(this, {
      previousActivity: 'pause',
      newActivity: 'start',
      blockedTasks,
    });

    let runState = session._globalRunState;
    if (runState && !runState.done()) {
      if (speechHandle && runState._unwatchHandle(speechHandle)) {
        suspendedHandles.push(speechHandle);
      }

      for (const task of blockedTasks) {
        if (runState._unwatchHandle(task)) {
          suspendedHandles.push(task);
        }
      }

      // It is OK to call _markDoneIfNeeded here: _updateActivity has started
      // the AgentTask activity, and any onEnter-generated speech is now watched.
      // Only call it when something was actually suspended — a run created
      // mid-transition has no watched handles yet (its generateReply is still
      // deferred behind the activity lock), and marking done on an empty
      // handle set would resolve it before it produced any events.
      if (suspendedHandles.length > 0) {
        runState._markDoneIfNeeded();
      }
    }

    try {
      return await this.future.await;
    } finally {
      // runState could have changed after future resolved
      runState = session._globalRunState;

      if (session.currentAgent !== this) {
        this.#logger.warn(
          `${this.constructor.name} completed, but the agent has changed in the meantime. ` +
            `Ignoring handoff to the previous agent, likely due to AgentSession.updateAgent being invoked.`,
        );
        await oldActivity.close();
      } else {
        if (runState && !runState.done()) {
          for (const handle of suspendedHandles) {
            runState._watchHandle(handle);
          }
        }

        const mergedChatCtx = oldAgent._chatCtx.merge(this._chatCtx, {
          excludeFunctionCall: !this._preserveFunctionCallHistory,
          excludeInstructions: true,
        });
        oldAgent._chatCtx.items = mergedChatCtx.items;

        await session._updateActivity(oldAgent, {
          previousActivity: 'close',
          newActivity: 'resume',
          waitOnEnter: false,
        });
      }
    }
  }
}
