from sqlalchemy.orm import Session

from packages.adapters.sqlalchemy_impl.models import GenerationTaskModel
from packages.domain import GenerationTask, GenerationTaskStatus


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

    def create(self, task: GenerationTask) -> GenerationTask:
        model = GenerationTaskModel(
            id=task.id,
            workspace_id=task.workspace_id,
            project_id=task.project_id,
            strategy_id=task.strategy_id,
            asset_library_id=task.asset_library_id,
            voice_library_id=task.voice_library_id,
            status=task.status.value,
            progress=task.progress,
            result_count=task.result_count,
            error_message=task.error_message,
            started_at=task.started_at,
            completed_at=task.completed_at,
            created_by_user_id=task.created_by_user_id,
            created_at=task.created_at,
        )
        self.session.add(model)
        self.session.commit()
        return task

    def get(self, task_id: str) -> GenerationTask | None:
        model = self.session.query(GenerationTaskModel).filter(GenerationTaskModel.id == task_id).first()
        if model is None:
            return None
        return GenerationTask(
            id=model.id,
            workspace_id=model.workspace_id,
            project_id=model.project_id,
            strategy_id=model.strategy_id,
            asset_library_id=model.asset_library_id,
            voice_library_id=model.voice_library_id,
            status=GenerationTaskStatus(model.status),
            progress=model.progress,
            result_count=int(model.result_count or 0),
            error_message=model.error_message,
            started_at=model.started_at,
            completed_at=model.completed_at,
            created_by_user_id=model.created_by_user_id,
            created_at=model.created_at,
        )

    def list_by_project(self, project_id: str) -> list[GenerationTask]:
        models = self.session.query(GenerationTaskModel).filter(GenerationTaskModel.project_id == project_id).all()
        return [self.get(model.id) for model in models if self.get(model.id) is not None]

    def update(self, task: GenerationTask) -> GenerationTask:
        model = self.session.query(GenerationTaskModel).filter(GenerationTaskModel.id == task.id).first()
        if model is None:
            raise ValueError(f"GenerationTask {task.id} not found")
        model.status = task.status.value
        model.progress = task.progress
        model.result_count = task.result_count
        model.error_message = task.error_message
        model.started_at = task.started_at
        model.completed_at = task.completed_at
        self.session.commit()
        return task