import { AsyncLocalStorage } from "node:async_hooks";
import { randomUUID } from "node:crypto";

import type { LettaCodeClientSessionOptions } from "@letta-ai/letta-agent-sdk";

import { ResultCollection } from "./collection.js";
import { parseOutput, structuredPrompt } from "./output.js";
import { SdkFlowClient } from "./sdk-client.js";
import { Semaphore } from "./semaphore.js";
import type {
  FlowClient,
  FlowCollection,
  FlowContext,
  FlowDefaults,
  FlowError,
  FlowEvent,
  FlowEventListener,
  FlowOperation,
  FlowOptions,
  FlowResult,
  FlowSession,
  FreshInput,
  FreshOptions,
  ForkOptions,
  ResumeOptions,
  RunMetadata,
} from "./types.js";

const DEFAULT_CONCURRENCY = 4;
const DEFAULT_STRAGGLER_GRACE_MS = 5000;

export function flow<T>(
  run: (context: FlowContext) => Promise<T>,
): Promise<T>;
export function flow<T>(
  options: FlowOptions,
  run: (context: FlowContext) => Promise<T>,
): Promise<T>;
export async function flow<T>(
  optionsOrRun: FlowOptions | ((context: FlowContext) => Promise<T>),
  maybeRun?: (context: FlowContext) => Promise<T>,
): Promise<T> {
  const options = typeof optionsOrRun === "function" ? {} : optionsOrRun;
  const run = typeof optionsOrRun === "function" ? optionsOrRun : maybeRun;
  if (!run) throw new Error("flow() requires a workflow callback");

  const env = options.env ?? process.env;
  const agentId = resolveAgentId(options, env);
  if (!agentId) {
    throw new Error(
      "flows requires an agentId, LETTA_AGENT_ID, or AGENT_ID; it will not silently create an agent",
    );
  }

  const scope = new FlowScope({ ...options, agentId });
  return scope.run(run);
}

function resolveAgentId(
  options: FlowOptions,
  env: Record<string, string | undefined>,
): string | undefined {
  if (options.agentId !== undefined) {
    const explicit = options.agentId.trim();
    if (!explicit) {
      throw new Error("flows received an empty agentId");
    }
    return explicit;
  }
  const lettaAgentId = env.LETTA_AGENT_ID?.trim();
  if (lettaAgentId) return lettaAgentId;
  const harnessAgentId = env.AGENT_ID?.trim();
  return harnessAgentId || undefined;
}

function assertConversationId(action: "resume" | "fork", value: string): void {
  if (!/^(?:conv-|local-conv-).+/.test(value)) {
    throw new Error(
      `${action}() requires a conversation ID beginning with conv- or local-conv-, got ${JSON.stringify(value)}`,
    );
  }
}

/** Deadline exceeded while the flow was still running. */
export class FlowDeadlineError extends Error {
  readonly flowId: string;
  constructor(flowId: string) {
    super(`Flow ${flowId} exceeded its time budget and was aborted`);
    this.name = "FlowDeadlineError";
    this.flowId = flowId;
  }
}

class FlowScope implements FlowContext {
  readonly flowId = `flow-${randomUUID()}`;
  readonly agentId: string;
  readonly #client: FlowClient;
  readonly #concurrency: number;
  readonly #semaphore: Semaphore;
  readonly #listener: FlowEventListener | undefined;
  readonly #defaults: FlowDefaults | undefined;
  readonly #budget:
    | { limitMs: number; turnDeadlineMs?: number; stragglerGraceMs: number }
    | undefined;
  readonly #phases = new AsyncLocalStorage<string>();
  readonly #sessions = new Set<FlowSession>();
  readonly #operations = new Set<Promise<unknown>>();
  #state: "running" | "closing" | "closed" = "running";
  #deadlineTimer: ReturnType<typeof setTimeout> | undefined;
  #deadlineFired = false;

