import { ImageInfo } from "../utils/image-utils.ts";
import { BUILTIN_IR_ROLES } from "./ir-roles.ts";
import type { BuiltinIRRole } from "./ir-roles.ts";
import type { CompilerUsage } from "./compilers/compiler-interface.ts";
import type {
  ToolCall,
  ToolExtensionIR,
  ToolFactoryRequirements,
  ToolMap,
  ToolSubagentNames,
} from "./tool-def.ts";

/*
 * LLM IR
 * -------------------------------------------------------------------------------------------------
 *
 * This defines a set of base IRs, which can be extended by callers using an Extra type.
 *
 * LLM compilers only accept the unextended base IRs. However, callers can use additional IR types
 * and convert them in a pre-compile pass to the lowered types, and tools can declare specific
 * IR extension requirements and can return those extended IRs, to help track richer information
 * that might be useful for pre-compile optimization passes.
 */

/*
 * IRs are all defined in terms of agents. Agents specify what tools they use, what subagents they
 * have, and what extended IRs they use.
 */
export type Agent<
  Extra extends ToolExtensionIR<any>,
  SubagentDirectory extends AgentDirectory,
  Tools extends ToolMap<Extract<keyof SubagentDirectory, string>, Extra>,
> = {
  tools: Tools;
  agents: SubagentDirectory;
};

// A named directory of agents
export type AgentDirectory = {
  [name: string]: Agent<any, any, any>;
};

// Helper function to define agents with compile-time safety guarantees. It's an identity function
// that runs compile-time validation on the passed-in agent and assigns the correct type + branding.
export function defineAgent<A extends { tools: ToolMap<any, any>; agents: AgentDirectory }>(
  a: A & ValidateAgentSubagents<A>,
): A {
  return a;
}

// An IR that defines sub-agent trajectories
export class AgentTrajectory<T extends AgentDirectory, Name extends keyof T> {
  readonly role = "trajectory";
  readonly ir: LlmIR<{ agents: T[Name]["agents"]; tools: T[Name]["tools"] }>[];
  private readonly name: Name;

  constructor(name: Name, ir: LlmIR<{ agents: T[Name]["agents"]; tools: T[Name]["tools"] }>[]) {
    this.name = name;
    this.ir = ir;
  }

  // A type guard to check whether an AgentTrajectory belongs to a specific named subagent. This is
  // useful since different subagents may have different tools, so you can narrow which tools an
  // if-statement needs to check for by first checking which subagent you're dealing with.
  // For example:
  //
  // if(agentIr.isNamed("research")) {
  //   // Only tools and subagents that the research subagent has access to are accessible here
  // }
  isNamed<N extends Name>(name: N): this is AgentTrajectory<T, N> {
    if (this.name === name) return true;
    return false;
  }

  // A helpful function for checking *all* possible subagent names, and narrowing each one down in a
  // callback. For example, if you had "explore" and "research" subagents, you might call:
  //
  // ir.cond({
  //   explore: exploreIr => {
  //     // exploreIr is guaranteed to be an explore subagent trajectory
  //     // any tools and subagents are narrowed to only what's accessible to the explore subagent
  //   },
  //   research: researchIr => {
  //     // researchIr is guaranteed to be a research subagent trajectory
  //     // any tools and subagents are narrowed to only what's accessible to the research subagent
  //   },
  // });
  //
  // Like Lisp-like `cond` expressions, `cond` returns whatever the cond arms return. It can take
  // either synchronous handlers, or async handlers. If it takes sync handlers, it returns a
  // non-promise; if it takes async handlers, it returns a promise that you can await.
  //
  // For example:
  //
  // const output = await ir.cond({
  //   explore: async (exploreIr) => {
  //     return someOutput;
  //   },
  //   research: async (researchIr) => {
  //     return someOtherOutput;
  //   },
  // });
  //
  // Note that the handlers must be all-async, or all-sync; you can't mix sync and async.
  cond<C extends CondHandlerMap<T, Name>>(
    conditions: C & CondHandlerAsyncValidation<C>,
  ): CondReturn<C> {
    for (const [k, v] of Object.entries(conditions) as Array<
      [keyof C, (self: AgentTrajectory<T, Name>) => unknown]
    >) {
      if (k === this.name) {
        return v(this) as CondReturn<C>;
      }
    }
    throw new Error("Impossible");
  }
}
type CondHandlerMap<T extends AgentDirectory, Name extends keyof T> = {
  [K in Name]: (self: AgentTrajectory<T, K>) => unknown;
};
type CondHandlerReturn<C> = C[keyof C] extends (...args: any) => infer Ret ? Ret : never;
declare const mixedCondHandlerReturns: unique symbol;

