// SPDX-FileCopyrightText: 2024 LiveKit, Inc.
//
// SPDX-License-Identifier: Apache-2.0
import type * as proto from '@livekit/protocol';
import type {
  E2EEOptions,
  LocalParticipant,
  RemoteParticipant,
  RemoteTrackPublication,
  Room,
  RtcConfiguration,
} from '@livekit/rtc-node';
import { ParticipantKind, RoomEvent, TrackKind } from '@livekit/rtc-node';
import { ThrowsPromise } from '@livekit/throws-transformer/throws';
import { RoomServiceClient } from 'livekit-server-sdk';
import { AsyncLocalStorage } from 'node:async_hooks';
import { mkdir, writeFile } from 'node:fs/promises';
import * as os from 'node:os';
import * as path from 'node:path';
import type { Logger } from 'pino';
import type { InferenceExecutor } from './ipc/inference_executor.js';
import { log } from './log.js';
import { flushOtelLogs, setupCloudTracer, uploadSessionReport } from './telemetry/index.js';
import { isCloud } from './utils.js';
import type { AgentSession, ResolvedRecordingOptions } from './voice/agent_session.js';
import { AgentsConsole } from './voice/console_io.js';
import { type SessionReport, createSessionReport, sessionReportToJSON } from './voice/report.js';

// AsyncLocalStorage for job context, similar to Python's contextvars
const jobContextStorage = new AsyncLocalStorage<JobContext<unknown>>();

/**
 * Returns the current job context.
 *
 * @param required - If true (default), throws when no context is found. If false, returns undefined.
 * @throws Error if no job context is found and required is true
 */
export function getJobContext<ProcessUserData = Record<string, unknown>>(
  required?: true,
): JobContext<ProcessUserData>;
export function getJobContext<ProcessUserData = Record<string, unknown>>(
  required: false,
): JobContext<ProcessUserData> | undefined;
export function getJobContext<ProcessUserData = Record<string, unknown>>(
  required = true,
): JobContext<ProcessUserData> | undefined {
  const ctx = jobContextStorage.getStore();
  if (!ctx && required) {
    throw new Error('no job context found, are you running this code inside a job entrypoint?');
  }
  return ctx as JobContext<ProcessUserData> | undefined;
}

/**
 * Runs a function within a job context, similar to Python's contextvars.
 * @internal
 */
export function runWithJobContext<T>(context: JobContext, fn: () => T): T {
  return jobContextStorage.run(context as JobContext<unknown>, fn);
}

/**
 * Runs an async function within a job context, similar to Python's contextvars.
 * @internal
 */
export function runWithJobContextAsync<T>(context: JobContext, fn: () => Promise<T>): Promise<T> {
  return jobContextStorage.run(context as JobContext<unknown>, fn);
}

/** Which tracks, if any, should the agent automatically subscribe to? */
export enum AutoSubscribe {
  SUBSCRIBE_ALL,
  SUBSCRIBE_NONE,
  VIDEO_ONLY,
  AUDIO_ONLY,
}

export type JobAcceptArguments = {
  name: string;
  identity: string;
  metadata: string;
  attributes?: { [key: string]: string };
};

export type RunningJobInfo = {
  acceptArguments: JobAcceptArguments;
  job: proto.Job;
  url: string;
  token: string;
  workerId: string;
  apiKey?: string;
  apiSecret?: string;
  /**
   * A locally-synthesized job that is not backed by a LiveKit room (e.g. the
   * `console` CLI runner). Room-bound operations (`connect`, `deleteRoom`,
   * recording uploads) become no-ops. Mirrors python `RunningJobInfo.fake_job`.
   */
  fakeJob?: boolean;
};

/** Attempted to add a function callback, but the function already exists. */
export class FunctionExistsError extends Error {
  constructor(msg?: string) {
    super(msg);
    Object.setPrototypeOf(this, new.target.prototype);
  }
}

