import type * as core from "../core/index.js";
import type * as JSONSchema from "./json-schema.js";
import { type $ZodRegistry, globalRegistry } from "./registries.js";
import type * as schemas from "./schemas.js";
import type { StandardJSONSchemaV1, StandardSchemaWithJSONProps } from "./standard-schema.js";

export type Processor<T extends schemas.$ZodType = schemas.$ZodType> = (
  schema: T,
  ctx: ToJSONSchemaContext,
  json: JSONSchema.BaseSchema,
  params: ProcessParams
) => void;

export interface JSONSchemaGeneratorParams {
  processors: Record<string, Processor>;
  /** A registry used to look up metadata for each schema. Any schema with an `id` property will be extracted as a $def.
   *  @default globalRegistry */
  metadata?: $ZodRegistry<Record<string, any>>;
  /** The JSON Schema version to target.
   * - `"draft-2020-12"` — Default. JSON Schema Draft 2020-12
   * - `"draft-07"` — JSON Schema Draft 7
   * - `"draft-04"` — JSON Schema Draft 4
   * - `"openapi-3.0"` — OpenAPI 3.0 Schema Object */
  target?: "draft-04" | "draft-07" | "draft-2020-12" | "openapi-3.0" | ({} & string) | undefined;
  /** How to handle unrepresentable types.
   * - `"throw"` — Default. Unrepresentable types throw an error
   * - `"any"` — Unrepresentable types become `{}` */
  unrepresentable?: "throw" | "any";
  /** Arbitrary custom logic that can be used to modify the generated JSON Schema. */
  override?: (ctx: {
    zodSchema: schemas.$ZodTypes;
    jsonSchema: JSONSchema.BaseSchema;
    path: (string | number)[];
  }) => void;
  /** Whether to extract the `"input"` or `"output"` type. Relevant to transforms, defaults, coerced primitives, etc.
   * - `"output"` — Default. Convert the output schema.
   * - `"input"` — Convert the input schema. */
  io?: "input" | "output";
  cycles?: "ref" | "throw";
  reused?: "ref" | "inline";
  external?:
    | {
        registry: $ZodRegistry<{ id?: string | undefined }>;
        uri?: ((id: string) => string) | undefined;
        defs: Record<string, JSONSchema.BaseSchema>;
      }
    | undefined;
}

/**
 * Parameters for the toJSONSchema function.
 */
export type ToJSONSchemaParams = Omit<JSONSchemaGeneratorParams, "processors" | "external">;

/**
 * Parameters for the toJSONSchema function when passing a registry.
 */
export interface RegistryToJSONSchemaParams extends ToJSONSchemaParams {
  uri?: (id: string) => string;
}

export interface ProcessParams {
  schemaPath: schemas.$ZodType[];
  path: (string | number)[];
}

export interface Seen {
  /** JSON Schema result for this Zod schema */
  schema: JSONSchema.BaseSchema;
  /** A cached version of the schema that doesn't get overwritten during ref resolution */
  def?: JSONSchema.BaseSchema;
  defId?: string | undefined;
  /** Number of times this schema was encountered during traversal */
  count: number;
  /** Cycle path */
  cycle?: (string | number)[] | undefined;
  isParent?: boolean | undefined;
  /** Schema to inherit JSON Schema properties from (set by processor for wrappers) */
  ref?: schemas.$ZodType | null;
  /** JSON Schema property path for this schema */
  path?: (string | number)[] | undefined;
}

export interface ToJSONSchemaContext {
  processors: Record<string, Processor>;
  metadataRegistry: $ZodRegistry<Record<string, any>>;
  target: "draft-04" | "draft-07" | "draft-2020-12" | "openapi-3.0" | ({} & string);
  unrepresentable: "throw" | "any";
  override: (ctx: {
    // must be schemas.$ZodType to prevent recursive type resolution error
    zodSchema: schemas.$ZodType;
    jsonSchema: JSONSchema.BaseSchema;
    path: (string | number)[];
  }) => void;
  io: "input" | "output";
  counter: number;
  seen: Map<schemas.$ZodType, Seen>;
  cycles: "ref" | "throw";
  reused: "ref" | "inline";
  external?:
    | {
        registry: $ZodRegistry<{ id?: string | undefined }>;
        uri?: ((id: string) => string) | undefined;
        defs: Record<string, JSONSchema.BaseSchema>;
      }
    | undefined;
}

