import { EventObject, AnyStateMachine, AnyEventObject, AnyActorRef, SnapshotFrom, EventFrom, PromiseActorLogic, ActorLogic, TransitionSnapshot, Values, ActorRefFrom, Subscription, StateValue, ObservableActorLogic } from 'xstate';
import { SomeZodObject, ZodType, TypeOf } from 'zod';
import { generateText, streamText, LanguageModel, CoreMessage, GenerateTextResult, StreamTextResult, CoreTool } from 'ai';

type ZodEventMapping = {
    [eventType: string]: SomeZodObject;
};
type ZodContextMapping = {
    [contextKey: string]: ZodType;
};

type GenerateTextOptions = Parameters<typeof generateText>[0];
type StreamTextOptions = Parameters<typeof streamText>[0];
type AgentPlanInput<TEvent extends EventObject> = Omit<GenerateTextOptions, 'prompt' | 'messages' | 'tools'> & {
    /**
     * The currently observed state.
     */
    state: ObservedState;
    /**
     * The goal for the agent to accomplish.
     * The agent will create a plan based on this goal.
     */
    goal: string;
    /**
     * The events that the agent can trigger. This is a mapping of
     * event types to Zod event schemas.
     */
    events: ZodEventMapping;
    /**
     * The state machine that represents the environment the agent
     * is interacting with.
     */
    machine?: AnyStateMachine;
    /**
     * The previous plan.
     */
    previousPlan?: AgentPlan<TEvent>;
};
type AgentPlan<TEvent extends EventObject> = {
    goal: string;
    state: ObservedState;
    content?: string;
    /**
     * Executes the plan based on the given `state` and resolves with
     * a potential next `event` to trigger to achieve the `goal`.
     */
    execute: (state: ObservedState) => Promise<TEvent | undefined>;
    nextEvent: TEvent | undefined;
    sessionId: string;
    timestamp: number;
};
interface TransitionData {
    eventType: string;
    description?: string;
    guard?: {
        type: string;
    };
    target?: any;
}
type PromptTemplate<TEvents extends EventObject> = (data: {
    goal: string;
    /**
     * The observed state
     */
    state?: ObservedState;
    /**
     * The context to provide.
     * This overrides the observed state.context, if provided.
     */
    context?: any;
    /**
     * The state machine model of the observed environment
     */
    machine?: unknown;
    /**
     * The potential next transitions that can be taken
     * in the state machine
     */
    transitions?: TransitionData[];
    /**
     * Past observations
     */
    observations?: AgentObservation<any>[];
    feedback?: AgentFeedback[];
    messages?: AgentMessage[];
    plans?: AgentPlan<TEvents>[];
}) => string;
type AgentPlanner<T extends AnyAgent> = (agent: T, input: AgentPlanInput<T['types']['events']>) => Promise<AgentPlan<T['types']['events']> | undefined>;
type AgentDecideOptions = {
    goal: string;
    model?: LanguageModel;
    context?: any;
    state: ObservedState;
    machine: AnyStateMachine;
    execute?: (event: AnyEventObject) => Promise<void>;
    planner?: AgentPlanner<any>;
    events?: ZodEventMapping;
} & Omit<Parameters<typeof generateText>[0], 'model' | 'tools' | 'prompt' | 'messages'>;
interface AgentFeedback {
    goal?: string;
    observationId?: string;
    /**
     * The message correlation that the feedback is relevant for
     */
    correlationId?: string;
    attributes: Record<string, any>;
    reward: number;
    timestamp: number;
    sessionId: string;
}
interface AgentFeedbackInput {
    goal?: string;
    observationId?: string;
    correlationId?: string;
    attributes?: Record<string, any>;
    timestamp?: number;
    reward?: number;
}
type AgentMessage = CoreMessage & {
    timestamp: number;
    id: string;
    /**
     * The response ID of the message, which references
     * which message this message is responding to, if any.
     */
    responseId?: string;
    result?: GenerateTextResult<any>;
    sessionId: string;
    correlationId: string;
    parentCorrelationId?: string;
};
type AgentMessageInput = CoreMessage & {
    timestamp?: number;
    id?: string;
    /**
     * The response ID of the message, which references
     * which message this message is responding to, if any.
     */
    responseId?: string;
    correlationId?: string;
    parentCorrelationId?: string;
    result?: GenerateTextResult<any>;
};
interface AgentObservation<TActor extends AnyActorRef> {
    id: string;
    prevState: SnapshotFrom<TActor> | undefined;
    event: EventFrom<TActor>;
    state: SnapshotFrom<TActor>;
    machineHash: string | undefined;
    sessionId: string;
    timestamp: number;
}
interface AgentObservationInput {
    id?: string;
    prevState: ObservedState | undefined;
    event: AnyEventObject;
    state: ObservedState;
    machine?: AnyStateMachine;
    timestamp?: number;
}
type AgentDecisionInput = {
    goal: string;
    model?: LanguageModel;
    context?: any;
} & Omit<Parameters<typeof generateText>[0], 'model' | 'tools' | 'prompt'>;
type AgentDecisionLogic<TEvents extends EventObject> = PromiseActorLogic<AgentPlan<TEvents> | undefined, AgentDecisionInput | string>;
type AgentEmitted<TEvents extends EventObject> = {
    type: 'feedback';
    feedback: AgentFeedback;
} | {
    type: 'observation';
    observation: AgentObservation<any>;
} | {
    type: 'message';
    message: AgentMessage;
} | {
    type: 'plan';
    plan: AgentPlan<TEvents>;
};
type AgentLogic<TEvents extends EventObject> = ActorLogic<TransitionSnapshot<AgentMemoryContext>, {
    type: 'agent.feedback';
    feedback: AgentFeedback;
} | {
    type: 'agent.observe';
    observation: AgentObservation<any>;
} | {
    type: 'agent.message';
    message: AgentMessage;
} | {
    type: 'agent.plan';
    plan: AgentPlan<TEvents>;
}, any, // TODO: input
any, AgentEmitted<TEvents>>;
type EventsFromZodEventMapping<TEventSchemas extends ZodEventMapping> = Values<{
    [K in keyof TEventSchemas & string]: {
        type: K;
    } & TypeOf<TEventSchemas[K]>;
}>;
type ContextFromZodContextMapping<TContextSchema extends ZodContextMapping> = {
    [K in keyof TContextSchema & string]: TypeOf<TContextSchema[K]>;
};
type Agent<TContext, TEvents extends EventObject> = ActorRefFrom<AgentLogic<TEvents>> & {
    /**
     * The name of the agent. All agents with the same name are related and
     * able to share experiences (observations, feedback) with each other.
     */
    name?: string;
    /**
     * The unique identifier for the agent.
     */
    id?: string;
    description?: string;
    events: ZodEventMapping;
    types: {
        events: TEvents;
        context: Compute<TContext>;
    };
    model: LanguageModel;
    defaultOptions: GenerateTextOptions;
    memory: AgentLongTermMemory | undefined;
    /**
     * The adapter used to perform LLM actions such as
     * `.generateText(…)` and `.streamText(…)`.
     *
     * Defaults to the Vercel AI SDK.
     */
    adapter: AIAdapter;
    /**
     * Resolves with an `AgentPlan` based on the information provided in the `options`, including:
     *
     * - The `goal` for the agent to achieve
     * - The observed current `state`
     * - The `machine` (e.g. a state machine) that specifies what can happen next
     * - Additional `context`
     */
    decide: (options: AgentDecideOptions) => Promise<AgentPlan<TEvents> | undefined>;
    generateText: (options: AgentGenerateTextOptions) => Promise<AgentGenerateTextResult>;
    streamText: (options: AgentStreamTextOptions) => Promise<AgentStreamTextResult>;
    addObservation: (observationInput: AgentObservationInput) => AgentObservation<any>;
    addMessage: (messageInput: AgentMessageInput) => AgentMessage;
    addFeedback: (feedbackInput: AgentFeedbackInput) => AgentFeedback;
    addPlan: (plan: AgentPlan<TEvents>) => void;
    /**
     * Called whenever the agent (LLM assistant) receives or sends a message.
     */
    onMessage: (callback: (message: AgentMessage) => void) => void;
    /**
     * Selects agent data from its context.
     *
     * @deprecated Select from `agent.getSnapshot().context` directly or:
     * - `agent.getMessages()`
     * - `agent.getObservations()`
     * - `agent.getFeedback()`
     * - `agent.getPlans()`
     */
    select: <T>(selector: (context: AgentMemoryContext) => T) => T;
    /**
     * Retrieves messages from the agent's short-term (local) memory.
     */
    getMessages: () => AgentMessage[];
    /**
     * Retrieves observations from the agent's short-term (local) memory.
     */
    getObservations: () => AgentObservation<Agent<TContext, TEvents>>[];
    /**
     * Retrieves feedback from the agent's short-term (local) memory.
     */
    getFeedback: () => AgentFeedback[];
    /**
     * Retrieves strategies from the agent's short-term (local) memory.
     */
    getPlans: () => AgentPlan<TEvents>[];
    /**
     * Interacts with this state machine actor by inspecting state transitions and storing them as observations.
     *
     * Observations contain the `prevState`, `event`, and current `state` of this
     * actor, as well as other properties that are useful when recalled.
     * These observations are stored in the `agent`'s short-term (local) memory
     * and can be retrieved via `agent.getObservations()`.
     *
     * @example
     * ```ts
     * // Only observes the actor's state transitions
     * agent.interact(actor);
     *
     * actor.start();
     * ```
     */
    interact<TActor extends AnyActorRef>(actorRef: TActor): Subscription;
    /**
     * Interacts with this state machine actor by:
     * 1. Inspecting state transitions and storing them as observations
     * 2. Deciding what to do next (which event to send the actor) based on
     * the agent input returned from `getInput(observation)`, if `getInput(…)` is provided as the 2nd argument.
     *
     * Observations contain the `prevState`, `event`, and current `state` of this
     * actor, as well as other properties that are useful when recalled.
     * These observations are stored in the `agent`'s short-term (local) memory
     * and can be retrieved via `agent.getObservations()`.
     *
     * @example
     * ```ts
     * // Observes the actor's state transitions and
     * // makes a decision if on the "summarize" state
     * agent.interact(actor, observed => {
     *   if (observed.state.matches('summarize')) {
     *     return {
     *       context: observed.state.context,
     *       goal: 'Summarize the message'
     *     }
     *   }
     * });
     *
     * actor.start();
     * ```
     */
    interact<TActor extends AnyActorRef>(actorRef: TActor, getInput: (observation: AgentObservation<TActor>) => AgentDecisionInput | undefined): Subscription;
};
type AnyAgent = Agent<any, any>;
type FromAgent<T> = T | ((agent: AnyAgent) => T | Promise<T>);
type CommonTextOptions = {
    prompt: FromAgent<string>;
    model?: LanguageModel;
    context?: Record<string, any>;
    messages?: FromAgent<CoreMessage[]>;
    template?: PromptTemplate<any>;
    correlationId?: string;
    parentCorrelationId?: string;
};
type TextResultMeta = {
    correlationId: string;
    parentCorrelationId?: string;
};
type AgentGenerateTextOptions = Omit<GenerateTextOptions, 'model' | 'prompt' | 'messages'> & CommonTextOptions;
type AgentGenerateTextResult = GenerateTextResult<any> & TextResultMeta;
type AgentStreamTextOptions = Omit<StreamTextOptions, 'model' | 'prompt' | 'messages'> & CommonTextOptions;
type AgentStreamTextResult = StreamTextResult<any> & TextResultMeta;
interface ObservedState {
    /**
     * The current state value of the state machine, e.g.
     * `"loading"` or `"processing"` or `"ready"`
     */
    value: StateValue;
    /**
     * Additional contextual data related to the current state
     */
    context: Record<string, unknown>;
}
type ObservedStateFrom<TActor extends AnyActorRef> = Pick<SnapshotFrom<TActor>, 'value' | 'context'>;
type AgentMemoryContext = {
    observations: AgentObservation<any>[];
    messages: AgentMessage[];
    plans: AgentPlan<any>[];
    feedback: AgentFeedback[];
};
type AgentMemory = AppendOnlyStorage<AgentMemoryContext>;
interface AppendOnlyStorage<T extends Record<string, any[]>> {
    append<K extends keyof T>(sessionId: string, key: K, item: T[K][0]): Promise<void>;
    getAll<K extends keyof T>(sessionId: string, key: K): Promise<T[K] | undefined>;
}
interface AgentLongTermMemory {
    get<K extends keyof AgentMemoryContext>(key: K): Promise<AgentMemoryContext[K]>;
    append<K extends keyof AgentMemoryContext>(key: K, item: AgentMemoryContext[K][0]): Promise<void>;
    set<K extends keyof AgentMemoryContext>(key: K, items: AgentMemoryContext[K]): Promise<void>;
}
interface AIAdapter {
    generateText: typeof generateText;
    streamText: typeof streamText;
}
type Compute<A extends any> = {
    [K in keyof A]: A[K];
} & unknown;

