import type { JsonValue } from "@earendil-works/chord";
import type {
  Api,
  AuthResult,
  Message,
  Model,
  Provider,
  Static,
  TSchema,
} from "@earendil-works/pi-ai";
import type { Environment } from "./environment.ts";
import type { TObject } from "typebox";
import { Value } from "typebox/value";
import type { EntryRecord } from "@earendil-works/pi-durable";

export type { EntryRecord, JsonValue };
export type { Environment, Process, Pty, PtyOptions, Socket, SpawnOptions } from "./environment.ts";

/** Who can call a tool: scripts in codemode (the harness's one direct tool), the model directly, or both. */
export type Exposure = "codemode" | "model" | "both";

export type Tool<
  S extends JsonValue = JsonValue,
  E extends JsonValue = never,
  P extends TSchema = TSchema,
  R extends JsonValue | undefined = JsonValue | undefined,
> = {
  readonly name: string;
  readonly description: string;
  readonly parameters: P;
  /** Type of the value scripts receive; only shapes the codemode declaration. */
  readonly returns?: TSchema;
  /** Default `"codemode"`. */
  readonly expose?: Exposure;
  /** Whether the tool is offered and callable in the conversation’s current environment. */
  enabled?(ctx: PromptContext): boolean | Promise<boolean>;
  /**
   * The value is what scripts receive. Called directly, strings become text, `{ type: "image", data, mimeType }`
   * an image, and anything else JSON text, unless `text` says otherwise.
   */
  execute(args: Static<P>, ctx: ToolContext<S, E>): R | Promise<R>;
  /** The value as text for the model's direct calls, where JSON would read poorly. */
  text?(value: R): string;
};

/** One conversation, as an extension sees it. */
export type ConversationContext<S extends JsonValue> = {
  readonly tid: string;
  /** The extension's state, folded from history. */
  readonly state: S;
};

export type ToolContext<S extends JsonValue, E extends JsonValue> = ConversationContext<S> & {
  /** The conversation's environment at its working directory, as this call found it. */
  readonly env: Environment;
  readonly signal: AbortSignal;
  /** Commit an event to history; `state` reflects it once this resolves. */
  emit(event: E): Promise<void>;
  /** Running output, for display. */
  output(chunk: string): void;
  /** Running details, for display; the last value stays with the result. */
  progress(details: JsonValue): void;
};

export type CommandContext<S extends JsonValue, E extends JsonValue> = ConversationContext<S> & {
  /** The conversation's environment at its working directory. */
  readonly env: Environment;
  /**
   * Record an event as the user. While the conversation is busy it lands at the next boundary, so the model never
   * sees it mid-turn; this resolves once it is admitted.
   */
  emit(event: E, options?: {
    /** Passive user-message context stored with the event, without starting a run. */
    content: string;
  }): Promise<void>;
  /**
   * Move the working directory, as the user, to the directory `path` names from it on the same machine. Lands like
   * `emit`; resolves with where it moves to.
   */
  cd(path: string, options?: {
    /** Add the destination as a workspace of the source directory's project, if any. */
    inheritProject?: boolean;
  }): Promise<string>;
  /** Start a manual compaction of this conversation; resolves once it has started, not once its summary is placed. */
  compact(instructions?: string): Promise<void>;
};

/** Something the user invokes, as `/name args`. */
export type Command<S extends JsonValue = JsonValue, E extends JsonValue = never> = {
  readonly name: string;
  readonly description: string;
  /** Whether the command is offered and callable in the conversation’s current environment. */
  enabled?(ctx: PromptContext): boolean | Promise<boolean>;
  /** `args` is the text after the command's name. The value is shown to the user. */
  execute(
    args: string,
    ctx: CommandContext<S, E>,
  ): JsonValue | undefined | Promise<JsonValue | undefined>;
};

/** What a handler answers with: its conversation, and the request's own abort signal. */
export type HandlerContext<S extends JsonValue> = ConversationContext<S> & {
  /** The conversation's environment at its working directory. */
  readonly env: Environment;
  /** Aborted once the frontend gives up on the request. */
  readonly signal: AbortSignal;
};

/** A request this extension's own frontend makes of it, such as for the data a view draws. */
export type Handler<
  S extends JsonValue = JsonValue,
  P extends TSchema = TSchema,
  R extends JsonValue | undefined = JsonValue | undefined,
