import { Mob, parseMobStatusResult } from './mob.js';
import { Session } from './session.js';
import type {
  RuntimeConfig,
  InitResult,
  SessionConfig,
  MobDefinition,
  JsonSchema,
  ToolCallback,
  MobStatus,
} from './types.js';
import type { MobListResult as WireMobListResult } from './generated/mob.js';

/** Expected WASM runtime version — must match the compiled binary. */
const EXPECTED_VERSION = '0.6.21';

/**
 * Convert a camelCase config object to the snake_case expected by WASM.
 *
 * Only converts known top-level keys to avoid mangling nested data.
 */
function toWasmConfig(config: RuntimeConfig): Record<string, unknown> {
  return {
    api_key: config.apiKey,
    anthropic_api_key: config.anthropicApiKey,
    openai_api_key: config.openaiApiKey,
    gemini_api_key: config.geminiApiKey,
    model: config.model,
    max_sessions: config.maxSessions,
    base_url: config.baseUrl,
    anthropic_base_url: config.anthropicBaseUrl,
    openai_base_url: config.openaiBaseUrl,
    gemini_base_url: config.geminiBaseUrl,
  };
}

function sessionToWasm(config: SessionConfig): Record<string, unknown> {
  // Plan §4d.wasm.2 + §6.13: per-session api_key / base_url fields are
  // deleted. Credentials come from bootstrap-populated config.realm or
  // the host's registered external-auth resolver.
  return {
    model: config.model,
    auth_binding: config.authBinding,
    system_prompt: config.systemPrompt,
    max_tokens: config.maxTokens,
    comms_name: config.commsName,
    keep_alive: config.keepAlive,
    labels: config.labels,
    additional_instructions: config.additionalInstructions,
    app_context: config.appContext,
  };
}

function isRecord(value: unknown): value is Record<string, unknown> {
  return typeof value === 'object' && value !== null && !Array.isArray(value);
}

function parseJsonPayload(json: string, context: string): unknown {
  try {
    return JSON.parse(json) as unknown;
  } catch (error) {
    throw new Error(`${context}: invalid JSON`);
  }
}

function parseMobListResult(raw: unknown): MobStatus[] {
  const context = 'Invalid mob/list response';
  if (!isRecord(raw)) {
    throw new Error(`${context}: malformed envelope`);
  }

  const record = raw as Partial<WireMobListResult> & Record<string, unknown>;
  if (!Array.isArray(record.mobs)) {
    throw new Error(`${context}: missing mobs`);
  }

  return record.mobs.map((mob, index) =>
    parseMobStatusResult(mob, `${context}.mobs[${index}]`),
  );
}

/**
 * The wasm-bindgen module type.
 *
 * This mirrors the exports from `meerkat_web_runtime.js` generated by wasm-pack.
 * We use a structural type so the wrapper compiles without the actual WASM build
 * artifacts present.
 */
export interface WasmModule {
  default: () => Promise<unknown>;
  runtime_version: () => string;
  init_runtime: (mobpackBytes: Uint8Array, credentialsJson: string) => string;
  init_runtime_from_config: (configJson: string) => string;
  register_tool_callback: (
    name: string,
    description: string,
    schemaJson: string,
    callback: (args: string) => Promise<string>,
  ) => void;
  register_js_tool: (
    name: string,
    description: string,
    schemaJson: string,
  ) => void;
  clear_tool_callbacks: () => void;
  destroy_runtime: () => void;
  create_session_simple: (configJson: string) => number;
  create_session: (mobpackBytes: Uint8Array, configJson: string) => number;
  start_turn: (handle: number, prompt: string) => Promise<string>;
  get_session_state: (handle: number) => string;
  destroy_session: (handle: number) => void;
  poll_events: (handle: number) => string;
  append_system_context: (handle: number, requestJson: string) => Promise<string>;
  inspect_mobpack: (mobpackBytes: Uint8Array) => string;
  mob_create: (definitionJson: string) => Promise<string>;
  mob_status: (mobId: string) => Promise<string>;
  mob_list: () => Promise<string>;
  mob_lifecycle: (mobId: string, action: string) => Promise<void>;
  mob_events: (mobId: string, afterCursor: number, limit: number) => Promise<string>;
  mob_spawn: (mobId: string, specsJson: string) => Promise<string>;
  mob_retire: (mobId: string, agentIdentity: string) => Promise<void>;
  mob_wire: (mobId: string, a: string, b: string) => Promise<void>;
  mob_unwire: (mobId: string, a: string, b: string) => Promise<void>;
  mob_wire_peer?: (mobId: string, member: string, peerJson: string) => Promise<void>;
  mob_unwire_peer?: (mobId: string, member: string, peerJson: string) => Promise<void>;
  mob_wire_target: (mobId: string, local: string, targetJson: string) => Promise<void>;
  mob_unwire_target: (mobId: string, local: string, targetJson: string) => Promise<void>;
  mob_list_members: (mobId: string) => Promise<string>;
  mob_append_system_context: (
    mobId: string,
    agentIdentity: string,
    requestJson: string,
  ) => Promise<string>;
  mob_member_send: (mobId: string, agentIdentity: string, requestJson: string) => Promise<string>;
  mob_member_status: (mobId: string, agentIdentity: string) => Promise<string>;
  mob_respawn: (mobId: string, agentIdentity: string, initialMessage?: string) => Promise<string>;
  mob_force_cancel: (mobId: string, agentIdentity: string) => Promise<void>;
  mob_spawn_helper: (mobId: string, requestJson: string) => Promise<string>;
  mob_fork_helper: (mobId: string, requestJson: string) => Promise<string>;
  mob_run_flow: (mobId: string, flowId: string, paramsJson: string) => Promise<string>;
  mob_flow_status: (mobId: string, runId: string) => Promise<string>;
  mob_cancel_flow: (mobId: string, runId: string) => Promise<void>;
  mob_member_subscribe: (mobId: string, agentIdentity: string) => Promise<number>;
  mob_subscribe_events: (mobId: string) => Promise<number>;
  poll_subscription: (handle: number) => string;
  close_subscription: (handle: number) => void;
  wire_cross_mob: (
    mobA: string,
    meerkatA: string,
    mobB: string,
    meerkatB: string,
  ) => Promise<void>;
  comms_peers: (sessionId: string) => Promise<string>;
  comms_send: (sessionId: string, paramsJson: string) => Promise<string>;
}

