import { SubstrateClient } from "./client.js";
import { configuration, validName } from "./config.js";
import { Workspace } from "./workspace.js";
import type { WorkspaceBinding } from "./types.js";

export interface SchemaBuilder {
  String(options?: Record<string, unknown>): unknown;
  Number(options?: Record<string, unknown>): unknown;
  Boolean(options?: Record<string, unknown>): unknown;
  Optional(schema: unknown): unknown;
  Object(properties: Record<string, unknown>): unknown;
}
interface Context {
  sessionManager: { getSessionId(): string; getBranch(): { type: string; customType?: string; data?: unknown }[] };
  hasUI: boolean;
  ui: {
    notify(message: string, level?: "info" | "warning" | "error"): void;
    setStatus(key: string, text: string | undefined): void;
    confirm(title: string, message: string): Promise<boolean>;
  };
}
interface ToolResult {
  content: { type: "text"; text: string }[];
  details: Record<string, unknown>;
}
export interface HostAPI {
  registerFlag(name: string, options: { description: string; type: "string" }): void;
  getFlag(name: string): unknown;
  registerCommand(name: string, options: { description: string; handler(args: string, ctx: Context): Promise<void> }): void;
  registerTool(tool: {
    name: string;
    label: string;
    description: string;
    parameters: unknown;
    execute(id: string, params: Record<string, unknown>, signal: AbortSignal | undefined, onUpdate: unknown, ctx: Context): Promise<ToolResult>;
  }): void;
  on(name: string, handler: (event: Record<string, unknown>, ctx: Context) => unknown): void;
  appendEntry(type: string, data: unknown): void;
  sendMessage(message: { customType: string; content: string; display: boolean }, options?: { triggerTurn: boolean }): void;
}

const STATE = "pi-agent-substrate.binding";
const HELP = "Substrate: /substrate create <id> [template] | attach <id> | status | suspend | detach | delete <id>. Configure --substrate-endpoint or SUBSTRATE_ENDPOINT first. All substrate_* tools execute remotely; other tools remain local.";

function savedBinding(value: unknown): WorkspaceBinding | null {
  if (value === null) return null;
  if (!value || typeof value !== "object") throw new Error("Invalid saved Substrate binding; detach before attaching explicitly.");
  const v = value as Record<string, unknown>;
  if (v.version !== 1 || typeof v.sessionId !== "string" || typeof v.endpoint !== "string" || typeof v.atespace !== "string" || typeof v.environmentId !== "string" || typeof v.uncertain !== "boolean" || !Array.isArray(v.processIds) || !v.processIds.every(p => typeof p === "string")) throw new Error("Unsupported saved Substrate binding; use /substrate detach and attach explicitly.");
  validName(v.atespace, "atespace");
  validName(v.environmentId);
  return v as unknown as WorkspaceBinding;
}

