"""Job / Experiment 的 SQLAlchemy 实现（Phase 4）。"""

from __future__ import annotations

from sqlalchemy import select
from sqlalchemy.orm import Session

from app.domain.entities.research import ExperimentRecord, JobRecord
from app.infrastructure.persistence.sqlalchemy.models.jobs import ExperimentModel, JobModel


def _job_to_model(job: JobRecord) -> JobModel:
    return JobModel(**job.model_dump())


class SqlAlchemyJobRepository:
    def __init__(self, session: Session) -> None:
        self._session = session

    def create(self, job: JobRecord) -> JobRecord:
        self._session.add(_job_to_model(job))
        self._session.flush()
        return job

    def get(self, job_id: str) -> JobRecord | None:
        row = self._session.get(JobModel, job_id)
        return JobRecord.model_validate(row, from_attributes=True) if row else None

    def update(self, job: JobRecord) -> None:
        row = self._session.get(JobModel, job.id)
        if row is None:
            raise KeyError(f"job {job.id} 不存在")
        for k, v in job.model_dump().items():
            setattr(row, k, v)

    def list_recent(self, kind: str | None = None, limit: int = 20) -> list[JobRecord]:
        stmt = select(JobModel).order_by(JobModel.created_at.desc()).limit(limit)
        if kind:
            stmt = stmt.where(JobModel.kind == kind)
        return [
            JobRecord.model_validate(r, from_attributes=True)
            for r in self._session.scalars(stmt).all()
        ]

    def list_by_status(self, status: str, limit: int = 100) -> list[JobRecord]:
        rows = self._session.scalars(
            select(JobModel)
            .where(JobModel.status == status)
            .order_by(JobModel.created_at)
            .limit(limit)
        ).all()
        return [JobRecord.model_validate(r, from_attributes=True) for r in rows]


class SqlAlchemyExperimentRepository:
    def __init__(self, session: Session) -> None:
        self._session = session

    def save(self, experiment: ExperimentRecord) -> ExperimentRecord:
        self._session.add(ExperimentModel(**experiment.model_dump()))
        self._session.flush()
        return experiment

    def get(self, experiment_id: str) -> ExperimentRecord | None:
        row = self._session.get(ExperimentModel, experiment_id)
        return ExperimentRecord.model_validate(row, from_attributes=True) if row else None

    def list_recent(self, limit: int = 50) -> list[ExperimentRecord]:
        rows = self._session.scalars(
            select(ExperimentModel).order_by(ExperimentModel.created_at.desc()).limit(limit)
        ).all()
        return [ExperimentRecord.model_validate(r, from_attributes=True) for r in rows]