from __future__ import annotations

from sqlalchemy import select
from sqlalchemy.orm import Session

from app.models.edge_node import EdgeNode
from app.models.edge_task import EdgeTask


class EdgeNodeRepository:
    def __init__(self, db: Session) -> None:
        self.db = db

    def add_or_update(self, node: EdgeNode) -> EdgeNode:
        self.db.add(node)
        self.db.commit()
        self.db.refresh(node)
        return node

    def get_by_edge_id(self, edge_id: str) -> EdgeNode | None:
        statement = select(EdgeNode).where(EdgeNode.edge_id == edge_id)
        return self.db.execute(statement).scalar_one_or_none()


class EdgeTaskRepository:
    def __init__(self, db: Session) -> None:
        self.db = db

    def add(self, edge_task: EdgeTask) -> EdgeTask:
        self.db.add(edge_task)
        self.db.commit()
        self.db.refresh(edge_task)
        return edge_task

    def update(self, edge_task: EdgeTask) -> EdgeTask:
        self.db.add(edge_task)
        self.db.commit()
        self.db.refresh(edge_task)
        return edge_task

    def get_by_step_id(self, step_id: str) -> EdgeTask | None:
        statement = select(EdgeTask).where(EdgeTask.step_id == step_id)
        return self.db.execute(statement).scalar_one_or_none()

    def list_by_task_id(self, task_id: str) -> list[EdgeTask]:
        statement = select(EdgeTask).where(EdgeTask.task_id == task_id).order_by(EdgeTask.created_at.desc())
        return list(self.db.execute(statement).scalars())

    def list_pending_by_edge_id(self, edge_id: str) -> list[EdgeTask]:
        statement = (
            select(EdgeTask)
            .where(EdgeTask.edge_id == edge_id)
            .where(EdgeTask.step_status == "PENDING")
            .order_by(EdgeTask.created_at.asc())
        )
        return list(self.db.execute(statement).scalars())

    def list_active_by_task_id(self, task_id: str) -> list[EdgeTask]:
        statement = (
            select(EdgeTask)
            .where(EdgeTask.task_id == task_id)
            .where(EdgeTask.step_status.in_(["PENDING", "RUNNING"]))
            .order_by(EdgeTask.created_at.asc())
        )
        return list(self.db.execute(statement).scalars())