import { describe, test, expect, vi } from "vitest";
import { z } from "zod";
import { ConnectError, Code } from "@connectrpc/connect";
import type { Client } from "@connectrpc/connect";
import type { grackle } from "@grackle-ai/common";
import { envTools } from "./env.js";

type GrackleClient = Client<typeof grackle.GrackleCore>;

/** Helper to find a tool definition by name. */
const getTool = (name: string) => envTools.find((t) => t.name === name)!;

/** Create a mock Grackle client with all env-related methods stubbed. */
function createMockClient(): GrackleClient {
  return {
    listEnvironments: vi.fn(),
    addEnvironment: vi.fn(),
    provisionEnvironment: vi.fn(),
    stopEnvironment: vi.fn(),
    destroyEnvironment: vi.fn(),
    removeEnvironment: vi.fn(),
    listDockerContainers: vi.fn(),
  } as unknown as GrackleClient;
}

describe("env_list", () => {
  /** Should return mapped environment summaries on success. */
  test("happy path returns environment list", async () => {
    const mockClient = createMockClient();
    (mockClient.listEnvironments as ReturnType<typeof vi.fn>).mockResolvedValue({
      environments: [
        {
          id: "env-1",
          displayName: "Dev Box",
          adapterType: "ssh",
          status: "running",
        },
      ],
    });

    const result = await getTool("env_list").handler({}, { core: mockClient });
    const parsed = JSON.parse(result.content[0].text);

    expect(parsed).toEqual([
      {
        id: "env-1",
        displayName: "Dev Box",
        adapterType: "ssh",
        status: "running",
      },
    ]);
    expect(result.isError).toBeUndefined();
  });

  /** Should return a structured error when the gRPC call fails with ConnectError. */
  test("ConnectError returns isError result", async () => {
    const mockClient = createMockClient();
    (mockClient.listEnvironments as ReturnType<typeof vi.fn>).mockRejectedValue(
      new ConnectError("not found", Code.NotFound),
    );

    const result = await getTool("env_list").handler({}, { core: mockClient });
    const parsed = JSON.parse(result.content[0].text);

    expect(result.isError).toBe(true);
    expect(parsed.code).toBe("NOT_FOUND");
  });
});

describe("env_list_docker_containers", () => {
  /** Should map running containers from the RPC result. */
  test("happy path returns mapped container list", async () => {
    const mockClient = createMockClient();
    (mockClient.listDockerContainers as ReturnType<typeof vi.fn>).mockResolvedValue({
      containers: [
        { id: "abc123", name: "demo-ext", image: "node:22", state: "running", status: "Up 3 minutes" },
      ],
      error: "",
    });

    const result = await getTool("env_list_docker_containers").handler({}, { core: mockClient });
    const parsed = JSON.parse(result.content[0].text);

    expect(parsed).toEqual({
      containers: [
        { id: "abc123", name: "demo-ext", image: "node:22", state: "running", status: "Up 3 minutes" },
      ],
      error: "",
    });
    expect(result.isError).toBeUndefined();
  });

  /** Should surface a non-fatal docker error from the RPC. */
  test("non-fatal docker error is returned in the payload", async () => {
    const mockClient = createMockClient();
    (mockClient.listDockerContainers as ReturnType<typeof vi.fn>).mockResolvedValue({
      containers: [],
      error: "docker: command not found",
    });

    const result = await getTool("env_list_docker_containers").handler({}, { core: mockClient });
    const parsed = JSON.parse(result.content[0].text);

    expect(parsed.error).toContain("command not found");
    expect(parsed.containers).toEqual([]);
  });

  /** Should return a structured error on ConnectError. */
  test("ConnectError returns isError result", async () => {
    const mockClient = createMockClient();
    (mockClient.listDockerContainers as ReturnType<typeof vi.fn>).mockRejectedValue(
      new ConnectError("unavailable", Code.Unavailable),
    );

    const result = await getTool("env_list_docker_containers").handler({}, { core: mockClient });

    expect(result.isError).toBe(true);
    const parsed = JSON.parse(result.content[0].text);
    expect(parsed.code).toBe("UNAVAILABLE");
  });

  /** Tool should be read-only. */
  test("is annotated read-only", () => {
    const tool = getTool("env_list_docker_containers");
    expect(tool.mutating).toBe(false);
    expect(tool.annotations?.readOnlyHint).toBe(true);
  });
});

