services/agent_service.py

import asyncio
import logging
import os
from typing import Any, Dict, List, Optional
from pydantic import BaseModel

from core.service_factory import ServiceFactory
from integrations.mcp_service import mcp_service

# Try to import Lux SDK, fallback to local model if available
try:
    from oagi import LuxAgent
    HAS_SDK = True
except ImportError:
    HAS_SDK = False

from ai.lux_model import LuxModel

logger = logging.getLogger(__name__)

class AgentTask(BaseModel):
    id: str
    goal: str
    mode: str
    status: str
    logs: List[str] = []
    result: Optional[str] = None

class ComputerUseAgent:
    """
    Service for managing Lux Computer Use Agents.
    Supports 'actor', 'thinker', and 'tasker' modes.
    """
    
    
    def __init__(self, tenant_id: str = "default"):
        self.tenant_id = tenant_id
        self.default_mode = os.getenv("LUX_MODEL_MODE", "thinker")
        self._active_tasks: Dict[str, AgentTask] = {}
        self.mcp = mcp_service  # MCP access for web search and web access
        
        # We now rely on local LuxModel if SDK is missing
        logger.info(f"ComputerUseAgent initialized for tenant {tenant_id}")
            
    async def execute_task(self, goal: str, mode: Optional[str] = None) -> Dict[str, Any]:
        """
        Start a computer use task.
        """
        task_id = f"task_{len(self._active_tasks) + 1}_{int(asyncio.get_event_loop().time())}"
        mode = mode or self.default_mode
        
        task = AgentTask(
            id=task_id,
            goal=goal,
            mode=mode,
            status="running",
            logs=[f"Task started in {mode} mode: {goal}"]
        )
        self._active_tasks[task_id] = task
        
        # Run in background to not block API
        asyncio.create_task(self._run_agent_loop(task_id))
        
        return task.dict()
        
    async def _run_agent_loop(self, task_id: str):
        """
        Internal method to run the agent loop.
        handles both Real SDK execution and Mock fallback.
        """
        task = self._active_tasks.get(task_id)
        if not task:
            return

        try:
            # Use ServiceFactory to get the LuxModel instance
            # which is now tenant-aware and uses LLMService
            task.logs.append("Initializing Lux Agent (Unified Infrastructure)...")
            
            # Real Lux Execution via ServiceFactory resolved model
            agent = await ServiceFactory.get_lux_model(tenant_id=self.tenant_id)
            
            if True: # Key check handled inside LuxModel/LLMService now
                
                # --- Governance Setup ---
                from core.agent_governance_service import AgentGovernanceService
                from core.database import SessionLocal

                # Define callback for governance checks
                async def check_governance(action_type: str, details: Dict) -> bool:
                    try:
                        db = SessionLocal()
                        service = AgentGovernanceService(db)
                        
                        # Register Computer Use Agent if missing
                        agent = service.register_or_update_agent(
                            name="Computer Use Agent",
                            category="Desktop Automation",
                            module_path="backend.services.agent_service",
                            class_name="ComputerUseAgent",
                            description="AI Agent capable of controlling desktop mouse and keyboard."
                        )
                        
                        # Check permission
                        check = service.enforce_action(agent.id, action_type)
                        
                        if check["proceed"]:
                            return True
                        else:
                            # Log detailed reason for blockage
                            reason = check.get("reason", "Action blocked by governance policies.")
                            task.logs.append(f"⛔ Governance Blocked Action '{action_type}': {reason}")
                            return False
                    except Exception as e:
                        logger.error(f"Governance check failed: {e}")
                        # Fail safe: Block if check fails
                        return False
                    finally:
                        db.close()

                try:
                    # Execute
                    result_data = await agent.execute_command(task.goal)
                    
                    if result_data.get("success"):
                        task.status = "completed"
                        task.result = f"Task completed: {json.dumps(result_data.get('actions', []), indent=2)}"
                        task.logs.append(f"Success. Actions taken: {len(result_data.get('actions', []))}")
                    else:
                        task.status = "failed"
                        task.result = f"Task failed: {result_data.get('error')}"
                        task.logs.append(f"Failure: {result_data.get('error')}")
                        
                except Exception as model_err:
                     task.status = "failed"
                     task.result = str(model_err)
                     task.logs.append(f"Model Execution Error: {model_err}")
                
        except Exception as e:
            logger.error(f"Agent task failed: {e}")
            task.status = "failed"
            task.logs.append(f"Error: {str(e)}")
            task.result = str(e)

    def get_task_status(self, task_id: str) -> Optional[Dict[str, Any]]:
        task = self._active_tasks.get(task_id)
        return task.dict() if task else None

    def stop_task(self, task_id: str) -> bool:
        task = self._active_tasks.get(task_id)
        if task and task.status == "running":
            task.status = "stopped"
            task.logs.append("Task stopped by user.")
            return True
        return False

# Singleton instance
agent_service = ComputerUseAgent()