/** Entry point for the Meerkat WASM runtime in the browser. */
export class MeerkatRuntime {
  private wasm: WasmModule;
  private destroyed = false;

  private constructor(wasm: WasmModule) {
    this.wasm = wasm;
  }

  /**
   * Initialize the WASM runtime with the given configuration.
   *
   * @param wasm - The WASM module (import from the wasm-pack output).
   * @param config - Runtime configuration including API keys and provider URLs.
   * @returns An initialized {@link MeerkatRuntime} instance.
   *
   * @example
   * ```typescript
   * import * as wasm from '@rkat/web/wasm/meerkat_web_runtime.js';
   *
   * const runtime = await MeerkatRuntime.init(wasm, {
   *   anthropicApiKey: 'sk-...',
   *   model: 'claude-sonnet-4-5',
   * });
   * ```
   */
  static async init(
    wasm: WasmModule,
    config: RuntimeConfig,
  ): Promise<MeerkatRuntime> {
    // Initialize the WASM module (loads the .wasm binary)
    await wasm.default();

    // Version check: ensure JS wrapper and WASM binary are from the same build
    const wasmVersion = wasm.runtime_version();
    if (wasmVersion !== EXPECTED_VERSION) {
      throw new Error(
        `WASM runtime version mismatch: expected ${EXPECTED_VERSION}, got ${wasmVersion}. ` +
          'Clear browser cache or rebuild the WASM runtime.',
      );
    }

    const configJson = JSON.stringify(toWasmConfig(config));
    const resultJson = wasm.init_runtime_from_config(configJson);
    const result = JSON.parse(resultJson) as InitResult;

    if (result.status !== 'initialized') {
      throw new Error(`Runtime initialization failed: ${resultJson}`);
    }

    return new MeerkatRuntime(wasm);
  }

  /**
   * Initialize from a mobpack archive.
   *
   * @param wasm - The WASM module.
   * @param mobpackBytes - The tar.gz mobpack archive bytes.
   * @param config - Credentials and optional overrides.
   */
  static async initFromMobpack(
    wasm: WasmModule,
    mobpackBytes: Uint8Array,
    config: RuntimeConfig,
  ): Promise<MeerkatRuntime> {
    await wasm.default();

    const wasmVersion = wasm.runtime_version();
    if (wasmVersion !== EXPECTED_VERSION) {
      throw new Error(
        `WASM runtime version mismatch: expected ${EXPECTED_VERSION}, got ${wasmVersion}. ` +
          'Clear browser cache or rebuild the WASM runtime.',
      );
    }

    const credentialsJson = JSON.stringify(toWasmConfig(config));
    const resultJson = wasm.init_runtime(mobpackBytes, credentialsJson);
    const result = JSON.parse(resultJson) as InitResult;

    if (result.status !== 'initialized') {
      throw new Error(`Runtime initialization failed: ${resultJson}`);
    }

    return new MeerkatRuntime(wasm);
  }

  /**
   * Register a tool callback on the WASM module.
   *
   * Runtime state must already be initialized.
   *
   * Prefer the instance method when possible so registration is clearly scoped
   * to the active runtime.
   *
   * @param wasm - The WASM module used by the active runtime.
   * @param name - Tool name.
   * @param description - Human-readable tool description.
   * @param schema - JSON Schema for tool input.
   * @param callback - Receives tool arguments as JSON string, returns result.
   */
  static registerTool(
    wasm: WasmModule,
    name: string,
    description: string,
    schema: JsonSchema,
    callback: ToolCallback,
  ): void {
    wasm.register_tool_callback(
      name,
      description,
      JSON.stringify(schema),
      async (args: string) => JSON.stringify(await callback(args)),
    );
  }