declare function createAgent<const TContextSchema extends ZodContextMapping, const TEventSchemas extends ZodEventMapping, TEvents extends EventObject = EventsFromZodEventMapping<TEventSchemas>, TContext = ContextFromZodContextMapping<TContextSchema>>({ name, description, model, events, context, planner, stringify, getMemory, logic, adapter, ...generateTextOptions }: {
    /**
     * The unique identifier for the agent.
     *
     * This should be the same across all sessions of a specific agent, as it can be
     * used to retrieve memory for this agent.
     *
     * @example
     * ```ts
     * const agent = createAgent({
     *  id: 'recipe-assistant',
     *  // ...
     * });
     * ```
     */
    id?: string;
    /**
     * The name of the agent
     */
    name?: string;
    /**
     * A description of the role of the agent
     */
    description?: string;
    /**
     * Events that the agent can cause (send) in an environment
     * that the agent knows about.
     */
    events: TEventSchemas;
    context?: TContextSchema;
    planner?: AgentPlanner<Agent<TContext, TEvents>>;
    stringify?: typeof JSON.stringify;
    /**
     * A function that retrieves the agent's long term memory
     */
    getMemory?: (agent: Agent<TContext, TEvents>) => AgentLongTermMemory;
    /**
     * Agent logic
     */
    logic?: AgentLogic<TEvents>;
    adapter?: AIAdapter;
} & GenerateTextOptions): Agent<TContext, TEvents>;

