import { access } from "node:fs/promises";
import { fileURLToPath } from "node:url";
import { describe, expect, it } from "vitest";
import { KDL_C_ABI_EXPORTS, KdlNativeAbi } from "../../src/kdl/nativeAbi.js";
import type { NativeKdlModule } from "../../src/kdl/nativeModule.js";
import { loadRobotFromUrdfModel } from "../../src/robot/urdfParser.js";

const BUILD_DIR = new URL("../../../build-wasm/", import.meta.url);
const WRAPPER_URL = new URL("kdl.js", BUILD_DIR);

type NativeFactory = (options?: {
  locateFile?: (path: string, prefix: string) => string;
}) => Promise<NativeKdlModule>;

async function loadNativeModule(): Promise<NativeKdlModule> {
  await access(fileURLToPath(WRAPPER_URL));
  const imported = (await import(/* @vite-ignore */ WRAPPER_URL.href)) as {
    default?: NativeFactory;
    createKdlModule?: NativeFactory;
  };
  const factory = imported.default ?? imported.createKdlModule;
  if (!factory) {
    throw new Error("kdl.js did not export createKdlModule");
  }
  return factory({
    locateFile: (path) => fileURLToPath(new URL(path, BUILD_DIR))
  });
}

const NATIVE_SOLVER_URDF = `
<robot name="native_solver">
  <link name="base_link"/>
  <link name="link_1"/>
  <link name="tool0"/>
  <joint name="joint_1" type="revolute">
    <parent link="base_link"/>
    <child link="link_1"/>
    <axis xyz="0 0 1"/>
    <limit lower="-3.141592653589793" upper="3.141592653589793" velocity="2" acceleration="4"/>
  </joint>
  <joint name="joint_2" type="prismatic">
    <parent link="link_1"/>
    <child link="tool0"/>
    <axis xyz="1 0 0"/>
    <limit lower="0" upper="1" velocity="0.5" acceleration="1"/>
  </joint>
</robot>
`;

function writeFloat64Array(native: NativeKdlModule, values: number[]): number {
  const bytes = values.length * Float64Array.BYTES_PER_ELEMENT;
  const ptr = native._malloc?.(bytes);
  if (!ptr) {
    throw new Error(`Failed to allocate ${bytes} bytes`);
  }
  native.HEAPF64?.set(values, ptr / Float64Array.BYTES_PER_ELEMENT);
  return ptr;
}

function readFloat64Array(native: NativeKdlModule, ptr: number, length: number): number[] {
  return Array.from(native.HEAPF64?.subarray(
    ptr / Float64Array.BYTES_PER_ELEMENT,
    ptr / Float64Array.BYTES_PER_ELEMENT + length
  ) ?? []);
}