// cond(...) preserves sync handlers as sync returns and async handlers as Promise returns. A plain
// conditional return type is not enough, because TypeScript can infer a mixed table as
// `string | Promise<string>`. This parameter-side validation rejects handler maps where some
// branches return PromiseLike values and others return non-Promise values.
type CondHandlerAsyncValidation<C> =
  Extract<CondHandlerReturn<C>, PromiseLike<unknown>> extends never
    ? unknown
    : Exclude<CondHandlerReturn<C>, PromiseLike<unknown>> extends never
      ? unknown
      : { readonly [mixedCondHandlerReturns]: never };

// Once mixed sync/async tables are rejected, cond(...) can return exactly what callers expect:
// sync tables return their handler value directly, async tables return one Promise for the awaited
// handler value union.
type CondReturn<C> =
  Extract<CondHandlerReturn<C>, PromiseLike<unknown>> extends never
    ? CondHandlerReturn<C>
    : Promise<Awaited<CondHandlerReturn<C>>>;

export type MalformedToolRequest = {
  type: "malformed-tool-request";
  error: string;
  call: {
    original: {
      name: string;
      arguments: any;
    };
  };
  toolCallId: string;
};

export type AnthropicAssistantData = {
  thinkingBlocks: Array<
    | {
        type: "thinking";
        thinking: string;
        signature: string;
      }
    | {
        type: "redacted_thinking";
        data: string;
      }
  >;
};

export type Content = {
  content: Array<
    | {
        type: "text";
        content: string;
      }
    | {
        type: "image";
        image: ImageInfo;
      }
  >;
};

export type Checkpoint = Content & {
  role: "checkpoint";
};

export type LoweredCheckpoint = Content & {
  role: "lowered-checkpoint";
};

export type AssistantMessage<T extends ToolMap<any, any>> = {
  role: "assistant";
  content: string;
  reasoningContent?: string | null;
  openai?: {
    encryptedReasoningContent?: string | null;
    reasoningId?: string;
  };
  anthropic?: AnthropicAssistantData;
  toolCalls?: Array<ToolCall<T> | MalformedToolRequest>;
  usage: CompilerUsage;
};

export type UserMessage = Content & {
  role: "user";
};

export type ToolOutputMessage<T extends ToolMap<any, any>> = Content & {
  role: "tool-output";
  toolCall: ToolCall<T>;
};

export type ToolRuntimeErrorMessage<T extends ToolMap<any, any>> = {
  role: "tool-runtime-error";
  toolCall: ToolCall<T>;
  error: string;
};

export type ToolValidationErrorMessage<T extends ToolMap<any, any>> = {
  role: "tool-validation-error";
  toolCall: ToolCall<T>;
  error: string;

  // TODO: remove this, if the validation is aborted treat it like an assistant message abort
  aborted: boolean;
};

export type ToolParseErrorMessage = {
  role: "tool-parse-error";
  malformedRequest: MalformedToolRequest;
};

export type ToolSkipOutputMessage<T extends ToolMap<any, any>> = {
  role: "tool-skip-output";
  toolCall: ToolCall<T>;
  reason: string;
};

export type ToolSubagentInvoke<T extends ToolMap<any, any>, SubagentName extends string> = {
  role: "tool-invoke-subagent";
  toolCall: ToolCall<T>;
  subagent: SubagentName;
};

/*
 * All compiler-ready base IR types, with no extension IR types and no subagent trajectories.
 *
 * Raw Checkpoint IR is intentionally not part of LoweredIR. Checkpoints only make sense before the
 * final lowering pass, because lower(...) must first discard everything before the most recent
 * checkpoint. To make that easy to enforce at compile time, callers target CheckpointedIR with raw
 * Checkpoints, and lower(...) converts the surviving checkpoint to LoweredCheckpoint before any
 * compiler can see it.
 */
export type LoweredIR<T extends ToolMap<any, any>> =
  | AssistantMessage<T>
  | UserMessage
  | ToolOutputMessage<T>
  | ToolRuntimeErrorMessage<T>
  | ToolValidationErrorMessage<T>
  | ToolParseErrorMessage
  | ToolSkipOutputMessage<T>
  | LoweredCheckpoint;

/*
 * LoweredIR with pre-compiler checkpoints.
 *
 * This is the shape user-space lowering passes should target after converting custom extension IRs,
 * but before calling libocto's final lower(...). It is identical to LoweredIR except that it carries
 * raw Checkpoints instead of LoweredCheckpoints, forcing the final checkpoint slicing/conversion pass
 * to happen before compiler use.
 */
export type CheckpointedIR<T extends ToolMap<any, any>> =
  | Exclude<LoweredIR<T>, LoweredCheckpoint>
  | Checkpoint;

/*
 * Compiler-ready IR plus subagent trajectories.
 *
 * This is the shape callers should produce after lowering their custom IR extensions, but before
 * deciding how to represent nested subagent trajectories for a concrete compiler.
 */
export type LoweredIRWithTrajectories<A extends Agent<any, any, any>> =
  | LoweredIR<A["tools"]>
  | AgentTrajectory<A["agents"], keyof A["agents"]>;

