import { expect, expectTypeOf, test } from "vitest";
import * as z from "zod/mini";

test("z.object", () => {
  const a = z.object({
    name: z.string(),
    age: z.number(),
    points: z.optional(z.number()),
    "test?": z.boolean(),
  });

  a._zod.def.shape["test?"];
  a._zod.def.shape.points._zod.optin;

  type a = z.output<typeof a>;

  expectTypeOf<a>().toEqualTypeOf<{
    name: string;
    age: number;
    points?: number | undefined;
    "test?": boolean;
  }>();
  expect(z.parse(a, { name: "john", age: 30, "test?": true })).toEqual({
    name: "john",
    age: 30,
    "test?": true,
  });
  // "test?" is required in ZodObject
  expect(() => z.parse(a, { name: "john", age: "30" })).toThrow();
  expect(() => z.parse(a, "hello")).toThrow();

  // null prototype
  const schema = z.object({ a: z.string() });
  const obj = Object.create(null);
  obj.a = "foo";
  expect(schema.parse(obj)).toEqual({ a: "foo" });
});

test("z.object().check()", () => {
  const a = z.object({
    name: z.string(),
    age: z.number(),
    points: z.optional(z.number()),
    "test?": z.boolean(),
  });

  type a = z.output<typeof a>;

  a.check(({ value }) => {
    expectTypeOf(value).toEqualTypeOf<a>();
  });
});

test("z.strictObject", () => {
  const a = z.strictObject({
    name: z.string(),
  });
  expect(z.parse(a, { name: "john" })).toEqual({ name: "john" });
  expect(() => z.parse(a, { name: "john", age: 30 })).toThrow();
  expect(() => z.parse(a, "hello")).toThrow();
});

test("z.looseObject", () => {
  const a = z.looseObject({
    name: z.string(),
    age: z.number(),
  });
  expect(z.parse(a, { name: "john", age: 30 })).toEqual({
    name: "john",
    age: 30,
  });
  expect(z.parse(a, { name: "john", age: 30, extra: true })).toEqual({
    name: "john",
    age: 30,
    extra: true,
  });
  expect(() => z.parse(a, "hello")).toThrow();
});

const userSchema = z.object({
  name: z.string(),
  age: z.number(),
  email: z.optional(z.string()),
});

test("z.keyof", () => {
  // z.keyof returns an enum schema of the keys of an object schema
  const userKeysSchema = z.keyof(userSchema);
  type UserKeys = z.infer<typeof userKeysSchema>;
  expectTypeOf<UserKeys>().toEqualTypeOf<"name" | "age" | "email">();
  expect(userKeysSchema).toBeDefined();
  expect(userKeysSchema._zod.def.type).toBe("enum");
  expect(userKeysSchema._zod.def.entries).toEqual({
    name: "name",
    age: "age",
    email: "email",
  });
  expect(z.safeParse(userKeysSchema, "name").success).toBe(true);
  expect(z.safeParse(userKeysSchema, "age").success).toBe(true);
  expect(z.safeParse(userKeysSchema, "email").success).toBe(true);
  expect(z.safeParse(userKeysSchema, "isAdmin").success).toBe(false);
});

test("z.extend", () => {
  const extendedSchema = z.extend(userSchema, {
    isAdmin: z.boolean(),
  });
  type ExtendedUser = z.infer<typeof extendedSchema>;
  expectTypeOf<ExtendedUser>().toEqualTypeOf<{
    name: string;
    age: number;
    email?: string | undefined;
    isAdmin: boolean;
  }>();
  expect(extendedSchema).toBeDefined();
  expect(z.safeParse(extendedSchema, { name: "John", age: 30, isAdmin: true }).success).toBe(true);
});

test("z.merge", () => {
  const mergedSchema = z.merge(userSchema, z.object({ isAdmin: z.boolean() }));
  type MergedUser = z.infer<typeof mergedSchema>;
  expectTypeOf<MergedUser>().toEqualTypeOf<{
    name: string;
    age: number;
    email?: string | undefined;
    isAdmin: boolean;
  }>();
  expect(z.safeParse(mergedSchema, { name: "John", age: 30, isAdmin: true }).success).toBe(true);

  // the second schema's catchall wins, matching classic
  const strict = z.merge(userSchema, z.strictObject({ isAdmin: z.boolean() }));
  expect(z.safeParse(strict, { name: "John", age: 30, isAdmin: true, extra: 1 }).success).toBe(false);

  expect(() =>
    // @ts-expect-error second argument is a schema, not a shape
    z.merge(userSchema, { isAdmin: z.boolean() })
  ).toThrow("Invalid input to merge");
});

test("z.safeExtend", () => {
  const extended = z.safeExtend(userSchema, { name: z.string() });
  expect(z.safeParse(extended, { name: "John", age: 30 }).success).toBe(true);
  type Extended = z.infer<typeof extended>;
  expectTypeOf<Extended>().toEqualTypeOf<{ name: string; age: number; email?: string | undefined }>();
  // @ts-expect-error
  z.safeExtend(userSchema, { name: z.number() });
});

test("z.pick", () => {
  const pickedSchema = z.pick(userSchema, { name: true, email: true });
  type PickedUser = z.infer<typeof pickedSchema>;
  expectTypeOf<PickedUser>().toEqualTypeOf<{ name: string; email?: string | undefined }>();
  expect(pickedSchema).toBeDefined();
  expect(z.safeParse(pickedSchema, { name: "John", email: "john@example.com" }).success).toBe(true);
});

