from enum import Enum
from typing import Annotated, Dict, Any, Optional, Type

from pydantic import BeforeValidator, BaseModel, ConfigDict, Field
from bson import ObjectId

PyObjectId = Annotated[str, BeforeValidator(str)]
   
class MongoBaseModel(BaseModel):
    id: Optional[PyObjectId] = Field(alias="_id", default_factory=ObjectId)

    model_config = ConfigDict(
        populate_by_name=True,
        arbitrary_types_allowed=True,
        json_encoders={PyObjectId: str},
    )

    def to_mongo(self):
        result = self.model_dump(by_alias=True, mode='json', exclude={'id',})
        result['_id'] = ObjectId(self.id)
        return result

    @classmethod
    def from_mongo(cls, data: Dict[str, Any]):
        def restore_enums(inst: Any, model_cls: Type[BaseModel]) -> None:
            for name, field in model_cls.__fields__.items():
                value = getattr(inst, name)
                if field and isinstance(field.annotation, type) and issubclass(field.annotation, Enum):
                    setattr(inst, name, field.annotation(value))
                elif isinstance(value, BaseModel):
                    restore_enums(value, value.__class__)
                elif isinstance(value, list):
                    for i, item in enumerate(value):
                        if isinstance(item, BaseModel):
                            restore_enums(item, item.__class__)
                        elif isinstance(field.annotation, type) and issubclass(field.annotation, Enum):
                            value[i] = field.annotation(item)
                elif isinstance(value, dict):
                    for k, v in value.items():
                        if isinstance(v, BaseModel):
                            restore_enums(v, v.__class__)
                        elif isinstance(field.annotation, type) and issubclass(field.annotation, Enum):
                            value[k] = field.annotation(v)

        if data is None:
            return None
        instance = cls(**data)
        restore_enums(instance, instance.__class__)
        return instance