from sqlalchemy.orm import Session

from app.crud.base import CRUDBase
from app.models.data_model import DataModel, DataModelField
from app.schemas.data_model import DataModelCreate, DataModelUpdate
from common_logging import get_logger

logger = get_logger(__name__)


class CRUDDataModel(CRUDBase[DataModel, DataModelCreate, DataModelUpdate]):

    def remove(self, db: Session, *, id: int) -> DataModel | None:
        return self.delete(db, id=id)

    def get_by_key(self, db: Session, *, key: str) -> DataModel | None:
        return db.query(DataModel).filter(DataModel.key == key).first()

    def create_with_fields(
        self, db: Session, *, obj_in: DataModelCreate, created_by: int, tenant_id: int = None
    ) -> DataModel:
        model_data = obj_in.dict(exclude={"fields"})
        model_data["created_by"] = created_by
        model_data["tenant_id"] = tenant_id
        db_model = DataModel(**model_data)
        db.add(db_model)
        db.flush()
        for field_data in obj_in.fields:
            field = DataModelField(**field_data.dict(), model_id=db_model.id)
            db.add(field)
        db.commit()
        db.refresh(db_model)
        logger.bind(data_model_id=db_model.id).info("Data model created with fields")
        return db_model

    def update_with_fields(
        self, db: Session, *, db_obj: DataModel, obj_in: DataModelUpdate
    ) -> DataModel:
        update_data = obj_in.dict(exclude_unset=True, exclude={"fields"})
        for field, value in update_data.items():
            setattr(db_obj, field, value)
        if obj_in.fields is not None:
            db.query(DataModelField).filter(DataModelField.model_id == db_obj.id).delete()
            for field_data in obj_in.fields:
                field = DataModelField(**field_data.dict(exclude={"id"}), model_id=db_obj.id)
                db.add(field)
        db.commit()
        db.refresh(db_obj)
        logger.bind(data_model_id=db_obj.id).info("Data model updated with fields")
        return db_obj


data_model_crud = CRUDDataModel(DataModel)