import Maybe, { just, nothing } from "true-myth/maybe";
import * as v from "@atcute/lexicons/validations";

type IsUnion<T, U = T> = T extends unknown ? ([U] extends [T] ? false : true) : never;

type MaybeifyUnion<T> = T extends unknown ? (T extends object ? Maybeify<T> : T) : never;

export type Maybeify<T> = {
  [K in keyof T]-?: undefined extends T[K]
    ? T[K] extends Array<infer U> | undefined
      ? Maybe<Maybeify<U>[]>
      : IsUnion<NonNullable<T[K]>> extends true
        ? Maybe<MaybeifyUnion<NonNullable<T[K]>>>
        : T[K] extends object | undefined
          ? Maybe<Maybeify<NonNullable<T[K]>>>
          : Maybe<NonNullable<T[K]>>
    : T[K] extends Array<infer U>
      ? Maybeify<U>[]
      : IsUnion<T[K]> extends true
        ? MaybeifyUnion<T[K]>
        : T[K] extends object
          ? Maybeify<T[K]>
          : T[K];
};

type AnySchema = v.BaseSchema;

function resolveMemberSchema(
  member: v.ObjectSchema | v.RecordSchema<v.ObjectSchema, v.RecordKeySchema>,
): v.ObjectSchema {
  return (member as { type: string }).type === "record"
    ? (member as v.RecordSchema<v.ObjectSchema, v.RecordKeySchema>).object
    : (member as v.ObjectSchema);
}

function resolveLiteralFromShape(shape: Record<string, AnySchema>): string | undefined {
  const typeSchema = shape["$type"] as AnySchema | undefined;
  if (typeSchema === undefined) {
    return undefined;
  }
  const inner =
    (typeSchema as { type: string }).type === "optional"
      ? (typeSchema as v.OptionalSchema<AnySchema, unknown>).wrapped
      : typeSchema;
  if ((inner as { type: string }).type === "literal") {
    return (inner as v.LiteralSchema<string>).expected;
  }
  return undefined;
}

export function maybeify<TSchema extends AnySchema>(
  schema: TSchema,
  value: v.InferOutput<TSchema>,
): Maybeify<v.InferOutput<TSchema>> {
  const schemaType = (schema as { type: string }).type;

  if (schemaType === "optional") {
    const wrapped = (schema as unknown as v.OptionalSchema<AnySchema, unknown>).wrapped;
    if (value === undefined) {
      return nothing() as unknown as Maybeify<v.InferOutput<TSchema>>;
    }
    return just(maybeify(wrapped, value as never)) as unknown as Maybeify<v.InferOutput<TSchema>>;
  }

  if (schemaType === "record") {
    const objectSchema = (schema as unknown as v.RecordSchema<v.ObjectSchema, v.RecordKeySchema>)
      .object;
    return maybeify(objectSchema, value as never) as unknown as Maybeify<v.InferOutput<TSchema>>;
  }

  if (schemaType === "object") {
    const objectSchema = schema as unknown as v.ObjectSchema<Record<string, AnySchema>>;
    const result: Record<string, unknown> = {};
    for (const [key, propSchema] of Object.entries(objectSchema.shape)) {
      const propValue = (value as Record<string, unknown>)[key];
      result[key] = maybeify(propSchema as AnySchema, propValue as never);
    }
    return result as unknown as Maybeify<v.InferOutput<TSchema>>;
  }

  if (schemaType === "array") {
    const arraySchema = schema as unknown as v.ArraySchema<AnySchema>;
    return (value as unknown[]).map((item) =>
      maybeify(arraySchema.item, item as never),
    ) as unknown as Maybeify<v.InferOutput<TSchema>>;
  }

  if (schemaType === "variant") {
    const variantSchema = schema as unknown as v.VariantSchema<any, any>;
    const typeValue = (value as Record<string, unknown>)["$type"] as string | undefined;
    if (typeValue !== undefined) {
      for (const member of variantSchema.members) {
        const memberSchema = resolveMemberSchema(member);
        const expectedType = resolveLiteralFromShape(
          memberSchema.shape as Record<string, AnySchema>,
        );
        if (expectedType === typeValue) {
          return maybeify(memberSchema, value as never) as unknown as Maybeify<
            v.InferOutput<TSchema>
          >;
        }
      }
    }
    return value as unknown as Maybeify<v.InferOutput<TSchema>>;
  }

  return value as unknown as Maybeify<v.InferOutput<TSchema>>;
}