from sqlmodel import Session
from sqlalchemy.dialects.mysql import insert
from typing import Optional

from app.models.regression_dataset import RegressionSample, SampleStatus
from app.models.validation_job import ValidationJob


class RegressionDatasetRepo:
    @staticmethod
    def create_from_job(db: Session, job: ValidationJob) -> int:
        sample = {
            "job_id": job.id,
            "message": job.message,
            "request": job.request,
            "steps": job.steps,
            "query_runs": job.query_runs,
        }

        stmt = insert(RegressionSample).values(sample)
        stmt = stmt.on_duplicate_key_update(**sample)
        db.execute(stmt)
        db.commit()

        sample_obj = (
            db.query(RegressionSample).filter(RegressionSample.job_id == job.id).first()
        )

        return sample_obj

    @staticmethod
    def get_samples(db: Session, sample_ids: list[int]) -> RegressionSample:
        return (
            db.query(RegressionSample).filter(RegressionSample.id.in_(sample_ids)).all()
        )

    @staticmethod
    def get_samples_paginated(db: Session, limit: int, offset: int) -> RegressionSample:
        return db.query(RegressionSample).offset(offset).limit(limit).all()

    @staticmethod
    def archive_samples(db: Session, sample_id: int):
        db.query(RegressionSample).filter(RegressionSample.id == sample_id).update(
            {"status": SampleStatus.ARCHIVED}
        )
        db.commit()

    @staticmethod
    def label_sample(db: Session, sample_id: int, gold_answer: dict, metadata: dict):
        db.query(RegressionSample).filter(RegressionSample.id == sample_id).update(
            {"gold_answer": gold_answer, "meta": metadata}
        )
        db.commit()