describe("env_add", () => {
  /** Should JSON.stringify adapterConfig. */
  test("happy path stringifies adapterConfig", async () => {
    const mockClient = createMockClient();
    (mockClient.addEnvironment as ReturnType<typeof vi.fn>).mockResolvedValue({
      id: "env-2",
      displayName: "My Env",
    });

    const result = await getTool("env_add").handler(
      {
        displayName: "My Env",
        adapterType: "codespace",
        adapterConfig: { codespaceName: "fluffy-spork" },
      },
      { core: mockClient },
    );

    expect(mockClient.addEnvironment).toHaveBeenCalledWith({
      displayName: "My Env",
      adapterType: "codespace",
      adapterConfig: JSON.stringify({ codespaceName: "fluffy-spork" }),
      githubAccountId: "",
    });

    const parsed = JSON.parse(result.content[0].text);
    expect(parsed.id).toBe("env-2");
    expect(result.isError).toBeUndefined();
  });

  /** Should pass an empty JSON object (not "") for adapterConfig when not provided, so the
   * server's later JSON.parse at provision time does not throw. */
  test("omitted adapterConfig sends empty JSON object", async () => {
    const mockClient = createMockClient();
    (mockClient.addEnvironment as ReturnType<typeof vi.fn>).mockResolvedValue({ id: "env-3" });

    await getTool("env_add").handler(
      { displayName: "Plain", adapterType: "local" },
      { core: mockClient },
    );

    expect(mockClient.addEnvironment).toHaveBeenCalledWith({
      displayName: "Plain",
      adapterType: "local",
      adapterConfig: "{}",
      githubAccountId: "",
    });
  });

  /** Should forward githubAccountId as the top-level request field for codespace. */
  test("forwards githubAccountId to the gRPC call", async () => {
    const mockClient = createMockClient();
    (mockClient.addEnvironment as ReturnType<typeof vi.fn>).mockResolvedValue({ id: "env-cs" });

    await getTool("env_add").handler(
      {
        displayName: "CS",
        adapterType: "codespace",
        adapterConfig: { codespaceName: "fluffy-spork" },
        githubAccountId: "acct-123",
      },
      { core: mockClient },
    );

    expect(mockClient.addEnvironment).toHaveBeenCalledWith({
      displayName: "CS",
      adapterType: "codespace",
      adapterConfig: JSON.stringify({ codespaceName: "fluffy-spork" }),
      githubAccountId: "acct-123",
    });
  });

  /** Should return a structured error on ConnectError. */
  test("ConnectError returns isError result", async () => {
    const mockClient = createMockClient();
    (mockClient.addEnvironment as ReturnType<typeof vi.fn>).mockRejectedValue(
      new ConnectError("already exists", Code.AlreadyExists),
    );

    const result = await getTool("env_add").handler(
      { displayName: "Dup", adapterType: "ssh", adapterConfig: { host: "example.com" } },
      { core: mockClient },
    );

    expect(result.isError).toBe(true);
    const parsed = JSON.parse(result.content[0].text);
    expect(parsed.code).toBe("ALREADY_EXISTS");
  });
});

