from enum import Enum

import pydantic.json

class EnumByName(Enum):
   """A custom Enum type for pydantic to validate by name.
   """
   # The reason for this class is there is a disconnect between how
   # SQLAlchemy stores Enum references (by name, not value; which is
   # also how we want our REST API to exchange enum references, by
   # name) and Pydantic which validates Enums by value.  We create a
   # Pydantic custom type which will validate an Enum reference by
   # name.

   # Ugliness: we need to monkeypatch pydantic's jsonification of Enums
   pydantic.json.ENCODERS_BY_TYPE[Enum] = lambda e: e.name

   @classmethod
   def __get_validators__(cls):
      # yield our validator
      yield cls._validate

   @classmethod
   def __modify_schema__(cls, schema):
      """Override pydantic using Enum.name for schema enum values"""
      schema['enum'] = list(cls.__members__.keys())

   @classmethod
   def _validate(cls, v):
      """Validate enum reference, `v`.  

      We check:
        1. If it is a member of this Enum
        2. If we can find it by name.
      """
      # is the value an enum member?
      try:
         if v in cls:
            return v
      except TypeError:
         pass

      # not a member...look up by name
      try:
         return cls[v]
      except KeyError:
         name = cls.__name__
         expected = list(cls.__members__.keys())
         raise ValueError(f'{v} not found for enum {name}. Expected one of: {expected}')

from enum import IntEnum
from typing import Annotated

import pytest
from pydantic import BaseModel, BeforeValidator, ValidationError


class LogLevel(IntEnum):
    DEBUG = 1
    INFO = 2


def log_level_by_name(v: str) -> LogLevel:
    try:
        return LogLevel[v]
    except KeyError:
        raise ValueError("invalid value")


LogLevelByName = Annotated[LogLevel, BeforeValidator(log_level_by_name)]


class AppConfig(BaseModel):
    log_level: LogLevelByName


def test_enum():
    c = AppConfig.model_validate({"log_level": "DEBUG"})
    assert c.log_level is LogLevel.DEBUG

    with pytest.raises(ValidationError):
        c = AppConfig.model_validate({"log_level": 1})

class MyEnum(enum.Enum):
    foo = enum.auto()
    bar = enum.auto()
    baz = enum.auto()

import enum
from typing import Any
from pydantic import GetCoreSchemaHandler
from pydantic_core import core_schema


def pydantic_enum[E: enum.Enum](enum_cls: type[E]) -> type[E]:
    def __get_pydantic_core_schema__(cls: type[E], source_type: Any, handler: GetCoreSchemaHandler):
        assert source_type is cls
        
        def get_enum(value: Any, validate_next: core_schema.ValidatorFunctionWrapHandler):
            if isinstance(value, cls):
                return value
            else:
                name: str = validate_next(value)
                return enum_cls[name]

        def serialize(enum: E):
            return enum.name
        
        expected = [member.name for member in cls]
        name_schema = core_schema.literal_schema(expected)
        
        return core_schema.no_info_wrap_validator_function(
            get_enum, name_schema, 
            ref=cls.__name__,
            serialization=core_schema.plain_serializer_function_ser_schema(serialize)
        )
    
    setattr(enum_cls, '__get_pydantic_core_schema__', classmethod(__get_pydantic_core_schema__))
    return enum_cls

from pydantic import BaseModel, ValidationError

@pydantic_enum
class MyEnum(enum.Enum):
    biba = enum.auto()
    boba = enum.auto()

class Owner(BaseModel):
    value: MyEnum

assert Owner.model_validate({
    'value': 'biba'
}).value == MyEnum.biba
assert Owner.model_validate({
    'value': MyEnum.boba
}).value == MyEnum.boba

try:
    Owner.model_validate({
        'value': 1
    })
    assert False
except ValidationError:
    pass

assert Owner(value=MyEnum.biba).model_dump(mode='json') == {'value': 'biba'}

class EnumByName:
    def __init__(self, *, ignore_case: bool = True):
        self.ignore_case = ignore_case

    def __get_pydantic_core_schema__(self, enum_cls: type[Enum], _handler: GetCoreSchemaHandler):
        name_enum = Enum('name_enum', {member.name: member.name for member in enum_cls})
        name_enum = cast(type[Enum], name_enum)

        def enum_or_name(value: Enum | str) -> Enum:
            if isinstance(value, str):
                if not self.ignore_case:
                    try:
                        return enum_cls[value]
                    except KeyError:
                        raise ValueError(f'Enum name not found: {value}')
                try:
                    return next(
                        member for member in enum_cls if member.name.lower() == value.lower()
                    )
                except StopIteration:
                    raise ValueError(f'Enum name not found: {value}')
            elif isinstance(value, enum_cls):
                return value
            raise ValueError(f'Expected enum member or name, got {type(value).__name__}: {value}')

        return core_schema.no_info_plain_validator_function(
            enum_or_name,
            json_schema_input_schema=core_schema.enum_schema(
                enum_cls, list(name_enum.__members__.values())
            ),
            ref=enum_cls.__name__,
            serialization=core_schema.plain_serializer_function_ser_schema(lambda e: e.name),
        )


class MyEnum(Enum):
    A = 1
    B = 2
    C = 3


class MyModel(BaseModel):
    my_enum: Annotated[MyEnum, EnumByName()]