> = {
  readonly name: string;
  /** What the frontend sends, validated with its defaults filled in. */
  readonly parameters: P;
  execute(args: Static<P>, ctx: HandlerContext<S>): R | Promise<R>;
};

/** What crosses a channel: JSON, or bytes as they are. */
export type ChannelMessage = JsonValue | Uint8Array;

/** One end of a channel between an extension's frontend and its server side. */
export interface Channel {
  /** Dropped once the channel has closed. */
  send(message: ChannelMessage): void;
  /** What the other end sends: it ends once either end closes the channel, and fails if it could not open. */
  readonly messages: ReadableStream<ChannelMessage>;
  close(): void;
}

/** What a channel opens with: its conversation, and a signal for its closing. */
export type ChannelContext<S extends JsonValue> = ConversationContext<S> & {
  /** The conversation's environment at its working directory. */
  readonly env: Environment;
  /** Aborted once the channel closes, from either end. */
  readonly signal: AbortSignal;
};

/** A long-lived connection this extension's frontend opens to it, such as a terminal's. */
export type ChannelHandler<S extends JsonValue = JsonValue, P extends TSchema = TSchema> = {
  readonly name: string;
  /** What the frontend opens it with, validated with its defaults filled in. */
  readonly parameters: P;
  /** The channel is the extension's from here, to keep and to close; a throw closes it with the message. */
  open(args: Static<P>, channel: Channel, ctx: ChannelContext<S>): void | Promise<void>;
};

/** A handler's or a channel's arguments as it receives them: validated, with their defaults filled in. */
export function handlerArguments<P extends TSchema>(
  parameters: P,
  args: JsonValue | undefined,
): Static<P> {
  const filled = Value.Default(parameters, structuredClone(args ?? {}));
  const [error] = Value.Errors(parameters, filled);
  if (error) throw new Error(`${error.instancePath.slice(1) || "arguments"}: ${error.message}`);
  return filled as Static<P>;
}

/** An editable projection of extension state, shared by new and existing conversations. */
export type Configuration<
  S extends JsonValue,
  E extends JsonValue,
  P extends TObject = TObject,
> = {
  /**
   * The questions, as a form draws them: string, number, integer and boolean properties, and unions of string
   * literals, with their `title`, `description` and `default`.
   */
  readonly parameters: P;
  values(state: S): Static<P> & Record<string, JsonValue>;
  /** Thrown errors refuse the values, with their message. */
  events(values: Static<P>, state: S): readonly E[] | Promise<readonly E[]>;
};

/** Global editable prose for this extension's main section, rendered as entered with `$USER` replaced; blank prose is omitted. */
export type PromptProse = {
  /** What the prose is unless settings override it. */
  readonly default: string;
};

/** Validate configuration values, their defaults filled in, and turn them into events without changing state. */
export async function configurationEvents<S extends JsonValue, E extends JsonValue>(
  configuration: Configuration<S, E>,
  values: JsonValue | undefined,
  state: S,
): Promise<readonly E[]> {
  const filled = Value.Default(configuration.parameters, structuredClone(values ?? {}));
  const [error] = Value.Errors(configuration.parameters, filled);
  if (error) throw new Error(`${error.instancePath.slice(1) || "values"}: ${error.message}`);
  return await configuration.events(filled as Static<TObject>, state);
}

/** What a prompt renders against. */
export type PromptContext = {
  /** The conversation's environment, unless it cannot be reached. */
  readonly env: Environment | undefined;
};

/** Text for this extension's system prompt section; `undefined` leaves it out. */
export type PromptRender = (
  ctx: PromptContext,
) => string | undefined | Promise<string | undefined>;

/** A named system prompt section, rendered against the conversation's current state. */
export type SectionRender<S extends JsonValue> = (
  ctx: ConversationContext<S> & PromptContext,
) => ReturnType<PromptRender>;

/** Inspect or replace outgoing messages without changing stored history. */
export type ContextHook = (
  messages: readonly Message[],
  model: Model<Api>,
) => readonly Message[] | undefined | Promise<readonly Message[] | undefined>;

/** Inspect or replace a provider's assembled request payload. */
export type PayloadHook = (
  payload: unknown,
  model: Model<Api>,
) => unknown | undefined | Promise<unknown | undefined>;