describe("env_add inputSchema (discriminated union)", () => {
  const schema = getTool("env_add").inputSchema;

  /** ssh requires a host; omitting it must fail with an issue on that field. */
  test("ssh without host fails validation", () => {
    const result = schema.safeParse({
      displayName: "Box",
      adapterType: "ssh",
      adapterConfig: {},
    });
    expect(result.success).toBe(false);
    if (!result.success) {
      expect(result.error.issues.some((i) => i.path.includes("host"))).toBe(true);
    }
  });

  /** ssh with a host is valid. */
  test("ssh with host passes validation", () => {
    const result = schema.safeParse({
      displayName: "Box",
      adapterType: "ssh",
      adapterConfig: { host: "example.com", sshPort: 2222 },
    });
    expect(result.success).toBe(true);
  });

  /** codespace requires codespaceName. */
  test("codespace without codespaceName fails validation", () => {
    const result = schema.safeParse({
      displayName: "CS",
      adapterType: "codespace",
      adapterConfig: {},
    });
    expect(result.success).toBe(false);
  });

  /** docker has no required config fields. */
  test("docker without adapterConfig passes validation", () => {
    const result = schema.safeParse({ displayName: "Box", adapterType: "docker" });
    expect(result.success).toBe(true);
  });

  /** local has no required config fields. */
  test("local without adapterConfig passes validation", () => {
    const result = schema.safeParse({ displayName: "Here", adapterType: "local" });
    expect(result.success).toBe(true);
  });

  /** Unknown adapter types are rejected by the discriminated union. */
  test("unknown adapterType fails validation", () => {
    const result = schema.safeParse({
      displayName: "K8s",
      adapterType: "kubernetes",
      adapterConfig: {},
    });
    expect(result.success).toBe(false);
  });

  /** Strict config objects reject unknown/typo'd fields instead of silently dropping them. */
  test("unknown config field fails validation (strict)", () => {
    const result = schema.safeParse({
      displayName: "Box",
      adapterType: "ssh",
      adapterConfig: { host: "example.com", bogus: true },
    });
    expect(result.success).toBe(false);
  });

  /**
   * The generated JSON Schema is what an MCP client (and the model) actually sees, and the
   * MCP server pre-computes it for every tool at session start — so it must not throw and must
   * expose one branch per adapter type. (Zod 4 may emit the literal as `const` or single-value
   * `enum`, so we assert on the serialized output rather than a fixed shape.)
   */
  test("renders to a JSON Schema with a branch per adapter type", () => {
    const json = z.toJSONSchema(schema) as Record<string, unknown>;
    const branches = (json.anyOf ?? json.oneOf) as unknown[] | undefined;
    expect(Array.isArray(branches)).toBe(true);
    expect(branches).toHaveLength(4);

    const serialized = JSON.stringify(json);
    for (const adapterType of ["local", "ssh", "codespace", "docker"]) {
      expect(serialized).toContain(`"${adapterType}"`);
    }
  });
});

