import { CoValueUniqueness } from "cojson";
import {
  Account,
  BranchDefinition,
  CoMap,
  DiscriminableCoValueSchemaDefinition,
  DiscriminableCoreCoValueSchema,
  Group,
  Settled,
  RefsToResolve,
  RefsToResolveStrict,
  Resolved,
  Simplify,
  SubscribeCallback,
  SubscribeListenerOptions,
  coOptionalDefiner,
  hydrateCoreCoValueSchema,
  isAnyCoValueSchema,
  unstable_mergeBranchWithResolve,
  withSchemaPermissions,
  isCoValueSchema,
  type Schema,
  CoValueCreateOptions,
  CoValueCursor,
  LoadCoValueCursorOption,
} from "../../../internal.js";
import { AnonymousJazzAgent } from "../../anonymousJazzAgent.js";
import { removeGetters, withSchemaResolveQuery } from "../../schemaUtils.js";
import { CoMapSchemaInit } from "../typeConverters/CoFieldSchemaInit.js";
import { InstanceOrPrimitiveOfSchema } from "../typeConverters/InstanceOrPrimitiveOfSchema.js";
import { InstanceOrPrimitiveOfSchemaCoValuesMaybeLoaded } from "../typeConverters/InstanceOrPrimitiveOfSchemaCoValuesMaybeLoaded.js";
import { z } from "../zodReExport.js";
import { AnyZodOrCoValueSchema, AnyZodSchema } from "../zodSchema.js";
import { CoOptionalSchema } from "./CoOptionalSchema.js";
import { CoreCoValueSchema, CoreResolveQuery } from "./CoValueSchema.js";
import {
  DEFAULT_SCHEMA_PERMISSIONS,
  SchemaPermissions,
} from "../schemaPermissions.js";
import {
  coValueValidationSchema,
  generateValidationSchemaFromItem,
} from "./schemaValidators.js";
import { resolveSchemaField } from "../runtimeConverters/schemaFieldToCoFieldDef.js";

type CoMapSchemaInstance<Shape extends z.core.$ZodLooseShape> = Simplify<
  CoMapInstanceCoValuesMaybeLoaded<Shape>
> &
  CoMap;

export class CoMapSchema<
  Shape extends z.core.$ZodLooseShape,
  CatchAll extends AnyZodOrCoValueSchema | unknown = unknown,
  Owner extends Account | Group = Account | Group,
  DefaultResolveQuery extends CoreResolveQuery = true,