/** A server-wide cap on effective context windows, so compaction starts at one threshold whatever the model's window. */
export type CompactionCap = {
  /** Blocking compaction once the context reaches this many tokens; smaller model windows keep their own policy. */
  readonly maxContextTokens: number;
  /** Background compaction starts this far below the blocking threshold; `0` disables it. */
  readonly backgroundTokens: number;
};

/** The tightest of the caps extensions report, ignoring missing and invalid ones. */
export function tightestCap(
  caps: readonly (CompactionCap | undefined)[],
): CompactionCap | undefined {
  let best: CompactionCap | undefined;
  for (const cap of caps) {
    if (cap === undefined) continue;
    const maxContextTokens = Math.floor(cap.maxContextTokens);
    const backgroundTokens = Math.floor(cap.backgroundTokens);
    if (!(maxContextTokens > 0) || !(backgroundTokens >= 0)) continue;
    if (best === undefined || maxContextTokens < best.maxContextTokens) {
      best = { maxContextTokens, backgroundTokens };
    }
  }
  return best;
}

/** What an extension contributes to one conversation; rebuilt whenever its state changes. */
export type ConversationApi<S extends JsonValue, E extends JsonValue> =
  & ConversationContext<S>
  & PromptContext
  & {
    /** Register a tool for this conversation, replacing a same-named tool from this extension's setup. */
    tool<P extends TSchema, R extends JsonValue | undefined>(tool: Tool<S, E, P, R>): void;
    prompt(render: PromptRender): void;
    onContext(hook: ContextHook): void;
    onPayload(hook: PayloadHook): void;
    command(command: Command<S, E>): void;
  };

export type ExtensionApi<S extends JsonValue, E extends JsonValue> = {
  /** Register a server-wide provider for this load; builtins may be replaced and are restored on release, but another extension's provider cannot be replaced. */
  provider(provider: Provider): void;
  /** Server-only provider auth, refreshed through the same credential store as model requests. */
  auth(provider: string, signal?: AbortSignal): Promise<AuthResult | undefined>;
  tool<P extends TSchema, R extends JsonValue | undefined>(tool: Tool<S, E, P, R>): void;
  /** Answer this extension's own frontend, in conversations that have the extension. */
  handle<P extends TSchema, R extends JsonValue | undefined>(handler: Handler<S, P, R>): void;
  /** Accept channels this extension's own frontend opens, in conversations that have the extension. */
  channel<P extends TSchema>(handler: ChannelHandler<S, P>): void;
  /** A separate section keyed as `<extension>_<name>`, outside the default conversation prompt. */
  prompt(name: string, render: SectionRender<S>): void;
  /** Global editable prose prepended to this extension's main section; a later call replaces an earlier one. */
  promptProse(prose: PromptProse): void;
  /** Cap effective context windows for this load; the tightest cap among loads wins, and it lapses on release. */
  compactionCap(getter: () => CompactionCap | undefined): void;
  /** Expose editable conversation configuration; a later call replaces an earlier one. */
  configuration<P extends TObject>(configuration: Configuration<S, E, P>): void;
  /** Runs when this load of the extension is replaced or shut down. */
  onDispose(dispose: () => void | Promise<void>): void;
};

export type Extension<S extends JsonValue = JsonValue, E extends JsonValue = never> = {
  /** `codemode` is the harness's own. */
  readonly name: string;
  readonly description?: string;
  /** Available in every conversation, outside its choice of extensions. */
  readonly alwaysOn?: boolean;
  /** What a new conversation starts with, until one is chosen from before. */
  readonly defaultEnabled?: boolean;
  /**
   * Where its configuration sorts among the others': lower first, and one without it after all of them, by name.
   */
  readonly sortOrder?: number;
  /**
   * Names it went by before: the events recorded under them are its own, and conversations that chose it by one of
   * them choose it by its name now.
   */
  readonly formerly?: readonly string[];
  initial(): S;
  /** Fold one of this extension's own events. Events are kept forever: stay compatible with old shapes. */
  apply?(state: S, event: E): S;
  /** Fold any history entry, including other extensions' events. */
  observe?(state: S, entry: EntryRecord): S;
  /** Runs once per load. The returned function contributes to each conversation. */
  setup(ext: ExtensionApi<S, E>): ((conv: ConversationApi<S, E>) => void) | void;
};

export function defineExtension<S extends JsonValue, E extends JsonValue = never>(
  extension: Extension<S, E>,
): Extension<S, E> {
  return extension;
}