// function initializeContext<T extends schemas.$ZodType>(inputs: JSONSchemaGeneratorParams<T>): ToJSONSchemaContext<T> {
//   return {
//     processor: inputs.processor,
//     metadataRegistry: inputs.metadata ?? globalRegistry,
//     target: inputs.target ?? "draft-2020-12",
//     unrepresentable: inputs.unrepresentable ?? "throw",
//   };
// }

export function initializeContext(params: JSONSchemaGeneratorParams): ToJSONSchemaContext {
  // Normalize target: convert old non-hyphenated versions to hyphenated versions
  let target: ToJSONSchemaContext["target"] = params?.target ?? "draft-2020-12";
  if (target === "draft-4") target = "draft-04";
  if (target === "draft-7") target = "draft-07";

  return {
    processors: params.processors ?? {},
    metadataRegistry: params?.metadata ?? globalRegistry,
    target,
    unrepresentable: params?.unrepresentable ?? "throw",
    override: (params?.override as any) ?? (() => {}),
    io: params?.io ?? "output",
    counter: 0,
    seen: new Map(),
    cycles: params?.cycles ?? "ref",
    reused: params?.reused ?? "inline",
    external: params?.external ?? undefined,
  };
}

export function process<T extends schemas.$ZodType>(
  schema: T,
  ctx: ToJSONSchemaContext,
  _params: ProcessParams = { path: [], schemaPath: [] }
): JSONSchema.BaseSchema {
  const def = schema._zod.def as schemas.$ZodTypes["_zod"]["def"];

  // check for schema in seens
  const seen = ctx.seen.get(schema);

  if (seen) {
    seen.count++;

    // check if cycle
    const isCycle = _params.schemaPath.includes(schema);
    if (isCycle) {
      seen.cycle = _params.path;
    }

    return seen.schema;
  }

  // initialize
  const result: Seen = { schema: {}, count: 1, cycle: undefined, path: _params.path };
  ctx.seen.set(schema, result);

  // custom method overrides default behavior
  const overrideSchema = schema._zod.toJSONSchema?.();
  if (overrideSchema) {
    result.schema = overrideSchema as any;
  } else {
    const params = {
      ..._params,
      schemaPath: [..._params.schemaPath, schema],
      path: _params.path,
    };

    if (schema._zod.processJSONSchema) {
      schema._zod.processJSONSchema(ctx, result.schema, params);
    } else {
      const _json = result.schema;
      const processor = ctx.processors[def.type];
      if (!processor) {
        throw new Error(`[toJSONSchema]: Non-representable type encountered: ${def.type}`);
      }
      processor(schema, ctx, _json, params);
    }

    const parent = schema._zod.parent as T;

    if (parent) {
      // Also set ref if processor didn't (for inheritance)
      if (!result.ref) result.ref = parent;
      process(parent, ctx, params);
      ctx.seen.get(parent)!.isParent = true;
    }
  }

  // metadata
  const meta = ctx.metadataRegistry.get(schema);
  if (meta) Object.assign(result.schema, meta);

  if (ctx.io === "input" && isTransforming(schema)) {
    // examples/defaults only apply to output type of pipe
    delete result.schema.examples;
    delete result.schema.default;
  }

  // set prefault as default
  if (ctx.io === "input" && result.schema._prefault) result.schema.default ??= result.schema._prefault;
  delete result.schema._prefault;

  // pulling fresh from ctx.seen in case it was overwritten
  const _result = ctx.seen.get(schema)!;

  return _result.schema;
}

