from sqlalchemy.orm import Session

from packages.adapters.sqlalchemy_impl.models import ClassificationJobModel
from packages.domain import ClassificationJob


class SQLAlchemyClassificationJobRepository:
    def __init__(self, session: Session):
        self.session = session

    def create(self, job: ClassificationJob) -> ClassificationJob:
        model = ClassificationJobModel(
            id=job.id,
            project_id=job.project_id,
            asset_id=job.asset_id,
            status=job.status,
            classification=job.classification,
            confidence=job.confidence,
            error_message=job.error_message,
            created_at=job.created_at,
            updated_at=job.updated_at,
        )
        self.session.add(model)
        self.session.commit()
        return job

    def get(self, job_id: str) -> ClassificationJob | None:
        model = self.session.query(ClassificationJobModel).filter(ClassificationJobModel.id == job_id).first()
        if model is None:
            return None
        return ClassificationJob(
            id=model.id,
            project_id=model.project_id,
            asset_id=model.asset_id,
            status=model.status,
            classification=model.classification,
            confidence=model.confidence,
            error_message=model.error_message,
            created_at=model.created_at,
            updated_at=model.updated_at,
        )

    def update(self, job: ClassificationJob) -> ClassificationJob:
        model = self.session.query(ClassificationJobModel).filter(ClassificationJobModel.id == job.id).first()
        if model is None:
            raise ValueError(f"ClassificationJob {job.id} not found")
        model.status = job.status
        model.classification = job.classification
        model.confidence = job.confidence
        model.error_message = job.error_message
        model.updated_at = job.updated_at
        self.session.commit()
        return job