  constructor(options: FlowOptions & { agentId: string }) {
    this.agentId = options.agentId;
    this.#concurrency = options.concurrency ?? DEFAULT_CONCURRENCY;
    this.#semaphore = new Semaphore(this.#concurrency);
    if (options.client && options.sdkClient) {
      throw new Error("Pass either client or sdkClient to flow(), not both");
    }
    this.#client =
      options.client ?? new SdkFlowClient(options.sdkClient, this.agentId);
    this.#listener = options.onEvent;
    this.#defaults = options.defaults;
    this.#budget = undefined;
    if (options.timeBudget !== undefined) {
      if (
        !Number.isFinite(options.timeBudget.limitMs) ||
        options.timeBudget.limitMs <= 0
      ) {
        throw new Error("timeBudget.limitMs must be a positive number");
      }
      if (
        options.timeBudget.turnDeadlineMs !== undefined &&
        (!Number.isFinite(options.timeBudget.turnDeadlineMs) ||
          options.timeBudget.turnDeadlineMs <= 0)
      ) {
        throw new Error("timeBudget.turnDeadlineMs must be a positive number");
      }
      if (
        options.timeBudget.stragglerGraceMs !== undefined &&
        (!Number.isFinite(options.timeBudget.stragglerGraceMs) ||
          options.timeBudget.stragglerGraceMs <= 0)
      ) {
        throw new Error(
          "timeBudget.stragglerGraceMs must be a positive number",
        );
      }
      this.#budget = {
        limitMs: options.timeBudget.limitMs,
        ...(options.timeBudget.turnDeadlineMs !== undefined
          ? { turnDeadlineMs: options.timeBudget.turnDeadlineMs }
          : {}),
        stragglerGraceMs:
          options.timeBudget.stragglerGraceMs ?? DEFAULT_STRAGGLER_GRACE_MS,
      };
    }
  }

  async run<T>(run: (context: FlowContext) => Promise<T>): Promise<T> {
    const started = Date.now();
    this.#emit({
      type: "flow.started",
      flowId: this.flowId,
      agentId: this.agentId,
      concurrency: this.#concurrency,
      timestamp: new Date(),
    });

    let value: T | undefined;
    let failure: unknown;
    try {
      value = await this.#withBudget(run);
    } catch (error) {
      failure = error;
    }

    this.#state = "closing";
    this.#clearDeadlineTimer();
    await this.#awaitOperations();
    await this.#cleanupSessions();
    this.#state = "closed";

    if (failure === undefined) {
      this.#emit({
        type: "flow.completed",
        flowId: this.flowId,
        agentId: this.agentId,
        durationMs: Date.now() - started,
        timestamp: new Date(),
      });
      return value as T;
    }

    this.#emit({
      type: "flow.failed",
      flowId: this.flowId,
      agentId: this.agentId,
      durationMs: Date.now() - started,
      error: failure,
      timestamp: new Date(),
    });
    throw failure;
  }

  async #withBudget<T>(run: (context: FlowContext) => Promise<T>): Promise<T> {
    const limitMs = this.#budget?.limitMs;
    if (limitMs === undefined) return run(this);

    return new Promise<T>((resolve, reject) => {
      this.#deadlineTimer = setTimeout(() => {
        this.#deadlineFired = true;
        this.#emit({
          type: "flow.deadline",
          flowId: this.flowId,
          agentId: this.agentId,
          timestamp: new Date(),
        });
        // Abort every tracked in-flight session; each prompt() stream then
        // delivers a terminal result, so operations settle and the flow's
        // cleanup can close sessions for real instead of detaching.
        void this.#abortInFlight();
        reject(new FlowDeadlineError(this.flowId));
      }, limitMs);
      this.#deadlineTimer.unref?.();
      run(this).then(resolve, reject);
    });
  }

  #clearDeadlineTimer(): void {
    if (this.#deadlineTimer !== undefined) {
      clearTimeout(this.#deadlineTimer);
      this.#deadlineTimer = undefined;
    }
  }

  /**
   * Wait for tracked operations to settle. Normally every operation settles
   * on its own. After a deadline abort, sessions that honor abort() settle
   * promptly (the abort makes the SDK stream deliver a terminal result), but
   * a session that ignores abort would otherwise hang the flow forever — so
   * the deadline path additionally races a bounded grace period and abandons
   * any stragglers. Abandoned operations are not cancelled (that is not
   * possible in JS); their sessions are still closed by cleanup, and their
   * eventual rejection is swallowed by the #track() bookkeeping.
   */
  async #awaitOperations(): Promise<void> {
    const operations = [...this.#operations];
    if (operations.length === 0) return;
    if (!this.#deadlineFired) {
      await Promise.allSettled(operations);
      return;
    }
    const graceMs = this.#budget?.stragglerGraceMs ?? DEFAULT_STRAGGLER_GRACE_MS;
    await Promise.race([
      Promise.allSettled(operations),
      new Promise<void>((resolve) => {
        const timer = setTimeout(resolve, graceMs);
        timer.unref?.();
      }),
    ]);
  }

  async #abortInFlight(): Promise<void> {
    await Promise.allSettled(
      [...this.#sessions].map((session) =>
        session.abort === undefined
          ? Promise.reject(new Error("session does not support abort"))
          : session.abort(),
      ),
    );
  }

  fresh<T = string>(input: FreshInput<T>): Promise<FlowResult<T>> {
    this.#assertRunning();
    const options: FreshOptions<T> =
      typeof input === "string" ? { prompt: input } : input;
    return this.#track(
      this.#semaphore.use(() => this.#turn("fresh", options)),
    );
  }

  resume<T = string>(input: ResumeOptions<T>): Promise<FlowResult<T>> {
    this.#assertRunning();
    assertConversationId("resume", input.conversationId);
    return this.#track(
      this.#semaphore.use(() =>
        this.#turn("resume", input, input.conversationId),
      ),
    );
  }

  fork<T = string>(input: ForkOptions<T>): Promise<FlowResult<T>> {
    this.#assertRunning();
    assertConversationId("fork", input.sourceConversationId);
    if (input.stateless === true) {
      throw new Error(
        "fork() cannot be stateless because stateless sessions do not load inherited conversation history",
      );
    }
    if (input.messageId !== undefined && input.messageId.trim().length === 0) {
      throw new Error("fork() requires a non-empty messageId when provided");
    }
    return this.#track(
      this.#semaphore.use(() =>
        this.#turn("fork", input, input.sourceConversationId),
      ),
    );
  }

  async parallel<T>(tasks: Array<() => Promise<T>>): Promise<FlowCollection<T>> {
    this.#assertRunning();
    const values = await Promise.all(tasks.map((task) => task()));
    return new ResultCollection(values);
  }

  async map<T, U>(
    items: readonly T[],
    mapper: (item: T, index: number) => Promise<U>,
  ): Promise<FlowCollection<U>> {
    this.#assertRunning();
    return this.parallel(items.map((item, index) => () => mapper(item, index)));
  }

  async phase<T>(name: string, run: () => Promise<T>): Promise<T> {
    this.#assertRunning();
    this.#emit({
      type: "phase.started",
      flowId: this.flowId,
      name,
      timestamp: new Date(),
    });
    try {
      return await this.#phases.run(name, run);
    } finally {
      this.#emit({
        type: "phase.completed",
        flowId: this.flowId,
        name,
        timestamp: new Date(),
      });
    }
  }

  async #turn<T>(
    operation: FlowOperation,
    options: FreshOptions<T>,
    conversationId?: string,
  ): Promise<FlowResult<T>> {
    const merged = this.#mergeDefaults(operation, options);
    const operationId = `operation-${randomUUID()}`;
    const phase = this.#phases.getStore();
    const startedAt = new Date();
    const base = {
      operationId,
      operation,
      agentId: this.agentId,
      ...(operation === "resume" && conversationId !== undefined
        ? { conversationId }
        : {}),
      ...(operation === "fork" && conversationId !== undefined
        ? { sourceConversationId: conversationId }
        : {}),
      ...(operation === "fork" && (merged as ForkOptions<T>).messageId !== undefined
        ? { sourceMessageId: (merged as ForkOptions<T>).messageId }
        : {}),
      ...(phase !== undefined ? { phase } : {}),
      ...(merged.label !== undefined ? { label: merged.label } : {}),
    };

    this.#emit({
      type: "operation.started",
      flowId: this.flowId,
      ...base,
      timestamp: startedAt,
    });

    let session: FlowSession | undefined;
    let finalResult: FlowResult<T> | undefined;
    let turnTimer: ReturnType<typeof setTimeout> | undefined;
    try {
      if (this.#deadlineFired) throw new FlowDeadlineError(this.flowId);
      const sessionOptions = toSessionOptions(merged);
      if (operation === "fresh") {
        session = await this.#client.fresh(this.agentId, sessionOptions);
      } else if (operation === "resume") {
        session = await this.#client.resume(conversationId!, sessionOptions);
      } else {
        const fork = merged as ForkOptions<T>;
        session = await this.#client.fork(
          this.agentId,
          fork.sourceConversationId,
          {
            ...(fork.messageId !== undefined ? { messageId: fork.messageId } : {}),
            ...(fork.hidden !== undefined ? { hidden: fork.hidden } : {}),
          },
          sessionOptions,
        );
      }
      this.#sessions.add(session);
      if (this.#deadlineFired) throw new FlowDeadlineError(this.flowId);

      const identity = await session.prepare();
      if (this.#deadlineFired) throw new FlowDeadlineError(this.flowId);
      if (identity.agentId !== this.agentId) {
        const completedAt = new Date();
        finalResult = {
          ok: false,
          error: {
            code: "agent_mismatch",
            message: `Conversation belongs to ${identity.agentId}, not root agent ${this.agentId}`,
          },
          run: {
            ...base,
            conversationId: identity.conversationId,
            ...(identity.model ? { model: identity.model } : {}),
            startedAt,
            completedAt,
            durationMs: completedAt.getTime() - startedAt.getTime(),
          },
        };
        return finalResult;
      }

      const prompt = merged.output
        ? structuredPrompt(merged.prompt)
        : merged.prompt;

      const turnDeadlineMs = this.#budget?.turnDeadlineMs;
      if (turnDeadlineMs !== undefined && session.abort !== undefined) {
        turnTimer = setTimeout(() => {
          void Promise.resolve(session!.abort!()).catch(() => {});
        }, turnDeadlineMs);
        turnTimer.unref?.();
      }

      const turn = await session.prompt(prompt);
      const metadata = completeMetadata(
        base,
        startedAt,
        session,
        turn,
        identity.model,
        merged.model,
      );

      let result: FlowResult<T>;
      if (!turn.success) {
        result = {
          ok: false,
          error: {
            code: "turn_failed",
            message:
              [turn.errorCode, turn.error, turn.errorDetail]
                .filter(Boolean)
                .join(": ") || "Agent turn failed",
            ...(turn.recoverable !== undefined
              ? { recoverable: turn.recoverable }
              : {}),
          },
          run: metadata,
        };
      } else if (turn.result === undefined) {
        result = {
          ok: false,
          error: { code: "no_output", message: "Agent turn completed without output" },
          run: metadata,
        };
      } else if (merged.output) {
        const parsed = parseOutput(turn.result, merged.output);
        result = parsed.ok
          ? { ok: true, value: parsed.value, run: metadata }
          : { ok: false, error: parsed.error, run: metadata };
      } else {
        result = { ok: true, value: turn.result as T, run: metadata };
      }

      finalResult = result;
      return result;
    } catch (cause) {
      const completedAt = new Date();
      const error: FlowError = {
        code:
          operation === "fork" && !session
            ? "fork_failed"
            : session
              ? "operation_failed"
              : "session_create_failed",
        message: cause instanceof Error ? cause.message : String(cause),
        cause,
      };
      const result: FlowResult<T> = {
        ok: false,
        error,
        run: {
          ...base,
          ...(session?.conversationId
            ? { conversationId: session.conversationId }
            : {}),
          startedAt,
          completedAt,
          durationMs: completedAt.getTime() - startedAt.getTime(),
        },
      };
      finalResult = result;
      return result;
    } finally {
      if (turnTimer !== undefined) clearTimeout(turnTimer);
      if (session) {
        try {
          await this.#boundedCleanup(() => session!.close());
          this.#sessions.delete(session);
        } catch (error) {
          this.#emit({
            type: "cleanup.failed",
            flowId: this.flowId,
            operationId,
            error,
            timestamp: new Date(),
          });
        }
      }
      if (finalResult) this.#emitCompletion(finalResult);
    }
  }

  /**
   * Merge flow-level defaults into per-operation options. Per-operation
   * fields always win; a per-operation `session` replaces the default
   * session wholesale so a turn's session shape is explicit, never merged.
   * Incompatible combinations (stateless fork) are rejected before any
   * session or conversation is created.
   */
  #mergeDefaults<T>(
    operation: FlowOperation,
    options: FreshOptions<T>,
  ): FreshOptions<T> {
    const defaults = this.#defaults;
    if (defaults === undefined) return options;

    if (operation === "fork") {
      const effectiveStateless = options.stateless ?? defaults.stateless;
      if (effectiveStateless === true) {
        throw new Error(
          "fork() cannot be stateless because stateless sessions do not load inherited conversation history",
        );
      }
    }

    return {
      ...options,
      ...(defaults.session !== undefined && options.session === undefined
        ? { session: defaults.session }
        : {}),
      ...(defaults.stateless !== undefined && options.stateless === undefined
        ? { stateless: defaults.stateless }
        : {}),
      ...(defaults.model !== undefined && options.model === undefined
        ? { model: defaults.model }
        : {}),
      ...(defaults.reasoningEffort !== undefined &&
      options.reasoningEffort === undefined
        ? { reasoningEffort: defaults.reasoningEffort }
        : {}),
    };
  }

  #track<T>(operation: Promise<T>): Promise<T> {
    this.#operations.add(operation);
    void operation.then(
      () => this.#operations.delete(operation),
      () => this.#operations.delete(operation),
    );
    return operation;
  }

  #assertRunning(): void {
    if (this.#state !== "running") {
      throw new Error(`Flow scope is ${this.#state}; no new work can be started`);
    }
  }

  async #cleanupSessions(): Promise<void> {
    await Promise.allSettled(
      [...this.#sessions].map(async (session) => {
        try {
          await this.#boundedCleanup(() => session.close());
          this.#sessions.delete(session);
        } catch (error) {
          this.#emit({
            type: "cleanup.failed",
            flowId: this.flowId,
            error,
            timestamp: new Date(),
          });
        }
      }),
    );
    await this.#disposeClient();
  }

  async #disposeClient(): Promise<void> {
    const client = this.#client as AsyncDisposable | { [Symbol.asyncDispose]?: unknown };
    const dispose = (client as { [Symbol.asyncDispose]?: () => Promise<void> })[
      Symbol.asyncDispose
    ];
    if (typeof dispose === "function") {
      try {
        await this.#boundedCleanup(() => dispose.call(client));
      } catch (error) {
        this.#emit({
          type: "cleanup.failed",
          flowId: this.flowId,
          error,
          timestamp: new Date(),
        });
      }
    }
  }

  async #boundedCleanup(action: () => void | Promise<void>): Promise<void> {
    let timer: ReturnType<typeof setTimeout> | undefined;
    try {
      await Promise.race([
        Promise.resolve().then(action),
        new Promise<never>((_, reject) => {
          timer = setTimeout(() => reject(new Error("Cleanup timed out; resource termination is unconfirmed")),
            this.#budget?.stragglerGraceMs ?? DEFAULT_STRAGGLER_GRACE_MS);
        }),
      ]);
    } finally {
      if (timer !== undefined) clearTimeout(timer);
    }
  }

  #emitCompletion<T>(result: FlowResult<T>): void {
    this.#emit({
      type: "operation.completed",
      flowId: this.flowId,
      result: result as FlowResult<unknown>,
      timestamp: new Date(),
    });
  }

  #emit(event: FlowEvent): void {
    try {
      this.#listener?.(event);
    } catch {
      // Observers must not alter workflow execution.
    }
  }
}