export function extractDefs<T extends schemas.$ZodType>(
  ctx: ToJSONSchemaContext,
  schema: T
  // params: EmitParams
): void {
  // iterate over seen map;
  const root = ctx.seen.get(schema);

  if (!root) throw new Error("Unprocessed schema. This is a bug in Zod.");

  // Track ids to detect duplicates across different schemas
  const idToSchema = new Map<string, schemas.$ZodType>();
  for (const entry of ctx.seen.entries()) {
    const id = ctx.metadataRegistry.get(entry[0])?.id;
    if (id) {
      const existing = idToSchema.get(id);
      if (existing && existing !== entry[0]) {
        throw new Error(
          `Duplicate schema id "${id}" detected during JSON Schema conversion. Two different schemas cannot share the same id when converted together.`
        );
      }
      idToSchema.set(id, entry[0]);
    }
  }

  // returns a ref to the schema
  // defId will be empty if the ref points to an external schema (or #)
  const makeURI = (entry: [schemas.$ZodType<unknown, unknown>, Seen]): { ref: string; defId?: string } => {
    // comparing the seen objects because sometimes
    // multiple schemas map to the same seen object.
    // e.g. lazy

    // external is configured
    const defsSegment = ctx.target === "draft-2020-12" ? "$defs" : "definitions";
    if (ctx.external) {
      const externalId = ctx.external.registry.get(entry[0])?.id; // ?? "__shared";// `__schema${ctx.counter++}`;

      // check if schema is in the external registry
      const uriGenerator = ctx.external.uri ?? ((id: string) => id);
      if (externalId) {
        return { ref: uriGenerator(externalId) };
      }

      // otherwise, add to __shared
      const id: string = entry[1].defId ?? (entry[1].schema.id as string) ?? `schema${ctx.counter++}`;
      entry[1].defId = id; // set defId so it will be reused if needed
      return { defId: id, ref: `${uriGenerator("__shared")}#/${defsSegment}/${id}` };
    }

    if (entry[1] === root) {
      return { ref: "#" };
    }

    // self-contained schema
    const uriPrefix = `#`;
    const defUriPrefix = `${uriPrefix}/${defsSegment}/`;
    const defId = entry[1].schema.id ?? `__schema${ctx.counter++}`;
    return { defId, ref: defUriPrefix + defId };
  };

  // stored cached version in `def` property
  // remove all properties, set $ref
  const extractToDef = (entry: [schemas.$ZodType<unknown, unknown>, Seen]): void => {
    // if the schema is already a reference, do not extract it
    if (entry[1].schema.$ref) {
      return;
    }
    const seen = entry[1];
    const { ref, defId } = makeURI(entry);

    seen.def = { ...seen.schema };
    // defId won't be set if the schema is a reference to an external schema
    // or if the schema is the root schema
    if (defId) seen.defId = defId;
    // wipe away all properties except $ref
    const schema = seen.schema;
    for (const key in schema) {
      delete schema[key];
    }
    schema.$ref = ref;
  };

  // throw on cycles

  // break cycles
  if (ctx.cycles === "throw") {
    for (const entry of ctx.seen.entries()) {
      const seen = entry[1];
      if (seen.cycle) {
        throw new Error(
          "Cycle detected: " +
            `#/${seen.cycle?.join("/")}/<root>` +
            '\n\nSet the `cycles` parameter to `"ref"` to resolve cyclical schemas with defs.'
        );
      }
    }
  }

  // extract schemas into $defs
  for (const entry of ctx.seen.entries()) {
    const seen = entry[1];

    // convert root schema to # $ref
    if (schema === entry[0]) {
      extractToDef(entry); // this has special handling for the root schema
      continue;
    }

    // extract schemas that are in the external registry
    if (ctx.external) {
      const ext = ctx.external.registry.get(entry[0])?.id;
      if (schema !== entry[0] && ext) {
        extractToDef(entry);
        continue;
      }
    }

    // extract schemas with `id` meta
    const id = ctx.metadataRegistry.get(entry[0])?.id;
    if (id) {
      extractToDef(entry);
      continue;
    }

    // break cycles
    if (seen.cycle) {
      // any
      extractToDef(entry);
      continue;
    }

    // extract reused schemas
    if (seen.count > 1) {
      if (ctx.reused === "ref") {
        extractToDef(entry);
        // biome-ignore lint:
        continue;
      }
    }
  }
}

