import { t } from "structural";
import type { Transport } from "../transports/transport-common.ts";
import { Result, ok } from "./result.ts";
import type { BuiltinIRRole } from "./ir-roles.ts";
import type { Content } from "./llm-ir.ts";

/*
 * Tool definitions
 * -------------------------------------------------------------------------------------------------
 *
 * Tools have to define a few things:
 *
 * 1. A tool name.
 * 2. A caller-facing description of what the tool does.
 * 3. A schema for the input arguments from the LLM.
 * 4. A schema that we parse the input into. If we don't need parse the LLM input, this can be
 *    autogenerated.
 * 5. A parse function that parses the input into the parsed form. If we don't need to parse the LLM
 *    input, this can be autogenerated.
 * 6. A validate function that validates the parsed data at runtime, e.g. "does this file still have
 *    this search string even though we performed two edits previously on it"
 * 7. A run function, which takes the tool data and runs it, returning some result
 */
export type ToolDef<
  Data,
  Name extends string,
  Arguments,
  Parsed,
  SubagentNames extends string,
  Extra,
> = {
  name: Name;
  description: string;
  ArgumentsSchema: t.Type<Arguments>;
  ParsedSchema: t.Type<Parsed>;
  parse: (
    args: ToolParseArgs<Arguments>,
  ) => Promise<Result<ParseResult<Arguments, Parsed>, string>>;
  validate: (
    abortSignal: AbortSignal,
    transport: Transport,
    t: {
      original: Schema<Name, Arguments>;
      parsed: Schema<Name, Parsed>;
    },
    data: Data,
  ) => Promise<Result<null, string>>;
  run: (
    args: ToolRunArgs<Data, Name, Arguments, Parsed>,
  ) => Promise<Result<ToolReturn<SubagentNames, Extra>, string>>;
};

export type ToolParseArgs<Arguments> = {
  signal: AbortSignal;
  transport: Transport;
  original: Arguments;
};

export type ToolRunArgs<Data, Name extends string, Arguments, Parsed> = {
  signal: AbortSignal;
  transport: Transport;
  toolCall: {
    toolCallId: string;
    original: Schema<Name, Arguments>;
    parsed: Schema<Name, Parsed>;
  };
  data: Data;
};

type ToolImplementationRunArgs<
  Data,
  Name extends string,
  Arguments,
  Parsed,
  CustomIR,
> = ToolRunArgs<Data, Name, Arguments, Parsed> & {
  customIR: CustomIR;
};

export type ParseResult<Arguments, Parsed> = {
  original: Arguments;
  parsed: Parsed;
};

/*
 * Basic schema for tool calling
 */
type Schema<Name extends string, Arguments> = {
  name: Name;
  arguments: Arguments;
};

/*
 * Tools can return:
 *
 * - output content
 * - subagent invocations (for their subagent dependencies)
 * - any custom LLM IR defined
 */
export type ToolReturn<SubagentName extends string, Extra> =
  | (Content & {
      type: "output";

      // The line count to show in the UI, if it's not just the number of lines in the content
      lines?: number;
    })
  | {
      type: "invoke-subagent";
      name: SubagentName;
    }
  | {
      type: "custom-ir";
      data: Extra;
    };

type RawToolFactory<
  Data,
  S extends AnyToolSchema,
  Parsed,
  SubagentNames extends string,
  Extra,
> = (args: {
  signal: AbortSignal;
  transport: Transport;
  data: Data;
}) => Promise<ToolDef<Data, S["name"], S["arguments"], Parsed, SubagentNames, Extra> | null>;

type SimpleRawToolFactory<
  Data,
  S extends AnyToolSchema,
  Parsed,
  SubagentNames extends string,
  Extra,
  CustomIR,