/** The job and environment context as seen by the agent, accessible by the entrypoint function. */
export class JobContext<ProcessUserData = Record<string, unknown>> {
  #proc: JobProcess<ProcessUserData>;
  #info: RunningJobInfo;
  #room: Room;
  #onConnect: () => void;
  #onShutdown: (s: string) => void;
  /** @internal */
  shutdownCallbacks: (() => Promise<void>)[] = [];
  #participantEntrypoints: ((
    job: JobContext<ProcessUserData>,
    p: RemoteParticipant,
  ) => Promise<void>)[] = [];
  #participantTasks: {
    [id: string]: {
      callback: (job: JobContext<ProcessUserData>, p: RemoteParticipant) => Promise<void>;
      result: Promise<void>;
    };
  } = {};
  #logger: Logger;
  #inferenceExecutor: InferenceExecutor;

  /** @internal */
  _primaryAgentSession?: AgentSession;

  /** @internal */
  _sessionDirectory: string;

  private connected: boolean = false;

  constructor(
    proc: JobProcess<ProcessUserData>,
    info: RunningJobInfo,
    room: Room,
    onConnect: () => void,
    onShutdown: (s: string) => void,
    inferenceExecutor: InferenceExecutor,
  ) {
    this.#proc = proc;
    this.#info = info;
    this.#room = room;
    this.#onConnect = onConnect;
    this.#onShutdown = onShutdown;
    this.onParticipantConnected = this.onParticipantConnected.bind(this);
    this.#room.on(RoomEvent.ParticipantConnected, this.onParticipantConnected);
    this.#logger = log().child({
      jobId: this.#info.job.id,
      roomName: this.#info.job.room?.name,
    });
    this.#inferenceExecutor = inferenceExecutor;
    // In console mode, recordings land in a local user-visible directory
    // (mirrors python's AgentsConsole); real jobs use a temp dir.
    const agentsConsole = AgentsConsole.getInstance();
    this._sessionDirectory = agentsConsole.enabled
      ? agentsConsole.sessionDirectory
      : path.join(os.tmpdir(), 'livekit-agents', `job-${this.#info.job.id}`);
  }

  get proc(): JobProcess<ProcessUserData> {
    return this.#proc;
  }

  get job(): proto.Job {
    return this.#info.job;
  }

  get workerId(): string {
    return this.#info.workerId;
  }

  /** @returns The room the agent was called into */
  get room(): Room {
    return this.#room;
  }

  get info(): RunningJobInfo {
    return this.#info;
  }

  /** @returns Whether this job is a locally-synthesized fake job (e.g. console mode). */
  get isFakeJob(): boolean {
    return this.#info.fakeJob ?? false;
  }

  /** @returns The agent's participant if connected to the room, otherwise `undefined` */
  get agent(): LocalParticipant | undefined {
    return this.#room.localParticipant;
  }

  /** @returns The global inference executor */
  get inferenceExecutor(): InferenceExecutor {
    return this.#inferenceExecutor;
  }

  /**
   * @returns The session directory for storing recordings and session data.
   */
  get sessionDirectory(): string {
    return this._sessionDirectory;
  }

  /** Adds a promise to be awaited when {@link JobContext.shutdown | shutdown} is called. */
  addShutdownCallback(callback: () => Promise<void>) {
    this.shutdownCallbacks.push(callback);
  }

  async waitForParticipant(identity?: string): Promise<RemoteParticipant> {
    if (!this.#room.isConnected) {
      throw new Error('room is not connected');
    }

    for (const p of this.#room.remoteParticipants.values()) {
      if ((!identity || p.identity === identity) && p.info.kind != ParticipantKind.AGENT) {
        return p;
      }
    }

    return new ThrowsPromise<RemoteParticipant, Error>((resolve, reject) => {
      const onParticipantConnected = (participant: RemoteParticipant) => {
        if (
          (!identity || participant.identity === identity) &&
          participant.info.kind != ParticipantKind.AGENT
        ) {
          clearHandlers();
          resolve(participant);
        }
      };
      const onDisconnected = () => {
        clearHandlers();
        reject(new Error('Room disconnected while waiting for participant'));
      };

      const clearHandlers = () => {
        this.#room.off(RoomEvent.ParticipantConnected, onParticipantConnected);
        this.#room.off(RoomEvent.Disconnected, onDisconnected);
      };

      this.#room.on(RoomEvent.ParticipantConnected, onParticipantConnected);
      this.#room.on(RoomEvent.Disconnected, onDisconnected);
    });
  }

  /**
   * Connects the agent to the room.
   *
   * @remarks
   * It is recommended to run this command as early in the function as possible, as executing it
   * later may cause noticeable delay between user and agent joins.
   *
   * @see {@link https://github.com/livekit/node-sdks/tree/main/packages/livekit-rtc#readme |
   * @livekit/rtc-node} for more information about the parameters.
   */
  async connect(
    e2ee?: E2EEOptions,
    autoSubscribe: AutoSubscribe = AutoSubscribe.SUBSCRIBE_ALL,
    rtcConfig?: RtcConfiguration,
  ) {
    if (this.connected) {
      return;
    }

    if (this.isFakeJob) {
      this.#logger.debug('job_ctx.connect() is a no-op in console mode');
      this.#onConnect();
      this.connected = true;
      return;
    }

    const opts = {
      e2ee,
      autoSubscribe: autoSubscribe == AutoSubscribe.SUBSCRIBE_ALL,
      rtcConfig,
      dynacast: false,
    };

    await this.#room.connect(this.#info.url, this.#info.token, opts);
    this.#onConnect();

    this.#room.remoteParticipants.forEach(this.onParticipantConnected);

    if ([AutoSubscribe.AUDIO_ONLY, AutoSubscribe.VIDEO_ONLY].includes(autoSubscribe)) {
      const desiredKind =
        autoSubscribe === AutoSubscribe.AUDIO_ONLY ? TrackKind.KIND_AUDIO : TrackKind.KIND_VIDEO;
      const subscribeIfMatch = (pub: RemoteTrackPublication) => {
        if (pub.kind === desiredKind) {
          pub.setSubscribed(true);
        }
      };

      // Subscribe to tracks already published at connect time.
      this.#room.remoteParticipants.forEach((p) => {
        p.trackPublications.forEach(subscribeIfMatch);
      });

      // Server-side autosubscribe is disabled for AUDIO_ONLY/VIDEO_ONLY, so
      // subscribe to matching tracks published AFTER connect as well (e.g. a
      // participant that unmutes or joins later). Without this, late-published
      // tracks are never subscribed. Mirrors Python's _apply_auto_subscribe_opts.
      this.#room.on(RoomEvent.TrackPublished, subscribeIfMatch);
    }
    this.connected = true;
  }

  /** Deletes the room and disconnects all participants. */
  async deleteRoom(roomName?: string): Promise<void> {
    if (this.isFakeJob) {
      this.#logger.warn('job_ctx.deleteRoom() is not executed in console mode');
      return;
    }

    const targetRoomName = roomName ?? this.#room.name;
    if (!targetRoomName) {
      this.#logger.warn('cannot delete room because room name is missing');
      return;
    }

    try {
      const client = new RoomServiceClient(this.#info.url, this.#info.apiKey, this.#info.apiSecret);
      await client.deleteRoom(targetRoomName);
      this.#logger.info({ roomName: targetRoomName }, 'room deleted');
    } catch (error) {
      this.#logger.warn({ error, roomName: targetRoomName }, 'error while deleting room');
    }
  }

  makeSessionReport(session?: AgentSession): SessionReport {
    const targetSession = session || this._primaryAgentSession;

    if (!targetSession) {
      throw new Error('Cannot prepare report, no AgentSession was found');
    }

    const recorderIO = targetSession._recorderIO;

    if (recorderIO && recorderIO.recording) {
      throw new Error('Cannot create the AgentSession report, the RecorderIO is still recording');
    }

    return createSessionReport({
      jobId: this.job.id,
      roomId: this.job.room?.sid || '',
      room: this.job.room?.name || '',
      options: targetSession.sessionOptions,
      events: targetSession._recordedEvents,
      enableRecording: targetSession._enableRecording,
      recordingOptions: targetSession._recordingOptions,
      chatHistory: targetSession.history.copy(),
      startedAt: targetSession._startedAt,
      audioRecordingPath: recorderIO?.outputPath,
      audioRecordingStartedAt: recorderIO?.recordingStartedAt,
      modelUsage: targetSession._usageCollector.flatten(),
    });
  }

  async _onSessionEnd(): Promise<void> {
    const session = this._primaryAgentSession;
    if (!session) {
      return;
    }

    const recorderIO = session._recorderIO;
    if (recorderIO?.recording) {
      this.#logger.warn('recorder_io is still recording at session end, closing it');
      await recorderIO.close();
    }

    const report = this.makeSessionReport(session);

    // Console recording: dump the session report to a local file (mirrors python).
    const agentsConsole = AgentsConsole.getInstance();
    if (agentsConsole.enabled && agentsConsole.record) {
      try {
        await mkdir(this._sessionDirectory, { recursive: true });
        await writeFile(
          path.join(this._sessionDirectory, 'session_report.json'),
          JSON.stringify(sessionReportToJSON(report), null, 2),
        );
      } catch (error) {
        this.#logger.error({ error }, 'failed to save the session report');
      }
    }

    // Upload session report to LiveKit Cloud if enabled. A fake job (console
    // mode) has no backing cloud URL, so skip the upload entirely.
    if (!this.isFakeJob) {
      const url = new URL(this.#info.url);

      if (report.enableRecording && isCloud(url)) {
        try {
          await uploadSessionReport({
            agentName: this.job.agentName,
            cloudHostname: url.hostname,
            report,
          });
          this.#logger.info(
            {
              jobId: report.jobId,
              roomId: report.roomId,
            },
            'Session report uploaded to LiveKit Cloud',
          );
        } catch (error) {
          this.#logger.error({ error }, 'Failed to upload session report');
        }
      }
    }

    this.#logger.debug(
      {
        jobId: report.jobId,
        roomId: report.roomId,
        eventsCount: report.events.length,
      },
      'Session ended, report generated',
    );

    // Explicitly clear the recorded events to avoid leaking memory
    session._recordedEvents = [];

    try {
      await flushOtelLogs();
    } catch (error) {
      this.#logger.error({ error }, 'Failed to flush OTEL logs');
    }
  }

  /**
   * Gracefully shuts down the job, and runs all shutdown promises.
   *
   * @param reason - Optional reason for shutdown
   */
  shutdown(reason = '') {
    this.#onShutdown(reason);
  }

  /** @internal */
  onParticipantConnected(p: RemoteParticipant) {
    for (const callback of this.#participantEntrypoints) {
      if (this.#participantTasks[p.identity!]?.callback == callback) {
        this.#logger.warn(
          'a participant has joined before a prior prticipant task matching the same identity has finished:',
          p.identity,
        );
      }
      const result = callback(this, p);
      result.finally(() => {
        if (this.#participantTasks[p.identity!]?.result === result) {
          delete this.#participantTasks[p.identity!];
        }
      });
      this.#participantTasks[p.identity!] = { callback, result };
    }
  }

  /**
   * Adds a promise to be awaited whenever a new participant joins the room.
   *
   * @throws {@link FunctionExistsError} if an entrypoint already exists
   */
  addParticipantEntrypoint(
    callback: (job: JobContext<ProcessUserData>, p: RemoteParticipant) => Promise<void>,
  ) {
    if (this.#participantEntrypoints.includes(callback)) {
      throw new FunctionExistsError('entrypoints cannot be added more than once');
    }

    this.#participantEntrypoints.push(callback);
  }

  async initRecording(options: ResolvedRecordingOptions) {
    if (this.isFakeJob) {
      return;
    }

    const url = new URL(this.#info.url);
    if (!isCloud(url)) {
      return;
    }

    // The cloud tracer handles trace spans and OTel logs; only configure it
    // when at least one of those categories is enabled.
    if (!options.traces && !options.logs) {
      return;
    }

    this.#logger.debug({ hostname: url.hostname }, 'Configuring session recording (cloud tracer)');
    await setupCloudTracer({
      roomId: this.job.room!.sid,
      jobId: this.job.id,
      cloudHostname: url.hostname,
      enableTraces: options.traces,
      enableLogs: options.logs,
    });
  }
}