export type CheckpointedIRWithTrajectories<A extends Agent<any, any, any>> =
  | CheckpointedIR<A["tools"]>
  | AgentTrajectory<A["agents"], keyof A["agents"]>;

/*
 * All IR types including extensions.
 *
 * Allows passing in arbitrary extra IR types via the Agent's tool map. Useful for IR types
 * that not all clients might use, e.g. file IO types which may have prompt optimizations.
 */
type ToolExtra<T> = T extends ToolFactoryRequirements<any, infer Extra> ? Extra : never;
type AgentExtra<A extends Agent<any, any, any>> = ToolExtra<A["tools"][keyof A["tools"]]>;

export type LlmIR<A extends Agent<any, any, any>> =
  | CheckpointedIRWithTrajectories<A>
  | AgentExtra<A>;

/*
 * Returns the tool call ID that an IR answers, or null if the IR is not tool-output-shaped.
 *
 * The parameter is expressed in terms of the genuinely generic IR types rather than a single
 * "all built-in IRs" union type: AgentTrajectory is invariant in its agent directory, so no
 * monomorphic AgentTrajectory instantiation (even AgentTrajectory<any, any>) accepts every
 * trajectory — the type parameters must be quantified at the function level.
 *
 * Every tool extension IR (see ToolExtensionIR) carries the tool call it answers by definition,
 * so non-built-in IRs are answered unconditionally. The built-in shapes are switched
 * exhaustively: adding a new built-in IR role breaks compilation of the default branch,
 * forcing an explicit decision here.
 */
function isBuiltinRole(role: string): role is BuiltinIRRole {
  return role in BUILTIN_IR_ROLES;
}

function isBuiltinIR<Role extends string, T extends AgentDirectory, Name extends keyof T>(
  ir: LoweredIR<any> | Checkpoint | AgentTrajectory<T, Name> | ToolExtensionIR<Role>,
): ir is LoweredIR<any> | Checkpoint | AgentTrajectory<T, Name> {
  return isBuiltinRole(ir.role);
}

export function answeredToolCallId<
  Role extends string,
  T extends AgentDirectory,
  Name extends keyof T,
>(
  ir: LoweredIR<any> | Checkpoint | AgentTrajectory<T, Name> | ToolExtensionIR<Role>,
): string | null {
  if (!isBuiltinIR(ir)) return ir.toolCall.toolCallId;
  switch (ir.role) {
    case "tool-parse-error":
      return ir.malformedRequest.toolCallId;
    case "assistant":
    case "user":
    case "checkpoint":
    case "lowered-checkpoint":
    case "trajectory":
      return null;
    case "tool-output":
    case "tool-runtime-error":
    case "tool-validation-error":
    case "tool-skip-output":
      return ir.toolCall.toolCallId;
    default: {
      const _exhaustive: never = ir;
      return _exhaustive;
    }
  }
}

type AssertNever<T extends never> = T;
// Keeps BUILTIN_IR_ROLES in sync with the roles of the built-in IR shapes handled above.
// Indexed access (rather than assignability) keeps this insensitive to AgentTrajectory's
// invariance in its agent directory.
type _BuiltinIRRolesMatch = AssertNever<
  | Exclude<
      LoweredIR<any>["role"] | Checkpoint["role"] | AgentTrajectory<any, any>["role"],
      BuiltinIRRole
    >
  | Exclude<
      BuiltinIRRole,
      LoweredIR<any>["role"] | Checkpoint["role"] | AgentTrajectory<any, any>["role"]
    >
>;

/*
 * Agent dependency compile-time validation/branding
 * -------------------------------------------------------------------------------------------------
 *
 * Checks that for a given agent with subagents, all tools that declare subagent dependencies have
 * those dependencies satisfied. For example, if a `research` tool expects to be able to invoke a
 * `research` subagent, and you use the research tool but don't define a research subagent, you'll
 * get a compile error.
 */
declare const missingToolSubagents: unique symbol;

type RequiredToolSubagentNames<Tools> = ToolSubagentNames<Tools[keyof Tools]>;

// Type used to validate that the given agent's tools have all of their subagent dependencies
// fulfilled. Recursively narrows the given type until it's either {} (all dependencies are
// satisfied), or an impossible-to-construct type via the non-existent missingToolSubagents unique
// symbol, which will cause a useful compile error.
type ValidateAgentSubagents<A> = A extends { tools: infer Tools; agents: infer SubagentDirectory }
  ? Exclude<
      RequiredToolSubagentNames<Tools>,
      Extract<keyof SubagentDirectory, string>
    > extends never
    ? {
        agents: { [K in keyof SubagentDirectory]: ValidateAgentSubagents<SubagentDirectory[K]> };
      }
    : {
        readonly [missingToolSubagents]: Exclude<
          RequiredToolSubagentNames<Tools>,
          Extract<keyof SubagentDirectory, string>
        >;
      }
  : never;