> implements CoreCoMapSchema<Shape, CatchAll>
{
  collaborative = true as const;
  builtin = "CoMap" as const;
  shape: Shape;
  catchAll?: CatchAll;
  #descriptorsSchema: CoMapDescriptorsSchema | undefined = undefined;
  getDefinition: () => CoMapSchemaDefinition;

  #validationSchema: z.ZodType | undefined = undefined;
  getValidationSchema = () => {
    if (this.#validationSchema) {
      return this.#validationSchema;
    }

    const plainShape: Record<string, z.ZodTypeAny> = {};

    for (const key in this.shape) {
      const item = this.shape[key];
      if (isCoValueSchema(item)) {
        // Inject as getter to avoid circularity issues
        Object.defineProperty(plainShape, key, {
          get: () => generateValidationSchemaFromItem(item),
          enumerable: true,
          configurable: true,
        });
      } else {
        plainShape[key] = generateValidationSchemaFromItem(item);
      }
    }

    let validationSchema = z.strictObject(plainShape);
    if (this.catchAll) {
      validationSchema = validationSchema.catchall(
        generateValidationSchemaFromItem(
          this.catchAll as unknown as AnyZodOrCoValueSchema,
        ),
      );
    }

    this.#validationSchema = coValueValidationSchema(validationSchema, CoMap);

    return this.#validationSchema;
  };

  /**
   * Default resolve query to be used when loading instances of this schema.
   * This resolve query will be used when no resolve query is provided to the load method.
   * @default true
   */
  resolveQuery: DefaultResolveQuery = true as DefaultResolveQuery;

  #permissions: SchemaPermissions | null = null;
  /**
   * Permissions to be used when creating or composing CoValues
   * @internal
   */
  get permissions(): SchemaPermissions {
    return this.#permissions ?? DEFAULT_SCHEMA_PERMISSIONS;
  }

  constructor(
    coreSchema: CoreCoMapSchema<Shape, CatchAll>,
    private coValueClass: typeof CoMap,
  ) {
    this.shape = coreSchema.shape;
    this.catchAll = coreSchema.catchAll;
    this.getDefinition = coreSchema.getDefinition;
  }

  getDescriptorsSchema = (): CoMapDescriptorsSchema => {
    if (this.#descriptorsSchema) {
      return this.#descriptorsSchema;
    }

    const descriptorShape: Record<string, Schema> = {};
    for (const key of Object.keys(this.shape)) {
      const field = this.shape[key as keyof Shape];
      descriptorShape[key] = resolveSchemaField(field as any);
    }

    const descriptorCatchall =
      this.catchAll === undefined
        ? undefined
        : resolveSchemaField(this.catchAll as any);

    this.#descriptorsSchema = {
      shape: descriptorShape,
      catchall: descriptorCatchall,
    };

    return this.#descriptorsSchema;
  };

  create(
    init: CoMapSchemaInit<Shape>,
    options?: CoValueCreateOptions,
  ): CoMapInstanceShape<Shape, CatchAll> & CoMap;
  /** @deprecated Creating CoValues with an Account as owner is deprecated. Use a Group instead. */
  create(
    init: CoMapSchemaInit<Shape>,
    options?: CoValueCreateOptions<{}, Account | Group>,
  ): CoMapInstanceShape<Shape, CatchAll> & CoMap;
  create(init: any, options?: any) {
    const optionsWithPermissions = withSchemaPermissions(
      options,
      this.permissions,
    );

    return this.coValueClass.create(init, optionsWithPermissions);
  }

  load<
    const R extends RefsToResolve<
      Simplify<CoMapInstanceCoValuesMaybeLoaded<Shape>> & CoMap
      // @ts-expect-error
    > = DefaultResolveQuery,
  >(
    id: string,
    options?: {
      resolve?: RefsToResolveStrict<
        Simplify<CoMapInstanceCoValuesMaybeLoaded<Shape>> & CoMap,
        R
      >;
      loadAs?: Account | AnonymousJazzAgent;
      skipRetry?: boolean;
      unstable_branch?: BranchDefinition;
      cursor?: LoadCoValueCursorOption;
    },
  ): Promise<
    Settled<
      Resolved<Simplify<CoMapInstanceCoValuesMaybeLoaded<Shape>> & CoMap, R>
    >
  > {
    // @ts-expect-error
    return this.coValueClass.load(
      id,
      // @ts-expect-error
      withSchemaResolveQuery(options, this.resolveQuery),
    );
  }

  unstable_merge<
    const R extends RefsToResolve<
      Simplify<CoMapInstanceCoValuesMaybeLoaded<Shape>> & CoMap
      // @ts-expect-error we can't statically enforce the schema's resolve query is a valid resolve query, but in practice it is
    > = DefaultResolveQuery,
  >(
    id: string,
    options: {
      resolve?: RefsToResolveStrict<
        Simplify<CoMapInstanceCoValuesMaybeLoaded<Shape>> & CoMap,
        R
      >;
      loadAs?: Account | AnonymousJazzAgent;
      branch: BranchDefinition;
    },
  ): Promise<void> {
    return unstable_mergeBranchWithResolve(
      this.coValueClass,
      id,
      // @ts-expect-error
      withSchemaResolveQuery(options, this.resolveQuery),
    );
  }

  subscribe<
    const R extends RefsToResolve<
      CoMapSchemaInstance<Shape>
      // @ts-expect-error we can't statically enforce the schema's resolve query is a valid resolve query, but in practice it is
    > = DefaultResolveQuery,
  >(
    id: string,
    listener: SubscribeCallback<Resolved<CoMapSchemaInstance<Shape>, R>>,
  ): () => void;
  subscribe<
    const R extends RefsToResolve<
      CoMapSchemaInstance<Shape>
      // @ts-expect-error we can't statically enforce the schema's resolve query is a valid resolve query, but in practice it is
    > = DefaultResolveQuery,
  >(
    id: string,
    options: SubscribeListenerOptions<CoMapSchemaInstance<Shape>, R>,
    listener: SubscribeCallback<Resolved<CoMapSchemaInstance<Shape>, R>>,
  ): () => void;
  subscribe<const R extends RefsToResolve<CoMapSchemaInstance<Shape>>>(
    id: string,
    optionsOrListener:
      | SubscribeListenerOptions<CoMapSchemaInstance<Shape>, R>
      | SubscribeCallback<Resolved<CoMapSchemaInstance<Shape>, R>>,
    maybeListener?: SubscribeCallback<Resolved<CoMapSchemaInstance<Shape>, R>>,
  ): () => void {
    if (typeof optionsOrListener === "function") {
      // @ts-expect-error
      return this.coValueClass.subscribe(
        id,
        withSchemaResolveQuery({}, this.resolveQuery),
        optionsOrListener,
      );
    }
    // @ts-expect-error
    return this.coValueClass.subscribe(
      id,
      withSchemaResolveQuery(optionsOrListener, this.resolveQuery),
      maybeListener,
    );
  }

  /** @deprecated Use `loadUnique` instead. */
  findUnique(
    unique: CoValueUniqueness["uniqueness"],
    ownerID: string,
    as?: Account | Group | AnonymousJazzAgent,
  ): string {
    return this.coValueClass.findUnique(unique, ownerID, as);
  }

  /**
   * Get an existing unique CoMap or create a new one if it doesn't exist.
   *
   * Unlike `upsertUnique`, this method does NOT update existing values with the provided value.
   * The provided value is only used when creating a new CoMap.
   *
   * @example
   * ```ts
   * const settings = await UserSettings.getOrCreateUnique({
   *   value: { theme: "dark", language: "en" },
   *   unique: "user-settings",
   *   owner: me,
   * });
   * ```
   *
   * @param options The options for creating or loading the CoMap.
   * @returns Either an existing CoMap (unchanged), or a new initialised CoMap if none exists.
   * @category Subscription & Loading
   */
  getOrCreateUnique<
    const R extends RefsToResolve<
      Simplify<CoMapInstanceCoValuesMaybeLoaded<Shape>> & CoMap
      // @ts-expect-error we can't statically enforce the schema's resolve query is a valid resolve query, but in practice it is
    > = DefaultResolveQuery,
  >(options: {
    value: Simplify<CoMapSchemaInit<Shape>>;
    unique: CoValueUniqueness["uniqueness"];
    owner: Owner;
    resolve?: RefsToResolveStrict<
      Simplify<CoMapInstanceCoValuesMaybeLoaded<Shape>> & CoMap,
      R
    >;
  }): Promise<
    Settled<
      Resolved<Simplify<CoMapInstanceCoValuesMaybeLoaded<Shape>> & CoMap, R>
    >
  > {
    // @ts-expect-error
    return this.coValueClass.getOrCreateUnique(
      // @ts-expect-error
      withSchemaResolveQuery(options, this.resolveQuery),
    );
  }

  /**
   * @deprecated Use `getOrCreateUnique` instead. Note: getOrCreateUnique does not update existing values.
   * If you need to update, use getOrCreateUnique followed by direct property assignment.
   */
  upsertUnique<
    const R extends RefsToResolve<
      Simplify<CoMapInstanceCoValuesMaybeLoaded<Shape>> & CoMap
      // @ts-expect-error we can't statically enforce the schema's resolve query is a valid resolve query, but in practice it is
    > = DefaultResolveQuery,
  >(options: {
    value: Simplify<CoMapSchemaInit<Shape>>;
    unique: CoValueUniqueness["uniqueness"];
    owner: Owner;
    resolve?: RefsToResolveStrict<
      Simplify<CoMapInstanceCoValuesMaybeLoaded<Shape>> & CoMap,
      R
    >;
  }): Promise<
    Settled<
      Resolved<Simplify<CoMapInstanceCoValuesMaybeLoaded<Shape>> & CoMap, R>
    >
  > {
    // @ts-expect-error
    return this.coValueClass.upsertUnique(
      // @ts-expect-error
      withSchemaResolveQuery(options, this.resolveQuery),
    );
  }

  loadUnique<
    const R extends RefsToResolve<
      Simplify<CoMapInstanceCoValuesMaybeLoaded<Shape>> & CoMap
      // @ts-expect-error we can't statically enforce the schema's resolve query is a valid resolve query, but in practice it is
    > = DefaultResolveQuery,
  >(
    unique: CoValueUniqueness["uniqueness"],
    ownerID: string,
    options?: {
      resolve?: RefsToResolveStrict<
        Simplify<CoMapInstanceCoValuesMaybeLoaded<Shape>> & CoMap,
        R
      >;
      loadAs?: Account | AnonymousJazzAgent;
    },
  ): Promise<
    Settled<
      Resolved<Simplify<CoMapInstanceCoValuesMaybeLoaded<Shape>> & CoMap, R>
    >
  > {
    // @ts-expect-error
    return this.coValueClass.loadUnique(
      unique,
      ownerID,
      // @ts-expect-error
      withSchemaResolveQuery(options, this.resolveQuery),
    );
  }

  /**
   * @deprecated `co.map().catchall` will be removed in an upcoming version.
   *
   * Use a `co.record` nested inside a `co.map` if you need to store key-value properties.
   *
   * @example
   * ```ts
   * // Instead of:
   * const Image = co.map({
   *   original: co.fileStream(),
   * }).catchall(co.fileStream());
   *
   * // Use:
   * const Image = co.map({
   *   original: co.fileStream(),
   *   resolutions: co.record(z.string(), co.fileStream()),
   * });
   * ```
   */
  catchall<T extends AnyZodOrCoValueSchema>(schema: T): CoMapSchema<Shape, T> {
    const schemaWithCatchAll = createCoreCoMapSchema(this.shape, schema);
    return hydrateCoreCoValueSchema(schemaWithCatchAll);
  }

  withMigration(migration: CoMapMigration<Shape>): this {
    // @ts-expect-error avoid exposing 'migrate' at the type level
    this.coValueClass.prototype.migrate = migration;
    return this;
  }

  getCoValueClass(): typeof CoMap {
    return this.coValueClass;
  }

  optional(): CoOptionalSchema<this> {
    return coOptionalDefiner(this);
  }

  /**
   * Creates a new CoMap schema by picking the specified keys from the original schema.
   * The new CoMap will **not** inherit any configurations (migration, catchall, permissions, default resolve query)
   *
   * @param keys - The keys to pick from the original schema.
   * @returns A new CoMap schema with the picked keys.
   */
  pick<Keys extends keyof Shape>(
    keys: { [key in Keys]: true },
  ): CoMapSchema<Simplify<Pick<Shape, Keys>>, unknown, Owner> {
    const keysSet = new Set(Object.keys(keys));
    const pickedShape: Record<string, AnyZodOrCoValueSchema> = {};

    for (const [key, value] of Object.entries(this.shape)) {
      if (keysSet.has(key)) {
        pickedShape[key] = value;
      }
    }

    return this.copy(
      {
        shape: pickedShape as Pick<Shape, Keys>,
        permissions: DEFAULT_SCHEMA_PERMISSIONS,
        resolveQuery: true,
      },
      {
        shouldCopyCatchAll: false,
        shouldCopyMigration: false,
      },
    );
  }

  /**
   * Creates a new CoMap schema by omitting the specified keys from the original schema.
   * The new CoMap will **not** inherit any configurations (migration, catchall, permissions, default resolve query)
   *
   * @param keys - The keys to omit from the original schema.
   * @returns A new CoMap schema with the omitted keys.
   */
  omit<Keys extends keyof Shape>(
    keys: { [key in Keys]: true },
  ): CoMapSchema<Simplify<Omit<Shape, Keys>>, unknown, Owner> {
    const newShape: Record<string, AnyZodOrCoValueSchema> = { ...this.shape };

    for (const key in keys) {
      delete newShape[key];
    }

    return this.copy(
      {
        shape: newShape as Omit<Shape, Keys>,
        permissions: DEFAULT_SCHEMA_PERMISSIONS,
        resolveQuery: true,
      },
      {
        shouldCopyCatchAll: false,
        shouldCopyMigration: false,
      },
    );
  }

  /**
   * Creates a new CoMap schema by making all fields optional.
   * The new CoMap will **not** inherit any configurations (migration, permissions, default resolve query) except catchall
   *
   * @returns A new CoMap schema with all fields optional.
   */
  partial<Keys extends keyof Shape = keyof Shape>(
    keys?: {
      [key in Keys]: true;
    },
  ): CoMapSchema<PartialShape<Shape, Keys>, CatchAll, Owner> {
    const partialShape: Record<string, AnyZodOrCoValueSchema> = {};

    for (const [key, value] of Object.entries(this.shape)) {
      if (keys && !keys[key as Keys]) {
        partialShape[key] = value;
        continue;
      }

      if (isAnyCoValueSchema(value)) {
        partialShape[key] = coOptionalDefiner(value);
      } else {
        partialShape[key] = z.optional(this.shape[key]);
      }
    }

    return this.copy(
      {
        shape: partialShape as PartialShape<Shape, Keys>,
        permissions: DEFAULT_SCHEMA_PERMISSIONS,
        resolveQuery: true,
      },
      {
        shouldCopyMigration: false,
      },
    );
  }

  /**
   * Creates a new CoMap schema by extending the current schema with additional fields.
   * The new CoMap **will** inherit all configurations (migration, catchall, permissions, default resolve query)
   *
   * @param shape - The shape object with additional fields to add to the schema.
   * @returns A new CoMap schema with the additional fields.
   */
  extend<ExtendShape extends z.core.$ZodLooseShape>(
    shape: ExtendShape,
  ): CoMapSchema<
    z.core.util.Extend<Shape, ExtendShape>,
    CatchAll,
    Owner,
    DefaultResolveQuery
  > {
    return this.copy({
      shape: {
        ...this.shape,
        ...shape,
      } as z.core.util.Extend<Shape, ExtendShape>,
    });
  }

  /**
   * Creates a new CoMap schema by extending the current schema with additional fields,
   * while also making sure that no existing fields are accidentally overridden.
   * The new CoMap **will** inherit all configurations (migration, catchall, permissions, default resolve query)
   *
   * @param shape - The shape object with additional fields to add to the schema.
   * @returns A new CoMap schema with the additional fields.
   */
  safeExtend<ExtendShape extends z.core.$ZodLooseShape>(
    shape: SafeExtendShape<Shape, ExtendShape> &
      Partial<Record<keyof Shape, AnyZodOrCoValueSchema>>,
  ): CoMapSchema<
    z.core.util.Extend<Shape, ExtendShape>,
    CatchAll,
    Owner,
    DefaultResolveQuery
  > {
    return this.copy({
      shape: {
        ...this.shape,
        ...shape,
      } as z.core.util.Extend<Shape, ExtendShape>,
    });
  }

  /**
   * Adds a default resolve query to be used when loading instances of this schema.
   * This resolve query will be used when no resolve query is provided to the load method.
   */
  resolved<
    const R extends RefsToResolve<
      Simplify<CoMapInstanceCoValuesMaybeLoaded<Shape>> & CoMap
    > = true,
  >(
    resolveQuery: RefsToResolveStrict<
      Simplify<CoMapInstanceCoValuesMaybeLoaded<Shape>> & CoMap,
      R
    >,
  ): CoMapSchema<Shape, CatchAll, Owner, R> {
    return this.copy({ resolveQuery: resolveQuery as R });
  }

  /**
   * Configure permissions to be used when creating or composing CoValues
   */
  withPermissions(
    permissions: Omit<SchemaPermissions, "writer">,
  ): CoMapSchema<Shape, CatchAll, Owner, DefaultResolveQuery> {
    return this.copy({ permissions });
  }

  /**
   * Creates a copy of this schema, preserving all previous configuration
   */
  private copy<
    ShapeOverride extends z.core.$ZodLooseShape = Shape,
    ResolveQuery extends CoreResolveQuery = DefaultResolveQuery,
  >(
    {
      shape,
      permissions,
      resolveQuery,
    }: {
      shape?: ShapeOverride;
      permissions?: SchemaPermissions;
      resolveQuery?: ResolveQuery;
    },
    options: {
      shouldCopyCatchAll?: boolean;
      shouldCopyMigration?: boolean;
    } = {},
  ): CoMapSchema<ShapeOverride, CatchAll, Owner, ResolveQuery> {
    const { shouldCopyCatchAll = true, shouldCopyMigration = true } = options;

    const coreSchema = createCoreCoMapSchema(
      shape ?? this.shape,
      shouldCopyCatchAll ? this.catchAll : (undefined as unknown),
    );

    const copy = hydrateCoreCoValueSchema(coreSchema) as CoMapSchema<
      ShapeOverride,
      CatchAll,
      Owner,
      ResolveQuery
    >;

    if (shouldCopyMigration) {
      // @ts-expect-error avoid exposing 'migrate' at the type level
      copy.coValueClass.prototype.migrate = this.coValueClass.prototype.migrate;
    }

    copy.#permissions = permissions ?? this.#permissions;
    copy.resolveQuery = (resolveQuery ?? this.resolveQuery) as ResolveQuery;
    return copy;
  }
}

