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

export type MaybePromise<T> = T | Promise<T>;
export type FlowOperation = "fresh" | "resume" | "fork";

export interface OutputParser<T> {
  parse(input: unknown): T;
}

export type FlowErrorCode =
  | "missing_agent_id"
  | "session_create_failed"
  | "fork_failed"
  | "agent_mismatch"
  | "turn_failed"
  | "no_output"
  | "invalid_json"
  | "invalid_output"
  | "operation_failed";

export interface FlowError {
  code: FlowErrorCode;
  message: string;
  cause?: unknown;
  recoverable?: boolean;
}

export interface RunMetadata {
  operationId: string;
  operation: FlowOperation;
  agentId: string;
  conversationId?: string;
  sourceConversationId?: string;
  sourceMessageId?: string;
  /** Model the server reported for this conversation, if it reported one. */
  model?: string;
  /** Model handle the flow configured for this operation, if any. */
  configuredModel?: string;
  phase?: string;
  label?: string;
  startedAt: Date;
  completedAt: Date;
  durationMs: number;
  runIds?: string[];
  stopReason?: string;
  costUsd?: number;
}

export type FlowResult<T> =
  | { ok: true; value: T; run: RunMetadata }
  | { ok: false; error: FlowError; run: RunMetadata };

export interface FreshOptions<T = string> {
  prompt: string;
  label?: string;
  output?: OutputParser<T>;
  stateless?: boolean;
  /** Sets a persistent model override on this new or resumed conversation. */
  model?: string;
  /** Sets the reasoning tier on this new or resumed conversation. */
  reasoningEffort?: ReasoningEffort;
  session?: Omit<
    LettaCodeClientSessionOptions,
    "stateless" | "model" | "reasoningEffort"
  >;
}

/**
 * Flow-wide defaults for agent turns. Every field can be overridden per
 * operation; a per-operation `session` object replaces the flow default
 * wholesale rather than merging, so what a turn receives is always explicit.
 */
export interface FlowDefaults {
  /** Default model override for fresh, fork, and resume turns. */
  model?: string;
  /** Default reasoning tier for fresh, fork, and resume turns. */
  reasoningEffort?: ReasoningEffort;
  /** Default session options for every turn. */
  session?: Omit<
    LettaCodeClientSessionOptions,
    "stateless" | "model" | "reasoningEffort"
  >;
  /**
   * Default stateless mode. Rejected for fork turns regardless of source,
   * because a stateless session would not load the inherited transcript.
   */
  stateless?: boolean;
}

/**
 * Wall-clock budget for the whole flow. When it elapses, in-flight turns are
 * aborted through the session (the runtime stops the model turn and the
 * stream delivers a result), the flow callback's continuation sees failures,
 * and the flow rejects with a deadline error. This closes sessions and
 * aborts work; it does not detach from still-running turns.
 */
export interface FlowTimeBudget {
  /** Total wall-clock milliseconds the flow may run. */
  limitMs: number;
  /**
   * Milliseconds a single turn may run before it is aborted. Turns aborted
   * by this deadline resolve as failed FlowResults, not flow rejections.
   */
  turnDeadlineMs?: number;
  /**
   * Grace period the flow waits for in-flight operations to settle after
   * the deadline aborts them. Real SDK sessions settle promptly (the abort
   * makes the stream deliver a terminal result); a session that ignores
   * abort is abandoned after this window instead of hanging the process.
   * Default 5000.
   */
  stragglerGraceMs?: number;
}

export type FreshInput<T = string> = string | FreshOptions<T>;

export interface ResumeOptions<T = string> extends FreshOptions<T> {
  conversationId: string;
}

export interface ForkOptions<T = string>
  extends Omit<FreshOptions<T>, "stateless"> {
  /** Forks must load their inherited transcript. */
  stateless?: never;
  /** Conversation whose in-context history should be copied. */
  sourceConversationId: string;
  /** Copy source history through this message, inclusive. Omit for full history. */
  messageId?: string;
  /** Hide the persistent fork from normal conversation listings. */
  hidden?: boolean;
}