declare function fromTextStream<T extends AnyAgent>(agent: T, defaultOptions?: AgentStreamTextOptions): ObservableActorLogic<{
    textDelta: string;
}, Omit<AgentStreamTextOptions, 'context'> & {
    context?: AgentStreamTextOptions['context'];
}>;
declare function fromText<T extends AnyAgent>(agent: T, defaultOptions?: AgentGenerateTextOptions): PromiseActorLogic<GenerateTextResult<Record<string, CoreTool<any, any>>>, Omit<AgentGenerateTextOptions, 'context'> & {
    context?: AgentGenerateTextOptions['context'];
}>;

declare function fromDecision(agent: AnyAgent, defaultInput?: AgentDecisionInput): AgentDecisionLogic<any>;

export { type AIAdapter, type Agent, type AgentDecideOptions, type AgentDecisionInput, type AgentDecisionLogic, type AgentEmitted, type AgentFeedback, type AgentFeedbackInput, type AgentGenerateTextOptions, type AgentGenerateTextResult, type AgentLogic, type AgentLongTermMemory, type AgentMemory, type AgentMemoryContext, type AgentMessage, type AgentMessageInput, type AgentObservation, type AgentObservationInput, type AgentPlan, type AgentPlanInput, type AgentPlanner, type AgentStreamTextOptions, type AgentStreamTextResult, type AnyAgent, type AppendOnlyStorage, type CommonTextOptions, type Compute, type ContextFromZodContextMapping, type EventsFromZodEventMapping, type FromAgent, type GenerateTextOptions, type ObservedState, type ObservedStateFrom, type PromptTemplate, type StreamTextOptions, type TextResultMeta, type TransitionData, createAgent, fromDecision, fromText, fromTextStream };