export function createCoreCoMapSchema<
  Shape extends z.core.$ZodLooseShape,
  CatchAll extends AnyZodOrCoValueSchema | unknown = unknown,
>(shape: Shape, catchAll?: CatchAll): CoreCoMapSchema<Shape, CatchAll> {
  let descriptorsSchema: CoMapDescriptorsSchema | undefined;

  return {
    collaborative: true as const,
    builtin: "CoMap" as const,
    shape,
    catchAll,
    getDescriptorsSchema: () => {
      if (descriptorsSchema) {
        return descriptorsSchema;
      }

      const descriptorShape: Record<string, Schema> = {};
      for (const key of Object.keys(shape)) {
        const field = shape[key as keyof Shape];
        descriptorShape[key] = resolveSchemaField(field as any);
      }

      const descriptorCatchall =
        catchAll === undefined
          ? undefined
          : resolveSchemaField(catchAll as any);

      descriptorsSchema = {
        shape: descriptorShape,
        catchall: descriptorCatchall,
      };

      return descriptorsSchema;
    },
    getDefinition: () => ({
      get shape() {
        return shape;
      },
      get catchall() {
        return catchAll;
      },
      get discriminatorMap() {
        const propValues: DiscriminableCoValueSchemaDefinition["discriminatorMap"] =
          {};
        // remove getters to avoid circularity issues. Getters are not used as discriminators
        for (const key in removeGetters(shape)) {
          if (isAnyCoValueSchema(shape[key])) {
            // CoValues cannot be used as discriminators either
            continue;
          }
          const field = shape[key]._zod;
          if (field.values) {
            propValues[key] ??= new Set();
            for (const v of field.values) propValues[key].add(v);
          }
        }
        return propValues;
      },
    }),
    resolveQuery: true as const,
    getValidationSchema: () => z.any(),
  };
}