export interface FlowForkSettings {
  messageId?: string;
  hidden?: boolean;
}

export type LazyTask<T> = () => Promise<T>;

export interface FlowCollection<T> {
  all(): readonly T[];
  values<V>(this: FlowCollection<FlowResult<V>>): V[];
  failures<V>(this: FlowCollection<FlowResult<V>>): Array<Extract<FlowResult<V>, { ok: false }>>;
}

export type FlowEvent =
  | {
      type: "flow.started";
      flowId: string;
      agentId: string;
      concurrency: number;
      timestamp: Date;
    }
  | {
      type: "flow.completed";
      flowId: string;
      agentId: string;
      durationMs: number;
      timestamp: Date;
    }
  | {
      type: "flow.failed";
      flowId: string;
      agentId: string;
      durationMs: number;
      error: unknown;
      timestamp: Date;
    }
  | {
      type: "flow.deadline";
      flowId: string;
      agentId: string;
      timestamp: Date;
    }
  | {
      type: "phase.started" | "phase.completed";
      flowId: string;
      name: string;
      timestamp: Date;
    }
  | {
      type: "operation.started";
      flowId: string;
      operationId: string;
      operation: FlowOperation;
      agentId: string;
      conversationId?: string;
      sourceConversationId?: string;
      sourceMessageId?: string;
      phase?: string;
      label?: string;
      timestamp: Date;
    }
  | {
      type: "operation.completed";
      flowId: string;
      result: FlowResult<unknown>;
      timestamp: Date;
    }
  | {
      type: "cleanup.failed";
      flowId: string;
      operationId?: string;
      error: unknown;
      timestamp: Date;
    };

export type FlowEventListener = (event: FlowEvent) => void;

export interface FlowOptions {
  agentId?: string;
  concurrency?: number;
  client?: FlowClient;
  sdkClient?: LettaAgentClient;
  env?: Record<string, string | undefined>;
  onEvent?: FlowEventListener;
  /** Flow-wide turn defaults; each operation can override every field. */
  defaults?: FlowDefaults;
  /** Wall-clock budget that aborts in-flight turns when it elapses. */
  timeBudget?: FlowTimeBudget;
}

export interface FlowTurnResult {
  success: boolean;
  result?: string;
  error?: string;
  errorCode?: string;
  errorDetail?: string;
  recoverable?: boolean;
  runIds?: string[];
  stopReason?: string;
  durationMs?: number;
  totalCostUsd?: number;
}

export interface FlowSession {
  readonly agentId?: string | null;
  readonly conversationId?: string | null;
  prepare(): Promise<{ agentId: string; conversationId: string; model?: string }>;
  prompt(message: string): Promise<FlowTurnResult>;
  /**
   * Ask the runtime to stop the in-flight turn. The prompt() stream then
   * delivers a terminal result instead of hanging. Optional for custom
   * clients; time budgets require it to actually abort.
   */
  abort?(): MaybePromise<void>;
  close(): MaybePromise<void>;
}

export interface FlowClient {
  fresh(agentId: string, options: LettaCodeClientSessionOptions): MaybePromise<FlowSession>;
  resume(conversationId: string, options: LettaCodeClientSessionOptions): MaybePromise<FlowSession>;
  /** Implementations must reject a source owned by another agent before forking it. */
  fork(
    agentId: string,
    sourceConversationId: string,
    fork: FlowForkSettings,
    options: LettaCodeClientSessionOptions,
  ): MaybePromise<FlowSession>;
}

export interface FlowContext {
  readonly flowId: string;
  readonly agentId: string;
  fresh<T = string>(input: FreshInput<T>): Promise<FlowResult<T>>;
  resume<T = string>(input: ResumeOptions<T>): Promise<FlowResult<T>>;
  fork<T = string>(input: ForkOptions<T>): Promise<FlowResult<T>>;
  parallel<T>(tasks: Array<LazyTask<T>>): Promise<FlowCollection<T>>;
  map<T, U>(items: readonly T[], mapper: (item: T, index: number) => Promise<U>): Promise<FlowCollection<U>>;
  phase<T>(name: string, run: () => Promise<T>): Promise<T>;
}