function toSessionOptions<T>(options: FreshOptions<T>): LettaCodeClientSessionOptions {
  return {
    ...options.session,
    ...(options.stateless !== undefined ? { stateless: options.stateless } : {}),
    ...(options.model !== undefined ? { model: options.model } : {}),
    ...(options.reasoningEffort !== undefined
      ? { reasoningEffort: options.reasoningEffort }
      : {}),
  };
}

function completeMetadata(
  base: Omit<RunMetadata, "startedAt" | "completedAt" | "durationMs">,
  startedAt: Date,
  session: FlowSession,
  turn: Awaited<ReturnType<FlowSession["prompt"]>>,
  model?: string,
  configuredModel?: string,
): RunMetadata {
  const completedAt = new Date();
  return {
    ...base,
    ...(session.agentId ? { agentId: session.agentId } : {}),
    ...(session.conversationId
      ? { conversationId: session.conversationId }
      : {}),
    ...(model ? { model } : {}),
    ...(configuredModel ? { configuredModel } : {}),
    startedAt,
    completedAt,
    durationMs: turn.durationMs ?? completedAt.getTime() - startedAt.getTime(),
    ...(turn.runIds !== undefined ? { runIds: turn.runIds } : {}),
    ...(turn.stopReason !== undefined ? { stopReason: turn.stopReason } : {}),
    ...(turn.totalCostUsd !== undefined ? { costUsd: turn.totalCostUsd } : {}),
  };
}