"""Custom Pydantic types for the application."""
from bson import ObjectId
from pydantic_core import core_schema

class PyObjectId(ObjectId):
    """MongoDB ObjectId type for Pydantic validation."""
    @classmethod
    def __get_pydantic_core_schema__(cls, _source, _handler):
        """Get the Pydantic core schema for ObjectId validation."""
        return core_schema.no_info_after_validator_function(
            cls.validate,
            core_schema.union_schema(
                [
                    core_schema.str_schema(),
                    core_schema.is_instance_schema(ObjectId),
                ]
            ),
        )

    @classmethod
    def validate(cls, v):
        """Validate and convert string to ObjectId."""
        if isinstance(v, ObjectId):
            return v
        if not ObjectId.is_valid(v):
            raise ValueError("Invalid ObjectId")
        return ObjectId(v)