export interface CoMapSchemaDefinition<
  Shape extends z.core.$ZodLooseShape = z.core.$ZodLooseShape,
  CatchAll extends AnyZodOrCoValueSchema | unknown = unknown,
> extends DiscriminableCoValueSchemaDefinition {
  shape: Shape;
  catchall?: CatchAll;
}

export type CoMapDescriptorsSchema = {
  shape: Record<string, Schema>;
  catchall?: Schema;
};

// less precise version to avoid circularity issues and allow matching against
export interface CoreCoMapSchema<
  Shape extends z.core.$ZodLooseShape = z.core.$ZodLooseShape,
  CatchAll extends AnyZodOrCoValueSchema | unknown = unknown,
> extends DiscriminableCoreCoValueSchema {
  builtin: "CoMap";
  shape: Shape;
  catchAll?: CatchAll;
  getDescriptorsSchema: () => CoMapDescriptorsSchema;
  getDefinition: () => CoMapSchemaDefinition;
}

type CoMapMigration<Shape extends z.core.$ZodLooseShape> = (
  value: Resolved<
    Simplify<CoMapInstanceCoValuesMaybeLoaded<Shape>> & CoMap,
    true
  >,
) => undefined;

export type CoMapInstanceShape<
  Shape extends z.core.$ZodLooseShape,
  CatchAll extends AnyZodOrCoValueSchema | unknown = unknown,