test("z.omit", () => {
  const omittedSchema = z.omit(userSchema, { age: true });
  type OmittedUser = z.infer<typeof omittedSchema>;
  expectTypeOf<OmittedUser>().toEqualTypeOf<{
    name: string;
    email?: string | undefined;
  }>();
  expect(omittedSchema).toBeDefined();
  expect(Reflect.ownKeys(omittedSchema._zod.def.shape)).toEqual(["name", "email"]);
  expect(z.safeParse(omittedSchema, { name: "John", email: "john@example.com" }).success).toBe(true);
});

test("z.partial", () => {
  const partialSchema = z.partial(userSchema);
  type PartialUser = z.infer<typeof partialSchema>;
  expectTypeOf<PartialUser>().toEqualTypeOf<{
    name?: string | undefined;
    age?: number | undefined;
    email?: string | undefined;
  }>();
  expect(z.safeParse(partialSchema, { name: "John" }).success).toBe(true);
});

test("z.partial with mask", () => {
  const partialSchemaWithMask = z.partial(userSchema, { name: true });
  type PartialUserWithMask = z.infer<typeof partialSchemaWithMask>;
  expectTypeOf<PartialUserWithMask>().toEqualTypeOf<{
    name?: string | undefined;
    age: number;
    email?: string | undefined;
  }>();
  expect(z.safeParse(partialSchemaWithMask, { age: 30 }).success).toBe(true);
  expect(z.safeParse(partialSchemaWithMask, { name: "John" }).success).toBe(false);
});

test("z.exactPartial", () => {
  const partialSchema = z.exactPartial(userSchema);
  type PartialUser = z.infer<typeof partialSchema>;
  expectTypeOf<PartialUser>().toEqualTypeOf<{
    name?: string;
    age?: number;
    email?: string | undefined;
  }>();
  expect(z.safeParse(partialSchema, { name: "John" }).success).toBe(true);
  expect(z.safeParse(partialSchema, {}).success).toBe(true);
  expect(z.safeParse(partialSchema, { name: undefined }).success).toBe(false);
  expect(z.safeParse(partialSchema, { email: undefined }).success).toBe(true);
});

test("z.exactPartial with mask", () => {
  const partialSchemaWithMask = z.exactPartial(userSchema, { name: true });
  type PartialUserWithMask = z.infer<typeof partialSchemaWithMask>;
  expectTypeOf<PartialUserWithMask>().toEqualTypeOf<{
    name?: string;
    age: number;
    email?: string | undefined;
  }>();
  expect(z.safeParse(partialSchemaWithMask, { age: 30 }).success).toBe(true);
  expect(z.safeParse(partialSchemaWithMask, { age: 30, name: undefined }).success).toBe(false);
});

test("z.required", () => {
  const requiredSchema = z.required(userSchema);
  type RequiredUser = z.infer<typeof requiredSchema>;
  expectTypeOf<RequiredUser>().toEqualTypeOf<{ name: string; age: number; email: string }>();
  expect(requiredSchema.shape.email).toBeInstanceOf(z.ZodMiniNonOptional);
  expect(z.safeParse(requiredSchema, { name: "John", age: 30 }).success).toBe(false);
  expect(z.safeParse(requiredSchema, { name: "John", age: 30, email: "john@example.com" }).success).toBe(true);
});

test("z.required with mask", () => {
  const schema = z.object({ name: z.optional(z.string()), age: z.optional(z.number()) });
  const requiredSchemaWithMask = z.required(schema, { name: true });
  type RequiredWithMask = z.infer<typeof requiredSchemaWithMask>;
  expectTypeOf<RequiredWithMask>().toEqualTypeOf<{ name: string; age?: number | undefined }>();
  expect(requiredSchemaWithMask.shape.name).toBeInstanceOf(z.ZodMiniNonOptional);
  expect(requiredSchemaWithMask.shape.age).toBeInstanceOf(z.ZodMiniOptional);
});

test("z.pick/omit/partial/required - do not allow unknown keys", () => {
  const schema = z.object({
    name: z.string(),
    age: z.number(),
  });

  // Mixed valid + invalid keys - throws at parse time (lazy evaluation)
  // @ts-expect-error
  expect(() => z.parse(z.pick(schema, { name: true, asdf: true }), {})).toThrow();
  // @ts-expect-error
  expect(() => z.parse(z.omit(schema, { name: true, asdf: true }), {})).toThrow();
  // @ts-expect-error
  expect(() => z.parse(z.partial(schema, { name: true, asdf: true }), {})).toThrow();
  // @ts-expect-error
  expect(() => z.parse(z.required(schema, { name: true, asdf: true }), {})).toThrow();

  // Only invalid keys
  // @ts-expect-error
  expect(() => z.parse(z.pick(schema, { $unknown: true }), {})).toThrow();
  // @ts-expect-error
  expect(() => z.parse(z.omit(schema, { $unknown: true }), {})).toThrow();
  // @ts-expect-error
  expect(() => z.parse(z.partial(schema, { $unknown: true }), {})).toThrow();
  // @ts-expect-error
  expect(() => z.parse(z.required(schema, { $unknown: true }), {})).toThrow();
});

test("z.catchall", () => {
  // z.catchall()
  const schema = z.catchall(
    z.object({
      name: z.string(),
      // age: z.number(),
    }),
    z.string()
  );

  type schemaIn = z.input<typeof schema>;
  type schemaOut = z.output<typeof schema>;
  expectTypeOf<schemaIn>().toEqualTypeOf<{
    name: string;
    [key: string]: string;
  }>();

  expectTypeOf<schemaOut>().toEqualTypeOf<{
    name: string;
    [key: string]: string;
  }>();

  schema.parse({
    name: "john",
    age: "30",
    extra: "extra value",
  });

  expect(() => schema.parse({ name: "john", age: 30 })).toThrow();
});