> = (args: { signal: AbortSignal; transport: Transport; data: Data }) => Promise<
  | (Omit<
      ToolDef<Data, S["name"], S["arguments"], Parsed, SubagentNames, Extra>,
      "ArgumentsSchema" | "ParsedSchema" | "description" | "name" | "validate" | "run"
    > & {
      run: (
        args: ToolImplementationRunArgs<Data, S["name"], S["arguments"], Parsed, CustomIR>,
      ) => Promise<Result<ToolReturn<SubagentNames, Extra>, string>>;
    } & Partial<
        Pick<ToolDef<Data, S["name"], S["arguments"], Parsed, SubagentNames, Extra>, "validate">
      >)
  | null
>;

type SimpleAutoParsedRawToolFactory<
  Data,
  S extends AnyToolSchema,
  SubagentNames extends string,
  Extra,
  CustomIR,
> = (args: { signal: AbortSignal; transport: Transport; data: Data }) => Promise<
  | (Omit<
      ToolDef<Data, S["name"], S["arguments"], S["arguments"], SubagentNames, Extra>,
      "ArgumentsSchema" | "ParsedSchema" | "description" | "name" | "parse" | "validate" | "run"
    > & {
      run: (
        args: ToolImplementationRunArgs<Data, S["name"], S["arguments"], S["arguments"], CustomIR>,
      ) => Promise<Result<ToolReturn<SubagentNames, Extra>, string>>;
    } & Partial<
        Pick<
          ToolDef<Data, S["name"], S["arguments"], S["arguments"], SubagentNames, Extra>,
          "validate"
        >
      >)
  | null
>;

type ToolFactoryArgs<Data> = {
  signal: AbortSignal;
  transport: Transport;
  data: Data;
};

type DynamicToolFactoryReturn<T extends (...args: any) => any> = ReturnType<
  Extract<Exclude<Awaited<ReturnType<T>>, null>, (...args: any) => any>
>;

type DynamicToolFactorySubagents<T extends (...args: any) => any> = ToolSubagentNames<
  Exclude<Awaited<ReturnType<T>>, null>
>;
type DynamicToolFactoryExtra<T extends (...args: any) => any> = ToolFactoryExtra<
  Exclude<Awaited<ReturnType<T>>, null>
>;

type RuntimeToolDefinition<Extra> = Omit<
  ToolDef<unknown, string, unknown, unknown, string, Extra>,
  "ArgumentsSchema" | "ParsedSchema" | "description" | "name" | "parse" | "validate" | "run"
> & {
  run: (
    args: ToolImplementationRunArgs<unknown, string, unknown, unknown, any>,
  ) => Promise<Result<ToolReturn<string, Extra>, string>>;
} & Partial<Pick<ToolDef<unknown, string, unknown, unknown, string, Extra>, "parse" | "validate">>;

type AnyToolSchema = Schema<any, any>;
type AnyParsedArguments = any;
type AnySubagentNames = any;
type AnyCustomIR = any;
type AnyToolData = any;

/*
 * Tool extension IRs are, by definition, tool-output-shaped IRs: they exist only as the result
 * of running a tool call, so they must carry the tool call they answer. Constraining extension
 * IRs to this shape lets libocto treat every extension IR as a tool output without clients
 * re-deriving that decision.
 */
export type ToolExtensionIR<Role extends string> = Role extends BuiltinIRRole
  ? never
  : {
      role: Role;
      toolCall: {
        toolCallId: string;
      };
    };

// Public tool factories are the precise raw factory type plus the capability brand. Keeping the raw
// factory in the intersection preserves literal tool names and parsed argument types.
export type ToolFactory<
  Data,
  S extends AnyToolSchema,
  Parsed,
  SubagentNames extends string,
  Extra extends ToolExtensionIR<any>,
> = RawToolFactory<Data, S, Parsed, SubagentNames, Extra> &
  ToolFactoryRequirements<SubagentNames, Extra>;

export class ToolBuilder<Data = unknown> {
  withData<NewData>(): ToolBuilder<NewData> {
    return new ToolBuilder<NewData>();
  }