> = {
  readonly [key in keyof Shape]: InstanceOrPrimitiveOfSchema<Shape[key]>;
} & (CatchAll extends AnyZodOrCoValueSchema
  ? {
      readonly [key: string]: InstanceOrPrimitiveOfSchema<CatchAll>;
    }
  : {});

export type CoMapInstanceCoValuesMaybeLoaded<
  Shape extends z.core.$ZodLooseShape,
> = {
  readonly [key in keyof Shape]: InstanceOrPrimitiveOfSchemaCoValuesMaybeLoaded<
    Shape[key]
  >;
};

export type PartialShape<
  Shape extends z.core.$ZodLooseShape,
  PartialKeys extends keyof Shape = keyof Shape,
> = Simplify<{
  -readonly [key in keyof Shape]: key extends PartialKeys
    ? Shape[key] extends AnyZodSchema
      ? z.ZodOptional<Shape[key]>
      : Shape[key] extends CoreCoValueSchema
        ? CoOptionalSchema<Shape[key]>
        : never
    : Shape[key];
}>;

export type SafeExtendShape<
  Base extends z.core.$ZodLooseShape,
  Ext extends z.core.$ZodLooseShape,
> = {
  [K in keyof Ext]: K extends keyof Base
    ? Base[K] extends z.core.SomeType
      ? Ext[K] extends z.core.SomeType
        ? InstanceOrPrimitiveOfSchema<
            Ext[K]
          > extends InstanceOrPrimitiveOfSchema<Base[K]>
          ? z.core.input<Ext[K]> extends z.core.input<Base[K]>
            ? Ext[K]
            : never
          : never
        : never
      : Ext[K] extends z.core.SomeType
        ? never
        : // both are CoValue schemas
          InstanceOrPrimitiveOfSchema<
              Ext[K]
            > extends InstanceOrPrimitiveOfSchema<Base[K]>
          ? Ext[K]
          : never
    : Ext[K];
};