import random
import string
from sqlmodel import select, Session, Field
from sqlmodel import SQLModel
from typing import Optional
from pydantic import validator
from database.db import get_session
def generate_unique_code(model: SQLModel, s: Session, field_name='id_code',  length: int = 6):
    if getattr(model, field_name, None) is None:
        raise KeyError("Invalid field name")
    while True:
        code = ''.join(random.choices(string.ascii_uppercase, k=length))
        is_unique = s.exec(select(model).where(
            getattr(model, field_name) == code)).first() == None
        if is_unique:
            break
    return code

class SQLModelWithIdCode(SQLModel):
    id: Optional[int] = Field(default=None, primary_key=True)
    id_code: Optional[str] = Field(default=None, unique=True, index=True)
    
    @validator('id_code', always=True)
    def generate_id_code(cls, value, values):
        if value is None:
            sessions = get_session()
            session = next(sessions)
            code = generate_unique_code(cls, s=session)
            return code
        return value