"""Agent run repository"""
from typing import Optional, List
from sqlalchemy.orm import Session
from app.core.db.models.agent_run import AgentRun, RunStatus
from datetime import datetime


class AgentRunRepository:
    """Repository for agent run operations"""
    
    def __init__(self, db: Session):
        self.db = db
    
    def create(self, agent_name: str) -> AgentRun:
        """Create a new agent run"""
        run = AgentRun(agent_name=agent_name, status=RunStatus.PENDING.value)
        self.db.add(run)
        self.db.commit()
        self.db.refresh(run)
        return run
    
    def get_by_id(self, run_id: int) -> Optional[AgentRun]:
        """Get run by ID"""
        return self.db.query(AgentRun).filter(AgentRun.id == run_id).first()
    
    def get_recent(self, agent_name: Optional[str] = None, limit: int = 50) -> List[AgentRun]:
        """Get recent runs"""
        query = self.db.query(AgentRun)
        if agent_name:
            query = query.filter(AgentRun.agent_name == agent_name)
        return query.order_by(AgentRun.started_at.desc()).limit(limit).all()
    
    def update_status(
        self,
        run_id: int,
        status: RunStatus,
        error: Optional[str] = None,
        logs: Optional[str] = None
    ) -> AgentRun:
        """Update run status"""
        run = self.get_by_id(run_id)
        if not run:
            raise ValueError(f"Run {run_id} not found")
        
        run.status = status.value
        if error:
            run.error = error
        if logs:
            run.logs = logs
        
        if status == RunStatus.COMPLETED.value or status == RunStatus.FAILED.value:
            run.finished_at = datetime.utcnow()
        
        self.db.commit()
        self.db.refresh(run)
        return run