export function finalize<T extends schemas.$ZodType>(
  ctx: ToJSONSchemaContext,
  schema: T
): ZodStandardJSONSchemaPayload<T> {
  const root = ctx.seen.get(schema);
  if (!root) throw new Error("Unprocessed schema. This is a bug in Zod.");

  // flatten refs - inherit properties from parent schemas
  const flattenRef = (zodSchema: schemas.$ZodType) => {
    const seen = ctx.seen.get(zodSchema)!;

    // already processed
    if (seen.ref === null) return;

    const schema = seen.def ?? seen.schema;
    const _cached = { ...schema };

    const ref = seen.ref;
    seen.ref = null; // prevent infinite recursion

    if (ref) {
      flattenRef(ref);

      const refSeen = ctx.seen.get(ref)!;
      const refSchema = refSeen.schema;

      // merge referenced schema into current
      if (refSchema.$ref && (ctx.target === "draft-07" || ctx.target === "draft-04" || ctx.target === "openapi-3.0")) {
        // older drafts can't combine $ref with other properties
        schema.allOf = schema.allOf ?? [];
        schema.allOf.push(refSchema);
      } else {
        Object.assign(schema, refSchema);
      }
      // restore child's own properties (child wins)
      Object.assign(schema, _cached);

      const isParentRef = zodSchema._zod.parent === ref;

      // For parent chain, child is a refinement - remove parent-only properties
      if (isParentRef) {
        for (const key in schema) {
          if (key === "$ref" || key === "allOf") continue;
          if (!(key in _cached)) {
            delete schema[key];
          }
        }
      }

      // When ref was extracted to $defs, remove properties that match the definition
      if (refSchema.$ref && refSeen.def) {
        for (const key in schema) {
          if (key === "$ref" || key === "allOf") continue;
          if (key in refSeen.def && JSON.stringify(schema[key]) === JSON.stringify(refSeen.def[key])) {
            delete schema[key];
          }
        }
      }
    }

    // If parent was extracted (has $ref), propagate $ref to this schema
    // This handles cases like: readonly().meta({id}).describe()
    // where processor sets ref to innerType but parent should be referenced
    const parent = zodSchema._zod.parent;
    if (parent && parent !== ref) {
      // Ensure parent is processed first so its def has inherited properties
      flattenRef(parent);
      const parentSeen = ctx.seen.get(parent);
      if (parentSeen?.schema.$ref) {
        schema.$ref = parentSeen.schema.$ref;
        // De-duplicate with parent's definition
        if (parentSeen.def) {
          for (const key in schema) {
            if (key === "$ref" || key === "allOf") continue;
            if (key in parentSeen.def && JSON.stringify(schema[key]) === JSON.stringify(parentSeen.def[key])) {
              delete schema[key];
            }
          }
        }
      }
    }

    // execute overrides
    ctx.override({
      zodSchema: zodSchema as schemas.$ZodTypes,
      jsonSchema: schema,
      path: seen.path ?? [],
    });
  };

  for (const entry of [...ctx.seen.entries()].reverse()) {
    flattenRef(entry[0]);
  }

  const result: JSONSchema.BaseSchema = {};
  if (ctx.target === "draft-2020-12") {
    result.$schema = "https://json-schema.org/draft/2020-12/schema";
  } else if (ctx.target === "draft-07") {
    result.$schema = "http://json-schema.org/draft-07/schema#";
  } else if (ctx.target === "draft-04") {
    result.$schema = "http://json-schema.org/draft-04/schema#";
  } else if (ctx.target === "openapi-3.0") {
    // OpenAPI 3.0 schema objects should not include a $schema property
  } else {
    // Arbitrary string values are allowed but won't have a $schema property set
  }

  if (ctx.external?.uri) {
    const id = ctx.external.registry.get(schema)?.id;
    if (!id) throw new Error("Schema is missing an `id` property");
    result.$id = ctx.external.uri(id);
  }

  Object.assign(result, root.def ?? root.schema);

  // build defs object
  const defs: JSONSchema.BaseSchema["$defs"] = ctx.external?.defs ?? {};
  for (const entry of ctx.seen.entries()) {
    const seen = entry[1];
    if (seen.def && seen.defId) {
      defs[seen.defId] = seen.def;
    }
  }

  // set definitions in result
  if (ctx.external) {
  } else {
    if (Object.keys(defs).length > 0) {
      if (ctx.target === "draft-2020-12") {
        result.$defs = defs;
      } else {
        result.definitions = defs;
      }
    }
  }

  try {
    // this "finalizes" this schema and ensures all cycles are removed
    // each call to finalize() is functionally independent
    // though the seen map is shared
    const finalized = JSON.parse(JSON.stringify(result));
    Object.defineProperty(finalized, "~standard", {
      value: {
        ...schema["~standard"],
        jsonSchema: {
          input: createStandardJSONSchemaMethod(schema, "input", ctx.processors),
          output: createStandardJSONSchemaMethod(schema, "output", ctx.processors),
        },
      },
      enumerable: false,
      writable: false,
    });

    return finalized;
  } catch (_err) {
    throw new Error("Error converting schema to JSON.");
  }
}