  /**
   * Register a fire-and-forget tool on the WASM module.
   *
   * Runtime state must already be initialized.
   *
   * When the agent calls this tool, it receives `"acknowledged"` immediately
   * and continues processing. The host should watch `ToolCallRequested` events
   * in the event stream to capture args and act on them. Any response (e.g.
   * human approval) comes back through the normal session or mob member send
   * APIs.
   *
   * @param wasm - The WASM module (same one passed to init).
   * @param name - Tool name the agent will call.
   * @param description - Human-readable description shown to the LLM.
   * @param schema - JSON Schema for tool input parameters.
   */
  static registerFireAndForgetTool(
    wasm: WasmModule,
    name: string,
    description: string,
    schema: JsonSchema,
  ): void {
    wasm.register_js_tool(name, description, JSON.stringify(schema));
  }

  /** Clear all registered tool callbacks on the active WASM runtime. */
  static clearTools(wasm: WasmModule): void {
    wasm.clear_tool_callbacks();
  }

  /** Tear down the embedded runtime and invalidate browser-side handles. */
  destroy(): void {
    if (this.destroyed) return;
    this.wasm.destroy_runtime();
    this.destroyed = true;
  }

  /** Register a tool callback on this initialized runtime instance. */
  registerTool(
    name: string,
    description: string,
    schema: JsonSchema,
    callback: ToolCallback,
  ): void {
    this.wasm.register_tool_callback(
      name,
      description,
      JSON.stringify(schema),
      async (args: string) => JSON.stringify(await callback(args)),
    );
  }

  /** Register a fire-and-forget tool on this initialized runtime instance. */
  registerFireAndForgetTool(
    name: string,
    description: string,
    schema: JsonSchema,
  ): void {
    this.wasm.register_js_tool(name, description, JSON.stringify(schema));
  }

  /** Clear all runtime-scoped JS tools from this initialized runtime instance. */
  clearTools(): void {
    this.wasm.clear_tool_callbacks();
  }

  /** Create a mob from a definition. */
  async createMob(definition: MobDefinition): Promise<Mob> {
    const mobId = await this.wasm.mob_create(JSON.stringify(definition));
    return new Mob(mobId, {
      mob_spawn: this.wasm.mob_spawn,
      mob_retire: this.wasm.mob_retire,
      mob_wire: this.wasm.mob_wire,
      mob_unwire: this.wasm.mob_unwire,
      mob_wire_target: this.wasm.mob_wire_target,
      mob_unwire_target: this.wasm.mob_unwire_target,
      mob_list_members: this.wasm.mob_list_members,
      mob_append_system_context: this.wasm.mob_append_system_context,
      mob_member_send: this.wasm.mob_member_send,
      mob_member_status: this.wasm.mob_member_status,
      mob_respawn: this.wasm.mob_respawn,
      mob_force_cancel: this.wasm.mob_force_cancel,
      mob_spawn_helper: this.wasm.mob_spawn_helper,
      mob_fork_helper: this.wasm.mob_fork_helper,
      mob_status: this.wasm.mob_status,
      mob_lifecycle: this.wasm.mob_lifecycle,
      mob_events: this.wasm.mob_events,
      mob_run_flow: this.wasm.mob_run_flow,
      mob_flow_status: this.wasm.mob_flow_status,
      mob_cancel_flow: this.wasm.mob_cancel_flow,
      mob_member_subscribe: this.wasm.mob_member_subscribe,
      mob_subscribe_events: this.wasm.mob_subscribe_events,
      poll_subscription: this.wasm.poll_subscription,
      close_subscription: this.wasm.close_subscription,
    });
  }

  /** List all active mobs. */
  async listMobs(): Promise<MobStatus[]> {
    const json = await this.wasm.mob_list();
    return parseMobListResult(parseJsonPayload(json, 'Invalid mob/list response'));
  }

  /** Wire two agents across different mobs for cross-mob comms. */
  async wireCrossMob(
    mobA: string,
    meerkatA: string,
    mobB: string,
    meerkatB: string,
  ): Promise<void> {
    await this.wasm.wire_cross_mob(mobA, meerkatA, mobB, meerkatB);
  }

  /** Create a direct session façade backed by a real runtime session identity. */
  createSession(config: SessionConfig): Session {
    const handle = this.wasm.create_session_simple(
      JSON.stringify(sessionToWasm(config)),
    );
    return new Session(
      handle,
      this.wasm.start_turn,
      this.wasm.get_session_state,
      this.wasm.destroy_session,
      this.wasm.poll_events,
      this.wasm.append_system_context,
    );
  }

  /** Inspect a mobpack archive without loading it. */
  inspectMobpack(mobpackBytes: Uint8Array): unknown {
    const json = this.wasm.inspect_mobpack(mobpackBytes);
    return JSON.parse(json);
  }

  /** Get the WASM runtime version string. */
  get version(): string {
    return this.wasm.runtime_version();
  }
}