  /*
   * Caller-facing tool declaration, explicit parse form.
   *
   * Use this overload when the LLM's raw arguments are not the shape the tool wants to run with.
   * define(...) must then supply parse({ original, ... }) to produce ParsedSchema.
   */
  declare<
    Name extends string,
    Arguments,
    Parsed,
    const SubagentNames extends string = never,
  >(partial: {
    name: Name;
    description: string;
    ArgumentsSchema: t.Type<Arguments>;
    ParsedSchema: t.Type<Parsed>;
    subagents?: readonly SubagentNames[];
  }): DeclaredTool<Data, Name, Arguments, Parsed, SubagentNames, never, true, {}>;

  /*
   * Caller-facing tool declaration, auto-parse form.
   *
   * Use this for the common case where the LLM's raw arguments are already the shape the tool runs
   * with. The caller omits ParsedSchema and parse(...). The resulting tool behaves as if
   * ParsedSchema === ArgumentsSchema and parse(...) returned { original: x, parsed: x }.
   */
  declare<Name extends string, Arguments, const SubagentNames extends string = never>(partial: {
    name: Name;
    description: string;
    ArgumentsSchema: t.Type<Arguments>;
    ParsedSchema?: undefined;
    subagents?: readonly SubagentNames[];
  }): DeclaredTool<Data, Name, Arguments, Arguments, SubagentNames, never, false, {}>;

  declare(partial: {
    name: string;
    description: string;
    ArgumentsSchema: t.Type<any>;
    ParsedSchema?: t.Type<any>;
    subagents?: readonly string[];
  }): any {
    return new DeclaredTool(partial);
  }

  /*
   * Caller-facing dynamic tool builder.
   *
   * Use this when the exact tool shape is not known until load time, for example when the schema
   * depends on data, transport capabilities, or discovered project state. The caller returns a
   * normal declare(...).define(...) result from the selector rather than repeating a second dynamic
   * API for name/schema/subagents/parse behavior.
   *
   * The selected tool carries the same type information as a static tool. This wrapper preserves
   * that information and forwards the runtime args into the selected factory.
   */
  dynamicDefineTool<
    T extends (
      args: ToolFactoryArgs<Data>,
    ) => Promise<ToolFactory<
      Data,
      AnyToolSchema,
      AnyParsedArguments,
      AnySubagentNames,
      AnyCustomIR
    > | null>,
  >(
    factory: T,
  ): ((args: ToolFactoryArgs<Data>) => DynamicToolFactoryReturn<T>) &
    ToolFactoryRequirements<DynamicToolFactorySubagents<T>, DynamicToolFactoryExtra<T>> {
    const selectTool = factory as unknown as (args: ToolFactoryArgs<Data>) => ReturnType<T>;

    const wrapped = async (args: ToolFactoryArgs<Data>) => {
      const selectedTool = await selectTool(args);
      const runSelectedTool = selectedTool as
        | null
        | ((args: ToolFactoryArgs<Data>) => DynamicToolFactoryReturn<T>);

      // One-level flatMap: the selector returns a tool factory; dynamicDefineTool must itself
      // behave like a tool factory that returns the selected tool definition.
      return runSelectedTool?.(args) ?? null;
    };

    // The brand has no runtime representation; it only makes defineAgent reject tools whose custom
    // IR output is not included in the target IR universe.
    return wrapped as unknown as ((args: ToolFactoryArgs<Data>) => DynamicToolFactoryReturn<T>) &
      ToolFactoryRequirements<DynamicToolFactorySubagents<T>, DynamicToolFactoryExtra<T>>;
  }
}

export class DeclaredTool<
  Data,
  Name extends string,
  Arguments,
  Parsed,
  SubagentNames extends string,
  Extra extends ToolExtensionIR<any>,
  RequiresParse extends boolean,
  CustomIR,
