"""Base agent class"""
from abc import ABC, abstractmethod
from sqlalchemy.orm import Session
from typing import Optional
import logging
from app.core.db.repositories.agent_run_repository import AgentRunRepository
from app.core.db.models.agent_run import RunStatus
from datetime import datetime

logger = logging.getLogger(__name__)


class BaseAgent(ABC):
    """Base class for all agents"""
    
    def __init__(self, db: Session, agent_name: str):
        self.db = db
        self.agent_name = agent_name
        self.run_repo = AgentRunRepository(db)
        self.current_run: Optional[int] = None
    
    def start_run(self) -> int:
        """Start a new agent run"""
        run = self.run_repo.create(self.agent_name)
        self.current_run = run.id
        self.run_repo.update_status(run.id, RunStatus.RUNNING)
        logger.info(f"Started {self.agent_name} run {run.id}")
        return run.id
    
    def finish_run(self, success: bool = True, error: Optional[str] = None, logs: Optional[str] = None):
        """Finish the current run"""
        if not self.current_run:
            return
        
        status = RunStatus.COMPLETED if success else RunStatus.FAILED
        self.run_repo.update_status(self.current_run, status, error=error, logs=logs)
        logger.info(f"Finished {self.agent_name} run {self.current_run} with status {status.value}")
        self.current_run = None
    
    @abstractmethod
    def run(self) -> bool:
        """Execute the agent's main logic"""
        pass