describe("KDL C ABI", () => {
  it("exports the stable P0 ABI names", async () => {
    const abi = new KdlNativeAbi(await loadNativeModule());

    expect(() => abi.assertExports()).not.toThrow();
    expect(KDL_C_ABI_EXPORTS).toEqual([
      "kdl_init",
      "kdl_create_robot",
      "kdl_destroy_robot",
      "kdl_get_robot_info",
      "kdl_fk",
      "kdl_fk_all_links",
      "kdl_jacobian",
      "kdl_ik",
      "kdl_plan_movej",
      "kdl_plan_movel",
      "kdl_plan_movec",
      "kdl_plan_path",
      "kdl_sample_trap",
      "kdl_last_error"
    ]);
  });

  it("initializes, caches model handles, returns JSON info, and destroys handles", async () => {
    const abi = new KdlNativeAbi(await loadNativeModule());
    const model = loadRobotFromUrdfModel(NATIVE_SOLVER_URDF, {
      robotId: "abi",
      baseLink: "base_link",
      tipLink: "tool0"
    });

    expect(abi.callNumber("kdl_init", ["string"], ["{}"])).toBe(0);
    const handle = abi.callNumber("kdl_create_robot", ["string"], [JSON.stringify(model)]);
    expect(handle).toBeGreaterThan(0);

    const info = abi.readJsonCall<{ handle: number; nativeState: string; dof: number }>(
      "kdl_get_robot_info",
      ["number"],
      [handle]
    );
    expect(info).toMatchObject({
      handle,
      dof: 2,
      nativeState: "kdl_chain"
    });

    expect(abi.callNumber("kdl_destroy_robot", ["number"], [handle])).toBe(0);
  });

  it("constructs a native KDL chain and returns real FK, fkAllLinks, Jacobian, and IK data", async () => {
    const native = await loadNativeModule();
    const abi = new KdlNativeAbi(native);
    const model = loadRobotFromUrdfModel(NATIVE_SOLVER_URDF, {
      robotId: "native",
      baseLink: "base_link",
      tipLink: "tool0"
    });

    expect(abi.callNumber("kdl_init", ["string"], ["{}"])).toBe(0);
    const handle = abi.callNumber("kdl_create_robot", ["string"], [JSON.stringify(model)]);
    expect(handle).toBeGreaterThan(0);
    expect(abi.readJsonCall("kdl_get_robot_info", ["number"], [handle])).toMatchObject({
      handle,
      dof: 2,
      jointNames: ["joint_1", "joint_2"],
      nativeState: "kdl_chain"
    });

    const joints = writeFloat64Array(native, [Math.PI / 2, 0.4]);
    const pose = native._malloc?.(7 * Float64Array.BYTES_PER_ELEMENT);
    const jacobian = native._malloc?.(12 * Float64Array.BYTES_PER_ELEMENT);
    const ikSeed = writeFloat64Array(native, [1.4, 0.3]);
    const ikOut = native._malloc?.(2 * Float64Array.BYTES_PER_ELEMENT);
    expect(pose).toBeTruthy();
    expect(jacobian).toBeTruthy();
    expect(ikOut).toBeTruthy();

    try {
      expect(abi.callNumber("kdl_fk", ["number", "number", "number", "number"], [handle, joints, 2, pose])).toBe(0);
      const pose7 = readFloat64Array(native, pose!, 7);
      expect(pose7[0]).toBeCloseTo(0);
      expect(pose7[1]).toBeCloseTo(0.4);
      expect(pose7[2]).toBeCloseTo(0);
      expect(pose7[5]).toBeCloseTo(Math.SQRT1_2);
      expect(pose7[6]).toBeCloseTo(Math.SQRT1_2);

      const links = abi.readJsonCall<{
        ok: boolean;
        linkPoses: Array<{ link: string; pose: { position: [number, number, number] } }>;
      }>("kdl_fk_all_links", ["number", "number", "number"], [handle, joints, 2]);
      expect(links.ok).toBe(true);
      expect(links.linkPoses.map((linkPose) => linkPose.link)).toEqual(["base_link", "link_1", "tool0"]);
      expect(links.linkPoses[2]?.pose.position[0]).toBeCloseTo(0);
      expect(links.linkPoses[2]?.pose.position[1]).toBeCloseTo(0.4);

      expect(abi.callNumber("kdl_jacobian", ["number", "number", "number", "number"], [handle, joints, 2, jacobian])).toBe(0);
      const jac = readFloat64Array(native, jacobian!, 12);
      expect(jac[0]).toBeCloseTo(-0.4, 4);
      expect(jac[1]).toBeCloseTo(0, 4);
      expect(jac[2]).toBeCloseTo(0, 4);
      expect(jac[3]).toBeCloseTo(1, 4);
      expect(jac[10]).toBeCloseTo(1, 4);

      expect(
        abi.callNumber(
          "kdl_ik",
          ["number", "number", "number", "number", "string", "number"],
          [handle, ikSeed, 2, pose, "{}", ikOut]
        )
      ).toBe(0);
      const ikJoints = readFloat64Array(native, ikOut!, 2);
      expect(ikJoints[0]).toBeCloseTo(Math.PI / 2, 4);
      expect(ikJoints[1]).toBeCloseTo(0.4, 4);
    } finally {
      native._free?.(joints);
      native._free?.(ikSeed);
      if (pose) {
        native._free?.(pose);
      }
      if (jacobian) {
        native._free?.(jacobian);
      }
      if (ikOut) {
        native._free?.(ikOut);
      }
      abi.callNumber("kdl_destroy_robot", ["number"], [handle]);
    }
  });

  it("normalizes C ABI failures through kdl_last_error", async () => {
    const native = await loadNativeModule();
    const abi = new KdlNativeAbi(native);
    const seed = writeFloat64Array(native, [0, 0]);
    const target = writeFloat64Array(native, [0, 0, 0, 0, 0, 0, 1]);
    const out = native._malloc?.(2 * Float64Array.BYTES_PER_ELEMENT);

    try {
      expect(abi.callNumber("kdl_init", ["string"], ["{}"])).toBe(0);
      const returnCode = abi.callNumber(
        "kdl_ik",
        ["number", "number", "number", "number", "string", "number"],
        [404, seed, 2, target, "{}", out]
      );

      expect(returnCode).toBe(-1);
      expect(abi.lastError()).toMatchObject({
        code: "KDL_INVALID_HANDLE",
        diagnostics: [
          {
            severity: "error",
            code: "KDL_INVALID_HANDLE"
          }
        ]
      });
      expect(() => abi.checkReturnCode(returnCode)).toThrowError(
        expect.objectContaining({
          code: "KDL_INVALID_HANDLE"
        })
      );
    } finally {
      native._free?.(seed);
      native._free?.(target);
      if (out) {
        native._free?.(out);
      }
    }
  });

  it("returns native trapezoid samples through kdl_sample_trap", async () => {
    const abi = new KdlNativeAbi(await loadNativeModule());

    expect(abi.callNumber("kdl_init", ["string"], ["{}"])).toBe(0);
    const profile = abi.readJsonCall<{
      ok: boolean;
      type: string;
      duration: number;
      samples: Array<{ index: number; time: number; s: number; sd: number; sdd: number }>;
      diagnostics: Array<{ code: string }>;
    }>(
      "kdl_sample_trap",
      ["number", "string"],
      [
        2,
        JSON.stringify({
          maxVelocity: 1,
          maxAcceleration: 1,
          sampleTime: 0.25
        })
      ]
    );

    expect(profile).toMatchObject({
      ok: true,
      type: "trapezoid",
      duration: 3,
      diagnostics: []
    });
    expect(profile.samples[0]).toMatchObject({ index: 0, time: 0, s: 0 });
    expect(profile.samples.find((sample) => sample.time === 1)?.s).toBeCloseTo(0.25);
    expect(profile.samples.find((sample) => sample.time === 1)?.sd).toBeCloseTo(0.5);
    expect(profile.samples.at(-1)).toMatchObject({ s: 1 });

    const triangle = abi.readJsonCall<{
      type: string;
      diagnostics: Array<{ code: string }>;
    }>(
      "kdl_sample_trap",
      ["number", "string"],
      [
        0.5,
        JSON.stringify({
          maxVelocity: 2,
          maxAcceleration: 1,
          sampleTime: 0.1
        })
      ]
    );
    expect(triangle.type).toBe("triangle");
    expect(triangle.diagnostics).toContainEqual(
      expect.objectContaining({
        code: "KDL_TRAP_TRIANGLE_PROFILE"
      })
    );

    const returnCode = abi.callNumber("kdl_sample_trap", ["number", "string", "number", "number"], [1, "{}", 0, 0]);
    expect(returnCode).toBe(-1);
    expect(abi.lastError()).toMatchObject({
      code: "KDL_INVALID_TRAP_PROFILE"
    });
  });

  it("reports JSON output buffer errors without raw strings", async () => {
    const native = await loadNativeModule();
    const abi = new KdlNativeAbi(native);
    const model = loadRobotFromUrdfModel(NATIVE_SOLVER_URDF, {
      robotId: "abi",
      baseLink: "base_link",
      tipLink: "tool0"
    });

    expect(abi.callNumber("kdl_init", ["string"], ["{}"])).toBe(0);
    const handle = abi.callNumber("kdl_create_robot", ["string"], [JSON.stringify(model)]);
    const ptr = native._malloc?.(4);
    expect(ptr).toBeTruthy();

    try {
      const returnCode = abi.callNumber("kdl_get_robot_info", ["number", "number", "number"], [handle, ptr, 4]);
      expect(returnCode).toBe(-1);
      expect(abi.lastError()).toMatchObject({
        code: "KDL_BUFFER_TOO_SMALL",
        diagnostics: [
          {
            severity: "error",
            code: "KDL_BUFFER_TOO_SMALL"
          }
        ]
      });
    } finally {
      if (ptr) {
        native._free?.(ptr);
      }
    }
  });
});