# Agent Scheduler - Cron-based tasks for agents
import os
import asyncio
import logging
from datetime import datetime, timedelta
from typing import Dict, List, Any, Optional
import httpx

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

GATEWAY_URL = os.getenv("GATEWAY_URL", "http://dagi-gateway-node1:9300")


class AgentTask:
    def __init__(self, name: str, agent_id: str, schedule: str, task_type: str, params: Dict = None):
        self.name = name
        self.agent_id = agent_id
        self.schedule = schedule
        self.task_type = task_type
        self.params = params or {}
        self.last_run = None
        self.next_run = None


SCHEDULED_TASKS = [
    AgentTask("daily_health_check", "*", "0 9 * * *", "health_check"),
    AgentTask("helion_energy_report", "helion", "0 8 * * 1", "generate_report", {"report_type": "energy_weekly"}),
    AgentTask("agromatrix_weather", "agromatrix", "0 6 * * *", "web_search", {"query": "weather forecast Ukraine"}),
    AgentTask("memory_cleanup", "*", "0 3 * * 0", "memory_cleanup", {"older_than_days": 90}),
]


async def execute_task(task: AgentTask) -> Dict:
    logger.info(f"Executing task: {task.name} for agent: {task.agent_id}")
    
    try:
        async with httpx.AsyncClient(timeout=60.0) as client:
            if task.task_type == "health_check":
                resp = await client.get(f"{GATEWAY_URL}/health")
                return {"status": "ok" if resp.status_code == 200 else "error"}
            
            elif task.task_type == "generate_report":
                resp = await client.post(
                    f"{GATEWAY_URL}/{task.agent_id}/internal/task",
                    json={"task": "generate_report", "params": task.params}
                )
                return resp.json() if resp.status_code == 200 else {"error": resp.status_code}
            
            elif task.task_type == "web_search":
                resp = await client.post(
                    f"{GATEWAY_URL}/{task.agent_id}/internal/task",
                    json={"task": "web_search", "params": task.params}
                )
                return resp.json() if resp.status_code == 200 else {"error": resp.status_code}
            
            elif task.task_type == "memory_cleanup":
                memory_url = os.getenv("MEMORY_SERVICE_URL", "http://memory-service:8000")
                resp = await client.post(f"{memory_url}/admin/cleanup", json=task.params)
                return resp.json() if resp.status_code == 200 else {"error": resp.status_code}
            
            else:
                return {"error": f"Unknown task type: {task.task_type}"}
    
    except Exception as e:
        logger.error(f"Task execution failed: {e}")
        return {"error": str(e)}


async def run_scheduler():
    logger.info("Agent Scheduler started")
    
    while True:
        now = datetime.now()
        
        for task in SCHEDULED_TASKS:
            if should_run(task, now):
                result = await execute_task(task)
                task.last_run = now
                logger.info(f"Task {task.name} completed: {result}")
        
        await asyncio.sleep(60)


def should_run(task: AgentTask, now: datetime) -> bool:
    if not task.last_run:
        return False
    if task.last_run and (now - task.last_run) < timedelta(hours=1):
        return False
    return False


if __name__ == "__main__":
    asyncio.run(run_scheduler())