> {
  constructor(
    private readonly partial: {
      name: Name;
      description: string;
      ArgumentsSchema: t.Type<Arguments>;
      ParsedSchema?: t.Type<Parsed>;
      subagents?: readonly SubagentNames[];
    },
    private readonly customIRBuilders: CustomIRBuilderMap | null = null,
  ) {}

  withCustomIR<Builders extends CustomIRBuilderMap>(
    builders: Builders,
  ): DeclaredTool<
    Data,
    Name,
    Arguments,
    Parsed,
    SubagentNames,
    CustomIRDataForTool<Data, Name, Arguments, Parsed, SubagentNames, Builders>,
    RequiresParse,
    BoundCustomIRForTool<Data, Name, Arguments, Parsed, SubagentNames, Builders>
  > {
    return new DeclaredTool(this.partial, builders);
  }

  define<
    T extends SimpleRawToolFactory<
      Data,
      Schema<Name, Arguments>,
      Parsed,
      SubagentNames,
      Extra,
      CustomIR
    >,
  >(
    this: DeclaredTool<Data, Name, Arguments, Parsed, SubagentNames, Extra, true, CustomIR>,
    factory: (args: Parameters<T>[0]) => Promise<Awaited<ReturnType<T>>>,
  ): ToolFactory<Data, Schema<Name, Arguments>, Parsed, SubagentNames, Extra>;

  define<
    T extends SimpleAutoParsedRawToolFactory<
      Data,
      Schema<Name, Arguments>,
      SubagentNames,
      Extra,
      CustomIR
    >,
  >(
    this: DeclaredTool<Data, Name, Arguments, Arguments, SubagentNames, Extra, false, CustomIR>,
    factory: (args: Parameters<T>[0]) => Promise<Awaited<ReturnType<T>>>,
  ): ToolFactory<Data, Schema<Name, Arguments>, Arguments, SubagentNames, Extra>;

  define(
    factory: (args: ToolFactoryArgs<Data>) => Promise<RuntimeToolDefinition<Extra> | null>,
  ): any {
    const partial = this.partial;
    const customIRBuilders = this.customIRBuilders;
    const wrapped = async (args: ToolFactoryArgs<Data>) => {
      const def = await factory(args);
      if (def === null) return null;

      return {
        ...def,
        name: partial.name,
        description: partial.description,
        ArgumentsSchema: partial.ArgumentsSchema,
        ParsedSchema: partial.ParsedSchema ?? partial.ArgumentsSchema,
        parse:
          def.parse ??
          (async ({ original }) => ({
            success: true as const,
            data: {
              original,
              parsed: original,
            },
          })),
        validate: def.validate ?? (async () => ({ success: true, data: null })),
        run: async (runArgs: ToolRunArgs<Data, Name, Arguments, Parsed>) =>
          def.run({
            ...runArgs,
            customIR: bindCustomIR(customIRBuilders, flattenToolCall(runArgs.toolCall)),
          }),
      };
    };

    // The brand has no runtime representation; it only makes defineAgent reject tools whose custom
    // IR output is not included in the target IR universe.
    return wrapped as unknown as ToolFactory<
      Data,
      Schema<Name, Arguments>,
      Parsed,
      SubagentNames,
      Extra
    >;
  }
}

// A shared tool builder for callers that do not need a dedicated builder instance.
export const TOOL_BUILDER = new ToolBuilder();

export type ToolMap<SubagentNames extends string, Extra extends ToolExtensionIR<any>> = {
  [key: string]: ToolFactory<AnyToolData, AnyToolSchema, AnyParsedArguments, SubagentNames, Extra>;
};

type CustomIRBuilderMap = Record<string, (toolCall: any) => (args: any) => ToolExtensionIR<any>>;

type DeclaredToolMap<
  Data,
  Name extends string,
  Arguments,
  Parsed,
  SubagentNames extends string,
  Extra extends ToolExtensionIR<any>,
> = {
  [K in Name]: ToolFactory<Data, { name: K; arguments: Arguments }, Parsed, SubagentNames, Extra>;
};

type DeclaredToolCall<
  Data,
  Name extends string,
  Arguments,
  Parsed,
  SubagentNames extends string,
> = ToolCall<DeclaredToolMap<Data, Name, Arguments, Parsed, SubagentNames, never>>;

