from pydantic import BaseModel, Field, EmailStr
from typing import Optional, List, Dict
from datetime import datetime
from bson import ObjectId
from enum import Enum

class UserRole(str, Enum):
    ADMIN = "admin"
    USER = "user"

class AuthMethod(str, Enum):
    LOCAL = "local"
    SSO = "sso"

class PyObjectId(ObjectId):
    @classmethod
    def __get_pydantic_core_schema__(cls, _source_type, _handler):
        from pydantic_core import core_schema
        
        def validate_from_str(input_value: str) -> ObjectId:
            return ObjectId(input_value)

        def validate_from_objectid(input_value: ObjectId) -> ObjectId:
            return input_value

        return core_schema.union_schema([
            core_schema.is_instance_schema(ObjectId),
            core_schema.no_info_plain_validator_function(
                validate_from_str,
                serialization=core_schema.to_string_ser_schema(),
            ),
        ])

    @classmethod
    def __get_pydantic_json_schema__(cls, core_schema, handler):
        json_schema = handler(core_schema)
        json_schema.update(type="string")
        return json_schema

class UserBase(BaseModel):
    email: EmailStr
    role: UserRole = UserRole.USER
    is_active: bool = True
    auth_method: AuthMethod = AuthMethod.LOCAL
    sso_provider: Optional[str] = None
    sso_user_id: Optional[str] = None
    sso_email: Optional[str] = None
    sso_name: Optional[str] = None
    sso_attributes: Optional[Dict] = None
    last_sso_login: Optional[datetime] = None
    created_at: Optional[datetime] = None
    updated_at: Optional[datetime] = None

class UserCreate(UserBase):
    password: str

class UserUpdate(BaseModel):
    email: Optional[EmailStr] = None
    role: Optional[UserRole] = None
    is_active: Optional[bool] = None
    password: Optional[str] = None
    auth_method: Optional[AuthMethod] = None
    sso_provider: Optional[str] = None
    sso_user_id: Optional[str] = None
    sso_email: Optional[str] = None
    sso_name: Optional[str] = None
    sso_attributes: Optional[Dict] = None
    last_sso_login: Optional[datetime] = None

class UserInDB(UserBase):
    id: PyObjectId = Field(default_factory=PyObjectId, alias="_id")
    hashed_password: Optional[str] = None  # Optional for SSO users
    index_access: List[str] = Field(default_factory=list)  # List of index IDs user has access to
    
    class Config:
        populate_by_name = True
        arbitrary_types_allowed = True
        json_encoders = {ObjectId: str}

class User(UserBase):
    id: PyObjectId = Field(default_factory=PyObjectId, alias="_id")
    index_access: List[str] = Field(default_factory=list)
    
    class Config:
        populate_by_name = True
        arbitrary_types_allowed = True
        json_encoders = {ObjectId: str}

class UserResponse(BaseModel):
    id: PyObjectId = Field(alias="_id")
    email: EmailStr
    role: UserRole
    is_active: bool
    auth_method: AuthMethod
    sso_provider: Optional[str] = None
    sso_name: Optional[str] = None
    last_sso_login: Optional[datetime] = None
    index_access: List[str]
    created_at: Optional[datetime] = None
    updated_at: Optional[datetime] = None
    
    class Config:
        populate_by_name = True
        arbitrary_types_allowed = True
        json_encoders = {ObjectId: str}