describe("env_provision", () => {
  /** Should collect streaming events and return finalStatus success. */
  test("happy path collects stream events", async () => {
    const mockClient = createMockClient();
    const mockStream = (async function* () {
      yield { stage: "init", message: "Starting", progress: 0 };
      yield { stage: "install", message: "Installing agent", progress: 50 };
      yield { stage: "done", message: "Complete", progress: 100 };
    })();
    (mockClient.provisionEnvironment as ReturnType<typeof vi.fn>).mockReturnValue(mockStream);

    const result = await getTool("env_provision").handler(
      { environmentId: "env-1" }, { core: mockClient },
    );
    const parsed = JSON.parse(result.content[0].text);

    expect(parsed.finalStatus).toBe("success");
    expect(parsed.events).toHaveLength(3);
    expect(parsed.events[0]).toEqual({ stage: "init", message: "Starting", progress: 0 });
    expect(parsed.events[2]).toEqual({ stage: "done", message: "Complete", progress: 100 });
    expect(result.isError).toBeUndefined();
  });

  /** Should return structured error on ConnectError during streaming. */
  test("ConnectError returns isError result", async () => {
    const mockClient = createMockClient();
    const mockStream = (async function* () {
      yield { stage: "init", message: "Starting", progress: 0 };
      throw new ConnectError("env not found", Code.NotFound);
    })();
    (mockClient.provisionEnvironment as ReturnType<typeof vi.fn>).mockReturnValue(mockStream);

    const result = await getTool("env_provision").handler(
      { environmentId: "env-missing" }, { core: mockClient },
    );
    const parsed = JSON.parse(result.content[0].text);

    expect(result.isError).toBe(true);
    expect(parsed.code).toBe("NOT_FOUND");
  });

  /** Should wrap non-ConnectError with collected events and finalStatus error. */
  test("non-ConnectError wraps with collected events", async () => {
    const mockClient = createMockClient();
    const mockStream = (async function* () {
      yield { stage: "init", message: "Starting", progress: 0 };
      throw new Error("unexpected failure");
    })();
    (mockClient.provisionEnvironment as ReturnType<typeof vi.fn>).mockReturnValue(mockStream);

    const result = await getTool("env_provision").handler(
      { environmentId: "env-1" }, { core: mockClient },
    );
    const parsed = JSON.parse(result.content[0].text);

    expect(result.isError).toBe(true);
    expect(parsed.finalStatus).toBe("error");
    expect(parsed.error).toBe("unexpected failure");
    expect(parsed.events).toHaveLength(1);
    expect(parsed.events[0].stage).toBe("init");
  });
  /** Should pass force flag through to provisionEnvironment. */
  test("force flag is passed through to gRPC call", async () => {
    const mockClient = createMockClient();
    const mockStream = (async function* () {
      yield { stage: "done", message: "Complete", progress: 100 };
    })();
    (mockClient.provisionEnvironment as ReturnType<typeof vi.fn>).mockReturnValue(mockStream);

    await getTool("env_provision").handler(
      { environmentId: "env-1", force: true },
      { core: mockClient },
    );

    expect(mockClient.provisionEnvironment).toHaveBeenCalledWith({
      id: "env-1",
      force: true,
    });
  });

  /** Should default force to false when omitted. */
  test("force defaults to false when omitted", async () => {
    const mockClient = createMockClient();
    const mockStream = (async function* () {
      yield { stage: "done", message: "Complete", progress: 100 };
    })();
    (mockClient.provisionEnvironment as ReturnType<typeof vi.fn>).mockReturnValue(mockStream);

    await getTool("env_provision").handler(
      { environmentId: "env-1" }, { core: mockClient },
    );

    expect(mockClient.provisionEnvironment).toHaveBeenCalledWith({
      id: "env-1",
      force: false,
    });
  });
});

describe("env_stop", () => {
  /** Should call stopEnvironment and return success. */
  test("happy path returns success", async () => {
    const mockClient = createMockClient();
    (mockClient.stopEnvironment as ReturnType<typeof vi.fn>).mockResolvedValue({});

    const result = await getTool("env_stop").handler(
      { environmentId: "env-1" }, { core: mockClient },
    );
    const parsed = JSON.parse(result.content[0].text);

    expect(parsed.success).toBe(true);
    expect(mockClient.stopEnvironment).toHaveBeenCalledWith({ id: "env-1" });
  });

  /** Should return a structured error on ConnectError. */
  test("ConnectError returns isError result", async () => {
    const mockClient = createMockClient();
    (mockClient.stopEnvironment as ReturnType<typeof vi.fn>).mockRejectedValue(
      new ConnectError("precondition failed", Code.FailedPrecondition),
    );

    const result = await getTool("env_stop").handler(
      { environmentId: "env-1" }, { core: mockClient },
    );

    expect(result.isError).toBe(true);
    const parsed = JSON.parse(result.content[0].text);
    expect(parsed.code).toBe("FAILED_PRECONDITION");
  });
});