function isTransforming(
  _schema: schemas.$ZodType,
  _ctx?: {
    seen: Set<schemas.$ZodType>;
  }
): boolean {
  const ctx = _ctx ?? { seen: new Set() };

  if (ctx.seen.has(_schema)) return false;
  ctx.seen.add(_schema);

  const def = (_schema as schemas.$ZodTypes)._zod.def;

  if (def.type === "transform") return true;

  if (def.type === "array") return isTransforming(def.element, ctx);
  if (def.type === "set") return isTransforming(def.valueType, ctx);
  if (def.type === "lazy") return isTransforming(def.getter(), ctx);

  if (
    def.type === "promise" ||
    def.type === "optional" ||
    def.type === "nonoptional" ||
    def.type === "nullable" ||
    def.type === "readonly" ||
    def.type === "default" ||
    def.type === "prefault"
  ) {
    return isTransforming(def.innerType, ctx);
  }

  if (def.type === "intersection") {
    return isTransforming(def.left, ctx) || isTransforming(def.right, ctx);
  }
  if (def.type === "record" || def.type === "map") {
    return isTransforming(def.keyType, ctx) || isTransforming(def.valueType, ctx);
  }
  if (def.type === "pipe") {
    return isTransforming(def.in, ctx) || isTransforming(def.out, ctx);
  }

  if (def.type === "object") {
    for (const key in def.shape) {
      if (isTransforming(def.shape[key]!, ctx)) return true;
    }
    return false;
  }
  if (def.type === "union") {
    for (const option of def.options) {
      if (isTransforming(option, ctx)) return true;
    }
    return false;
  }
  if (def.type === "tuple") {
    for (const item of def.items) {
      if (isTransforming(item, ctx)) return true;
    }
    if (def.rest && isTransforming(def.rest, ctx)) return true;
    return false;
  }

  return false;
}

export type ZodStandardSchemaWithJSON<T> = StandardSchemaWithJSONProps<core.input<T>, core.output<T>>;
export interface ZodStandardJSONSchemaPayload<T> extends JSONSchema.BaseSchema {
  "~standard": ZodStandardSchemaWithJSON<T>;
}

/**
 * Creates a toJSONSchema method for a schema instance.
 * This encapsulates the logic of initializing context, processing, extracting defs, and finalizing.
 */
export const createToJSONSchemaMethod =
  <T extends schemas.$ZodType>(schema: T, processors: Record<string, Processor> = {}) =>
  (params?: ToJSONSchemaParams): ZodStandardJSONSchemaPayload<T> => {
    const ctx = initializeContext({ ...params, processors });
    process(schema, ctx);
    extractDefs(ctx, schema);
    return finalize(ctx, schema);
  };

/**
 * Creates a toJSONSchema method for a schema instance.
 * This encapsulates the logic of initializing context, processing, extracting defs, and finalizing.
 */
type StandardJSONSchemaMethodParams = Parameters<StandardJSONSchemaV1["~standard"]["jsonSchema"]["input"]>[0];
export const createStandardJSONSchemaMethod =
  <T extends schemas.$ZodType>(schema: T, io: "input" | "output", processors: Record<string, Processor> = {}) =>
  (params?: StandardJSONSchemaMethodParams): JSONSchema.BaseSchema => {
    const { libraryOptions, target } = params ?? {};
    const ctx = initializeContext({ ...(libraryOptions ?? {}), target, io, processors });
    process(schema, ctx);
    extractDefs(ctx, schema);
    return finalize(ctx, schema);
  };