export class JobProcess<UserData = Record<string, unknown>> {
  #pid = process.pid;
  userData = {} as UserData;

  get pid(): number {
    return this.#pid;
  }
}

/**
 * A request sent by the server to spawn a new agent job.
 *
 * @remarks
 * For most applications, this is best left to the default, which simply accepts the job and
 * handles the logic inside the entrypoint function. This class is useful for vetting which
 * requests should fill idle processes and which should be outright rejected.
 */
export class JobRequest {
  #job: proto.Job;
  #onReject: () => Promise<void>;
  #onAccept: (args: JobAcceptArguments) => Promise<void>;

  /** @internal */
  constructor(
    job: proto.Job,
    onReject: () => Promise<void>,
    onAccept: (args: JobAcceptArguments) => Promise<void>,
  ) {
    this.#job = job;
    this.#onReject = onReject;
    this.#onAccept = onAccept;
  }

  /** @returns The ID of the job, set by the LiveKit server */
  get id(): string {
    return this.#job.id;
  }

  /** @see {@link https://www.npmjs.com/package/@livekit/protocol | @livekit/protocol} */
  get job(): proto.Job {
    return this.#job;
  }

  /** @see {@link https://www.npmjs.com/package/@livekit/protocol | @livekit/protocol} */
  get room(): proto.Room | undefined {
    return this.#job.room;
  }

  /** @see {@link https://www.npmjs.com/package/@livekit/protocol | @livekit/protocol} */
  get publisher(): proto.ParticipantInfo | undefined {
    return this.#job.participant;
  }

  /** @returns The agent's name, as set in {@link WorkerOptions} */
  get agentName(): string {
    return this.#job.agentName;
  }

  /** Rejects the job. */
  async reject() {
    await this.#onReject();
  }

  /** Accepts the job, launching it on an idle child process. */
  async accept(name = '', identity = '', metadata = '', attributes?: { [key: string]: string }) {
    if (identity === '') identity = 'agent-' + this.id;

    this.#onAccept({ name, identity, metadata, attributes });
  }
}