describe("env_destroy", () => {
  /** Should call destroyEnvironment and return success. */
  test("happy path returns success", async () => {
    const mockClient = createMockClient();
    (mockClient.destroyEnvironment as ReturnType<typeof vi.fn>).mockResolvedValue({});

    const result = await getTool("env_destroy").handler(
      { environmentId: "env-1" }, { core: mockClient },
    );
    const parsed = JSON.parse(result.content[0].text);

    expect(parsed.success).toBe(true);
    expect(mockClient.destroyEnvironment).toHaveBeenCalledWith({ id: "env-1" });
  });

  /** Should return a structured error on ConnectError. */
  test("ConnectError returns isError result", async () => {
    const mockClient = createMockClient();
    (mockClient.destroyEnvironment as ReturnType<typeof vi.fn>).mockRejectedValue(
      new ConnectError("internal", Code.Internal),
    );

    const result = await getTool("env_destroy").handler(
      { environmentId: "env-1" }, { core: mockClient },
    );

    expect(result.isError).toBe(true);
  });
});

describe("env_remove", () => {
  /** Should call removeEnvironment and return success. */
  test("happy path returns success", async () => {
    const mockClient = createMockClient();
    (mockClient.removeEnvironment as ReturnType<typeof vi.fn>).mockResolvedValue({});

    const result = await getTool("env_remove").handler(
      { environmentId: "env-1" }, { core: mockClient },
    );
    const parsed = JSON.parse(result.content[0].text);

    expect(parsed.success).toBe(true);
    expect(mockClient.removeEnvironment).toHaveBeenCalledWith({ id: "env-1" });
  });

  /** Should return a structured error on ConnectError. */
  test("ConnectError returns isError result", async () => {
    const mockClient = createMockClient();
    (mockClient.removeEnvironment as ReturnType<typeof vi.fn>).mockRejectedValue(
      new ConnectError("precondition", Code.FailedPrecondition),
    );

    const result = await getTool("env_remove").handler(
      { environmentId: "env-1" }, { core: mockClient },
    );

    expect(result.isError).toBe(true);
    const parsed = JSON.parse(result.content[0].text);
    expect(parsed.code).toBe("FAILED_PRECONDITION");
  });
});

describe("env_wake", () => {
  /** Should collect streaming events and return finalStatus success. */
  test("happy path collects stream events", async () => {
    const mockClient = createMockClient();
    const mockStream = (async function* () {
      yield { stage: "wake", message: "Waking up", progress: 0 };
      yield { stage: "done", message: "Ready", progress: 100 };
    })();
    (mockClient.provisionEnvironment as ReturnType<typeof vi.fn>).mockReturnValue(mockStream);

    const result = await getTool("env_wake").handler(
      { environmentId: "env-1" }, { core: mockClient },
    );
    const parsed = JSON.parse(result.content[0].text);

    expect(parsed.finalStatus).toBe("success");
    expect(parsed.events).toHaveLength(2);
    expect(result.isError).toBeUndefined();
  });

  /** Should return structured error on ConnectError. */
  test("ConnectError returns isError result", async () => {
    const mockClient = createMockClient();
    const mockStream = (async function* () {
      throw new ConnectError("unavailable", Code.Unavailable);
    })();
    (mockClient.provisionEnvironment as ReturnType<typeof vi.fn>).mockReturnValue(mockStream);

    const result = await getTool("env_wake").handler(
      { environmentId: "env-1" }, { core: mockClient },
    );

    expect(result.isError).toBe(true);
    const parsed = JSON.parse(result.content[0].text);
    expect(parsed.code).toBe("UNAVAILABLE");
  });

  /** Should wrap non-ConnectError with collected events and finalStatus error. */
  test("non-ConnectError wraps with collected events", async () => {
    const mockClient = createMockClient();
    const mockStream = (async function* () {
      yield { stage: "wake", message: "Waking", progress: 10 };
      throw new Error("connection lost");
    })();
    (mockClient.provisionEnvironment as ReturnType<typeof vi.fn>).mockReturnValue(mockStream);

    const result = await getTool("env_wake").handler(
      { environmentId: "env-1" }, { core: mockClient },
    );
    const parsed = JSON.parse(result.content[0].text);

    expect(result.isError).toBe(true);
    expect(parsed.finalStatus).toBe("error");
    expect(parsed.error).toBe("connection lost");
    expect(parsed.events).toHaveLength(1);
  });
});
