import { z } from 'zod'
import { type Identity, Node, type NodeShape } from '../../types/index.ts'

function identity<T extends z.ZodString>(this: T) {
    return this.refine((v) => Node.isValid(v)).transform((v) => v as Identity)
}

function node<T extends z.ZodString>(this: T) {
    return this.refine((v) => Node.isValid(v)).transform((v) => Node.from(v as NodeShape))
}

z.ZodString.prototype.identity = identity
z.ZodString.prototype.node = node

// From https://github.com/colinhacks/zod/discussions/4445#discussioncomment-13236525
declare module 'zod' {
    interface ZodString {
        identity(): ReturnType<typeof identity<this>>
        node(): ReturnType<typeof node<this>>
    }
}