type CustomIRData<Builders extends CustomIRBuilderMap, Call> = {
  [K in keyof Builders]: Builders[K] extends (toolCall: Call) => (args: any) => infer IR
    ? IR extends ToolExtensionIR<any>
      ? IR
      : never
    : never;
}[keyof Builders];

type CustomIRDataForTool<
  Data,
  Name extends string,
  Arguments,
  Parsed,
  SubagentNames extends string,
  Builders extends CustomIRBuilderMap,
> = CustomIRData<Builders, DeclaredToolCall<Data, Name, Arguments, Parsed, SubagentNames>>;

type BoundCustomIR<Builders extends CustomIRBuilderMap, Call> = {
  [K in keyof Builders]: Builders[K] extends (toolCall: Call) => (args: infer Args) => infer IR
    ? (args: Args) => Result<{ type: "custom-ir"; data: IR }, string>
    : never;
};

type BoundCustomIRForTool<
  Data,
  Name extends string,
  Arguments,
  Parsed,
  SubagentNames extends string,
  Builders extends CustomIRBuilderMap,
> = BoundCustomIR<Builders, DeclaredToolCall<Data, Name, Arguments, Parsed, SubagentNames>>;

export type LoadedTools<T extends ToolMap<any, any>> = {
  [K in keyof T]: Exclude<Awaited<ReturnType<T[K]>>, null>;
};
export type ToolCall<T extends ToolMap<any, any>> = {
  [K in keyof LoadedTools<T>]: {
    type: "tool-call";
    name: LoadedTools<T>[K]["name"];
    toolCallId: string;
    parsed: t.GetType<LoadedTools<T>[K]["ParsedSchema"]>;
    original: t.GetType<LoadedTools<T>[K]["ArgumentsSchema"]>;
  };
}[keyof LoadedTools<T>];

function flattenToolCall<Name extends string, Arguments, Parsed>(toolCall: {
  toolCallId: string;
  original: Schema<Name, Arguments>;
  parsed: Schema<Name, Parsed>;
}) {
  return {
    type: "tool-call" as const,
    name: toolCall.parsed.name,
    toolCallId: toolCall.toolCallId,
    original: toolCall.original.arguments,
    parsed: toolCall.parsed.arguments,
  };
}

function bindCustomIR(builders: CustomIRBuilderMap | null, toolCall: unknown) {
  const bound: Record<
    string,
    (args: unknown) => Result<{ type: "custom-ir"; data: unknown }, string>
  > = {};
  for (const [name, builder] of Object.entries(builders ?? {})) {
    const build = builder(toolCall);
    bound[name] = args => ok({ type: "custom-ir", data: build(args) });
  }
  return bound;
}

/*
 * ALERT ALERT
 *
 * Type system bullshit used for branding and validating IRs with respect to tools and subagent
 * dependencies.
 */

// Tool factories are structurally just functions, so TypeScript can otherwise forget which IR
// extension set they were created for if the implementation does not currently return custom IR.
// These unique-symbol fields are a compile-time-only brand attached by declare(...).define(...).
declare const toolFactoryExtra: unique symbol;
declare const toolFactorySubagents: unique symbol;

// The brand should be covariant: a base tool that emits no custom IR (never) is valid in a richer
// IR context, but a richer tool is not valid in a base IR context.
type Covariant<T> = () => T;

// Required phantom fields keep Extra and SubagentNames visible when checking assignability to a
// ToolMap. They are never read or written at runtime.
export type ToolFactoryRequirements<
  SubagentNames extends string,
  Extra extends ToolExtensionIR<any>,
> = {
  readonly [toolFactoryExtra]: Covariant<Extra>;
  readonly [toolFactorySubagents]: Covariant<SubagentNames>;
};

type ToolFactoryExtra<T> = T extends ToolFactoryRequirements<any, infer Extra> ? Extra : never;

export type ToolSubagentNames<T> =
  T extends ToolFactoryRequirements<infer SubagentNames, any> ? SubagentNames : never;