export default function install(api: HostAPI, schema: SchemaBuilder): void {
  for (const [name, description] of [
    ["substrate-endpoint", "Private ate-env-api URL; HTTPS or loopback HTTP only"],
    ["substrate-atespace", "Substrate atespace (default ate-env)"],
    ["substrate-env", "Existing environment to attach explicitly on startup"],
    ["substrate-idle-seconds", "Opt-in idle suspension for a dedicated environment; 0 disables"],
  ]) api.registerFlag(name, { description, type: "string" });

  let workspace: Workspace | undefined;
  let context: Context | undefined;
  let initialization: Promise<void> = Promise.resolve();
  let generation = 0;
  const report = (message: string) => api.sendMessage({ customType: "substrate.status", content: message, display: true }, { triggerTurn: false });

  function reset(ctx: Context): void {
    generation++;
    workspace?.dispose();
    workspace = undefined;
    context = ctx;
    ctx.ui.setStatus("substrate", undefined);
  }

  function open(ctx: Context): Workspace {
    const config = configuration(name => api.getFlag(name));
    const sessionId = ctx.sessionManager.getSessionId();
    const epoch = generation;
    const client = new SubstrateClient(config.connection);
    workspace = new Workspace(client, {
      sessionId,
      endpoint: config.connection.endpoint,
      atespace: config.connection.atespace,
      idleMs: config.idleMs,
      persist: binding => { if (epoch === generation) api.appendEntry(STATE, binding); },
      notify: (message, level) => { if (epoch === generation) ctx.ui.notify(message, level); },
      changed: binding => {
        if (epoch === generation) ctx.ui.setStatus("substrate", binding ? `remote: ${binding.atespace}/${binding.environmentId}${binding.uncertain ? " (check processes)" : ""}` : undefined);
      },
    });
    return workspace;
  }

  async function initialize(ctx: Context, explicitStartup: boolean): Promise<void> {
    reset(ctx);
    // Loading the package without a configured backend is safe and network-free.
    if (!api.getFlag("substrate-endpoint") && !process.env.SUBSTRATE_ENDPOINT) return;
    const current = open(ctx);
    const entries = ctx.sessionManager.getBranch().filter(e => e.type === "custom" && e.customType === STATE);
    const last = entries.at(-1);
    if (last) current.restore(savedBinding(last.data));
    // A new branch must not inherit the parent's remote workspace implicitly.
    const requested = configuration(name => api.getFlag(name)).environmentId;
    if (explicitStartup && !last && requested) await current.attach(requested);
  }

  function scheduleInitialization(ctx: Context, explicitStartup: boolean): Promise<void> {
    const pending = initialize(ctx, explicitStartup);
    const epoch = generation;
    initialization = pending.catch(error => {
      if (epoch !== generation) return;
      workspace?.dispose();
      workspace = undefined;
      ctx.ui.notify(`Substrate initialization failed: ${String(error)}`, "error");
    });
    return initialization;
  }

  api.on("session_start", (event, ctx) => scheduleInitialization(ctx, event.reason === undefined || event.reason === "startup" || event.reason === "reload"));
  api.on("session_switch", (_event, ctx) => scheduleInitialization(ctx, false));
  api.on("session_branch", (_event, ctx) => scheduleInitialization(ctx, false));
  api.on("session_tree", (_event, ctx) => scheduleInitialization(ctx, false));
  api.on("session_shutdown", () => { generation++; workspace?.dispose(); workspace = undefined; context = undefined; });

  async function ready(ctx: Context): Promise<Workspace> {
    await initialization;
    if (context?.sessionManager.getSessionId() !== ctx.sessionManager.getSessionId()) await initialize(ctx, false);
    return workspace ?? open(ctx);
  }

  api.on("before_agent_start", async (event) => {
    await initialization;
    if (!workspace?.binding || typeof event.systemPrompt !== "string") return;
    const { atespace, environmentId } = workspace.binding;
    return { systemPrompt: `${event.systemPrompt}\n\nSubstrate workspace: ${atespace}/${environmentId}. Use substrate_read, substrate_write, substrate_shell and substrate_process for this REMOTE workspace. Paths belong to the remote environment, not the local cwd. Use substrate_shell for remote editing, search and Git. Other tools, local shell shortcuts and other extensions are NOT routed or sandboxed by this extension. Do not substitute local execution after a remote failure.` };
  });

  api.registerCommand("substrate", {
    description: "Create, attach, inspect, suspend or delete a remote Substrate workspace",
    async handler(args, ctx) {
      const [action = "help", id, template, ...extra] = args.trim().split(/\s+/).filter(Boolean);
      if (action === "help") { report(HELP); return; }
      try {
        if (extra.length || (action !== "create" && template)) throw new Error(HELP);
        if (["status", "suspend", "detach"].includes(action) && id) throw new Error(HELP);
        // Detach can also recover from invalid saved state without constructing a client.
        if (action === "detach") {
          await initialization;
          if (workspace) workspace.detach(); else api.appendEntry(STATE, null);
          report("Detached. The remote environment and any processes remain intact.");
          return;
        }
        const current = await ready(ctx);
        switch (action) {
          case "create":
            if (!id) throw new Error(HELP);
            await current.create(validName(id), template ? { name: validName(template, "template"), atespace: configuration(name => api.getFlag(name)).connection.atespace } : undefined);
            report(`Created and attached ${id}. Use substrate_* tools for remote work.`);
            break;
          case "attach":
            if (!id) throw new Error(HELP);
            await current.attach(validName(id));
            report(`Attached ${id}. Local tools remain local.`);
            break;
          case "status": report(JSON.stringify({ environment: await current.status(), binding: current.binding }, null, 2)); break;
          case "suspend":
            await current.suspend();
            report("Suspended the remote workspace. The next remote operation requests wake-up; background processes remain paused until then.");
            break;
          case "delete":
            if (!id || id !== current.binding?.environmentId) throw new Error("Specify the exact attached environment ID: /substrate delete <id>.");
            if (!ctx.hasUI) throw new Error("Deletion requires an interactive confirmation. No environment was deleted.");
            if (!await ctx.ui.confirm("Permanently delete remote workspace?", `Delete ${current.binding.atespace}/${id}, including its files, processes and snapshot?`)) return;
            await current.delete(id);
            report(`Deleted ${id}.`);
            break;
          default: throw new Error(HELP);
        }
      } catch (error) { report(`Substrate: ${error instanceof Error ? error.message : String(error)}`); }
    },
  });

  const text = (description: string) => schema.String({ description });
  const optional = (description: string) => schema.Optional(text(description));
  const result = (value: string, details: Record<string, unknown> = {}): ToolResult => ({ content: [{ type: "text", text: value }], details });
  const string = (args: Record<string, unknown>, key: string, required = true): string | undefined => {
    const value = args[key];
    if (value === undefined && !required) return undefined;
    if (typeof value !== "string" || (required && !value.length)) throw new Error(`${key} must be a nonempty string.`);
    return value;
  };

  api.registerTool({
    name: "substrate_read", label: "Substrate Read",
    description: "Read a file from the attached REMOTE workspace (1 MiB limit). Choose base64 for binary content. Does not access local files.",
    parameters: schema.Object({ path: text("Remote file path"), encoding: schema.Optional(schema.String({ enum: ["utf8", "base64"] })) }),
    async execute(_id, args, signal, _update, ctx) {
      const encoding = args.encoding ?? "utf8";
      if (encoding !== "utf8" && encoding !== "base64") throw new Error("encoding must be utf8 or base64.");
      const data = await (await ready(ctx)).read(string(args, "path")!, signal);
      const content = encoding === "base64" ? data.toString("base64") : new TextDecoder("utf-8", { fatal: true }).decode(data);
      return result(content, { bytes: data.length, encoding });
    },
  });

  api.registerTool({
    name: "substrate_write", label: "Substrate Write",
    description: "Create or overwrite a file in the attached REMOTE workspace (1 MiB decoded limit). No local writes. Use the remote shell for targeted edits.",
    parameters: schema.Object({ path: text("Remote file path"), content: schema.String({ description: "Complete file content" }), encoding: schema.Optional(schema.String({ enum: ["utf8", "base64"] })) }),
    async execute(_id, args, signal, _update, ctx) {
      const encoding = args.encoding ?? "utf8";
      if (encoding !== "utf8" && encoding !== "base64") throw new Error("encoding must be utf8 or base64.");
      if (typeof args.content !== "string") throw new Error("content must be a string.");
      if (Buffer.byteLength(args.content) > 1_398_104) throw new Error("Write exceeds the 1 MiB decoded limit. Split the file into bounded remote operations.");
      if (encoding === "base64" && !/^(?:[A-Za-z0-9+/]{4})*(?:[A-Za-z0-9+/]{2}==|[A-Za-z0-9+/]{3}=)?$/.test(args.content)) throw new Error("content is not canonical padded base64.");
      const bytes = await (await ready(ctx)).write(string(args, "path")!, Buffer.from(args.content, encoding), signal);
      return result(`Wrote ${bytes} bytes remotely.`, { bytes });
    },
  });

  api.registerTool({
    name: "substrate_shell", label: "Substrate Shell",
    description: "Execute sh -c in the attached REMOTE workspace. Use for remote editing, search, builds and Git. Set background=true for a long-running job, then inspect/stop it with substrate_process. Foreground output is capped at 1 MiB; overflow or cancellation triggers termination, not local fallback.",
    parameters: schema.Object({ command: text("Remote shell command"), cwd: optional("Remote working directory"), background: schema.Optional(schema.Boolean()), timeoutSeconds: schema.Optional(schema.Number({ minimum: 1, maximum: 120 })) }),
    async execute(_id, args, signal, _update, ctx) {
      const request = { command: ["sh", "-c", string(args, "command")!], cwd: string(args, "cwd", false) };
      if (args.background !== undefined && typeof args.background !== "boolean") throw new Error("background must be boolean.");
      const seconds = args.timeoutSeconds ?? 120;
      if (typeof seconds !== "number" || !Number.isFinite(seconds) || seconds < 1 || seconds > 120) throw new Error("timeoutSeconds must be between 1 and 120. Use background for longer jobs.");
      const deadline = AbortSignal.timeout(Math.ceil(seconds * 1000));
      const bounded = signal ? AbortSignal.any([signal, deadline]) : deadline;
      const current = await ready(ctx);
      if (args.background) {
        const processId = await current.start(request, bounded);
        return result(`Started remote process ${processId}. Inspect or kill it using substrate_process. Automatic suspension is blocked until it is confirmed finished.`, { processId });
      }
      const { process, output } = await current.shell(request, bounded);
      const outputText = `${output.stdout}${output.stderr ? `\n[stderr]\n${output.stderr}` : ""}\n[exit ${process.exitCode}]`;
      if (process.exitCode !== 0) throw new Error(outputText);
      return result(outputText, { process, ...output });
    },
  });

  api.registerTool({
    name: "substrate_process", label: "Substrate Process",
    description: "Inspect, read logs or terminate a tracked process in the attached REMOTE environment. Logs return byte offsets for paging; status reconciles completed jobs for idle suspension.",
    parameters: schema.Object({ processId: text("Remote process ID"), action: schema.String({ enum: ["status", "logs", "kill"] }), stdoutOffset: schema.Optional(schema.Number({ minimum: 0 })), stderrOffset: schema.Optional(schema.Number({ minimum: 0 })) }),
    async execute(_id, args, signal, _update, ctx) {
      const action = args.action;
      if (action !== "status" && action !== "logs" && action !== "kill") throw new Error("action must be status, logs or kill.");
      for (const name of ["stdoutOffset", "stderrOffset"] as const) if (args[name] !== undefined && (!Number.isSafeInteger(args[name]) || Number(args[name]) < 0)) throw new Error(`${name} must be a nonnegative safe integer.`);
      const value = await (await ready(ctx)).process(string(args, "processId")!, action, signal, { stdoutOffset: args.stdoutOffset as number | undefined, stderrOffset: args.stderrOffset as number | undefined });
      return result(JSON.stringify(value, null, 2));
    },
  });
}