"""
Multi-Agent Scheduler for Autonomous Agent Coordination

Provides coordinated scheduling and execution of multiple autonomous agents:
- MasterOrchestratorAgent: Pipeline health and ETL orchestration
- HealerAgent: Self-healing and fault tolerance
- GovernorAgent: Data governance and quality

Features:
- Configurable schedules per agent
- Inter-agent communication and triggers
- Priority-based task queue
- Concurrent execution with conflict detection
- Event-driven agent activation

Token Zero Design:
- Scheduler logic in pure Python (0 tokens)
- Agent intent understanding (~50 tokens per agent)
- No LLM calls in coordination layer
"""

from __future__ import annotations

import asyncio
import logging
import os
import uuid
from dataclasses import dataclass, field
from datetime import datetime, timezone, timedelta
from enum import Enum
from typing import Any, Callable, TYPE_CHECKING

if TYPE_CHECKING:
    from .master_orchestrator_agent import MasterOrchestratorAgent
    from .healer_agent import HealerAgent
    from .governor_agent import GovernorAgent

from .types import (
    AgentExecutionResult,
    AgentPriority,
    TriggerType,
)

logger = logging.getLogger(__name__)


# ============================================================================
# Types
# ============================================================================


class AgentType(str, Enum):
    """Types of autonomous agents"""
    ORCHESTRATOR = "orchestrator"
    HEALER = "healer"
    GOVERNOR = "governor"


class TaskStatus(str, Enum):
    """Status of a scheduled task"""
    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELLED = "cancelled"


@dataclass
class ScheduledAgentTask:
    """A scheduled task for an agent"""
    task_id: str
    agent_type: AgentType
    task_name: str
    priority: AgentPriority = AgentPriority.MEDIUM
    trigger_type: TriggerType = TriggerType.SCHEDULED
    parameters: dict[str, Any] = field(default_factory=dict)
    status: TaskStatus = TaskStatus.PENDING
    created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
    started_at: datetime | None = None
    completed_at: datetime | None = None
    result: AgentExecutionResult | None = None
    error: str | None = None
    triggered_by: str | None = None  # ID of task that triggered this


@dataclass
class AgentScheduleConfig:
    """Configuration for agent scheduling"""
    agent_type: AgentType
    enabled: bool = True
    interval_hours: float = 1.0
    default_task: str = "health_check"
    priority: AgentPriority = AgentPriority.MEDIUM
    max_concurrent: int = 1
    triggers_on_failure: list[AgentType] = field(default_factory=list)
    triggers_on_issues: list[AgentType] = field(default_factory=list)


@dataclass
class SchedulerStats:
    """Statistics for the scheduler"""
    total_tasks_scheduled: int = 0
    total_tasks_completed: int = 0
    total_tasks_failed: int = 0
    tasks_by_agent: dict[str, int] = field(default_factory=dict)
    average_duration_ms: float = 0.0
    last_run_time: datetime | None = None
    uptime_seconds: float = 0.0


# ============================================================================
# Multi-Agent Scheduler
# ============================================================================


class MultiAgentScheduler:
    """
    Coordinated scheduler for multiple autonomous agents.

    Responsibilities:
    - Schedule agents based on configurable intervals
    - Handle inter-agent triggers (e.g., Governor triggers Healer on quality issues)
    - Manage task queue with priorities
    - Track execution statistics
    - Prevent conflicts (e.g., don't run ETL while governance audit)

    Example usage:
        scheduler = MultiAgentScheduler()
        scheduler.register_agent(AgentType.ORCHESTRATOR, orchestrator_agent)
        scheduler.register_agent(AgentType.HEALER, healer_agent)
        scheduler.register_agent(AgentType.GOVERNOR, governor_agent)
        await scheduler.start()
    """

    # Default schedules for each agent type
    DEFAULT_SCHEDULES = {
        AgentType.ORCHESTRATOR: AgentScheduleConfig(
            agent_type=AgentType.ORCHESTRATOR,
            interval_hours=1.0,
            default_task="health_check",
            priority=AgentPriority.HIGH,
            triggers_on_failure=[AgentType.HEALER],
        ),
        AgentType.HEALER: AgentScheduleConfig(
            agent_type=AgentType.HEALER,
            interval_hours=4.0,
            default_task="assess_health",
            priority=AgentPriority.HIGH,
            triggers_on_issues=[AgentType.ORCHESTRATOR],
        ),
        AgentType.GOVERNOR: AgentScheduleConfig(
            agent_type=AgentType.GOVERNOR,
            interval_hours=6.0,
            default_task="check_governance_status",
            priority=AgentPriority.MEDIUM,
            triggers_on_issues=[AgentType.HEALER],
        ),
    }

    def __init__(self):
        """Initialize the multi-agent scheduler."""
        self._agents: dict[AgentType, Any] = {}
        self._schedules: dict[AgentType, AgentScheduleConfig] = {}
        self._task_queue: list[ScheduledAgentTask] = []
        self._running_tasks: dict[str, ScheduledAgentTask] = {}
        self._completed_tasks: list[ScheduledAgentTask] = []
        self._running = False
        self._started_at: datetime | None = None
        self._stats = SchedulerStats()
        self._task_counter = 0
        self._lock = asyncio.Lock()

        logger.info("MultiAgentScheduler initialized")

    # =========================================================================
    # Agent Registration
    # =========================================================================

    def register_agent(
        self,
        agent_type: AgentType,
        agent: Any,
        schedule: AgentScheduleConfig | None = None,
    ) -> None:
        """
        Register an agent with the scheduler.

        Args:
            agent_type: Type of agent
            agent: Agent instance
            schedule: Optional custom schedule (uses default if not provided)
        """
        self._agents[agent_type] = agent
        self._schedules[agent_type] = schedule or self.DEFAULT_SCHEDULES.get(
            agent_type,
            AgentScheduleConfig(agent_type=agent_type),
        )

        logger.info(
            f"Registered {agent_type.value} agent with "
            f"{self._schedules[agent_type].interval_hours}h interval"
        )

    def update_schedule(
        self,
        agent_type: AgentType,
        **kwargs,
    ) -> None:
        """
        Update schedule configuration for an agent.

        Args:
            agent_type: Type of agent
            **kwargs: Schedule parameters to update
        """
        if agent_type not in self._schedules:
            raise ValueError(f"Agent {agent_type.value} not registered")

        schedule = self._schedules[agent_type]
        for key, value in kwargs.items():
            if hasattr(schedule, key):
                setattr(schedule, key, value)

        logger.info(f"Updated schedule for {agent_type.value}: {kwargs}")

    # =========================================================================
    # Task Management
    # =========================================================================

    def _generate_task_id(self) -> str:
        """Generate a unique task ID."""
        self._task_counter += 1
        return f"task_{self._task_counter}_{datetime.now(timezone.utc).strftime('%Y%m%d%H%M%S')}"

    async def schedule_task(
        self,
        agent_type: AgentType,
        task_name: str,
        priority: AgentPriority = AgentPriority.MEDIUM,
        trigger_type: TriggerType = TriggerType.MANUAL,
        parameters: dict[str, Any] | None = None,
        triggered_by: str | None = None,
    ) -> ScheduledAgentTask:
        """
        Schedule a task for an agent.

        Args:
            agent_type: Type of agent to run
            task_name: Name of the task (method name)
            priority: Task priority
            trigger_type: How the task was triggered
            parameters: Optional parameters for the task
            triggered_by: ID of task that triggered this one

        Returns:
            The scheduled task
        """
        async with self._lock:
            task = ScheduledAgentTask(
                task_id=self._generate_task_id(),
                agent_type=agent_type,
                task_name=task_name,
                priority=priority,
                trigger_type=trigger_type,
                parameters=parameters or {},
                triggered_by=triggered_by,
            )

            # Insert by priority (higher priority first)
            priority_order = {
                AgentPriority.CRITICAL: 0,
                AgentPriority.HIGH: 1,
                AgentPriority.MEDIUM: 2,
                AgentPriority.LOW: 3,
            }

            insert_idx = len(self._task_queue)
            for i, existing in enumerate(self._task_queue):
                if priority_order[task.priority] < priority_order[existing.priority]:
                    insert_idx = i
                    break

            self._task_queue.insert(insert_idx, task)
            self._stats.total_tasks_scheduled += 1

            logger.info(
                f"Scheduled task {task.task_id}: "
                f"{agent_type.value}.{task_name} (priority={priority.value})"
            )

            return task

    async def _get_next_task(self) -> ScheduledAgentTask | None:
        """Get the next task to execute."""
        async with self._lock:
            if not self._task_queue:
                return None

            # Find first task that can run (no conflicts)
            for i, task in enumerate(self._task_queue):
                if self._can_run_task(task):
                    return self._task_queue.pop(i)

            return None

    def _can_run_task(self, task: ScheduledAgentTask) -> bool:
        """Check if a task can run (no conflicts)."""
        schedule = self._schedules.get(task.agent_type)
        if not schedule:
            return False

        # Check max concurrent for this agent type
        running_same_type = sum(
            1 for t in self._running_tasks.values()
            if t.agent_type == task.agent_type
        )

        if running_same_type >= schedule.max_concurrent:
            return False

        # Add conflict rules here if needed
        # e.g., don't run Governor while Orchestrator is running ETL

        return True

    # =========================================================================
    # Database Logging
    # =========================================================================

    async def _log_execution_to_db(self, task: ScheduledAgentTask) -> None:
        """Log agent execution to the monitoring database."""
        try:
            import asyncpg

            db_url = os.environ.get(
                "TIMESCALEDB_URL",
                "postgresql://sweetnight:sweetnight123@localhost:5433/sweetnight_geo"
            )
            # asyncpg requires plain postgresql:// scheme
            db_url = db_url.replace("postgresql+asyncpg://", "postgresql://")

            # Use execution ID from task (set during execution)
            execution_id = getattr(self._agents.get(task.agent_type), '_execution_id', None)
            if not execution_id:
                execution_id = f"{task.agent_type.value}_{task.task_id}_{uuid.uuid4().hex[:8]}"

            # Calculate duration
            duration_ms = None
            if task.started_at and task.completed_at:
                duration_ms = (task.completed_at - task.started_at).total_seconds() * 1000

            # Extract metrics from result
            tokens_used = 0
            turns = 0
            tools_called = 0
            errors_count = 0
            result_summary = None

            if task.result:
                tokens_used = getattr(task.result, 'tokens_used', 0) or 0
                # Get turns and tools_called from result_data if available
                result_data = getattr(task.result, 'result_data', {}) or {}
                turns = result_data.get('turns', 0)
                tools_called = result_data.get('tools_called', 0)
                errors = getattr(task.result, 'errors', []) or []
                errors_count = len(errors)
                # Try to get a summary from result_data or action
                result_summary = result_data.get('summary', getattr(task.result, 'action', None))
                if result_summary and len(result_summary) > 500:
                    result_summary = result_summary[:500]

            conn = await asyncpg.connect(db_url)
            try:
                await conn.execute(
                    """
                    INSERT INTO dwd.agent_executions (
                        execution_id, agent_type, task_name, status, priority,
                        trigger_type, success, duration_ms, tokens_used, turns,
                        tools_called, errors_count, error_message, result_summary,
                        started_at, completed_at
                    ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16)
                    """,
                    execution_id,
                    task.agent_type.value,
                    task.task_name,
                    task.status.value,
                    task.priority.value,
                    task.trigger_type.value if task.trigger_type else None,
                    task.result.success if task.result else False,
                    duration_ms,
                    tokens_used,
                    turns,
                    tools_called,
                    1 if task.error else 0,
                    task.error,
                    result_summary,
                    task.started_at,
                    task.completed_at,
                )
                logger.info(f"Logged execution {execution_id} to database")
            finally:
                await conn.close()

        except Exception as e:
            logger.warning(f"Failed to log execution to database: {e}")

    # =========================================================================
    # Task Execution
    # =========================================================================

    async def _execute_task(self, task: ScheduledAgentTask) -> None:
        """Execute a single task."""
        agent = self._agents.get(task.agent_type)
        if not agent:
            task.status = TaskStatus.FAILED
            task.error = f"Agent {task.agent_type.value} not registered"
            logger.error(task.error)
            return

        task.status = TaskStatus.RUNNING
        task.started_at = datetime.now(timezone.utc)
        self._running_tasks[task.task_id] = task

        logger.info(f"Executing task {task.task_id}: {task.agent_type.value}.{task.task_name}")

        # Generate execution ID and set on agent for tool logging
        execution_id = f"{task.agent_type.value}_{task.task_id}_{uuid.uuid4().hex[:8]}"
        if hasattr(agent, '_execution_id'):
            agent._execution_id = execution_id

        try:
            # Get the method to call
            method_name = task.task_name
            if hasattr(agent, method_name):
                method = getattr(agent, method_name)
                if asyncio.iscoroutinefunction(method):
                    result = await method(**task.parameters)
                else:
                    result = method(**task.parameters)
            elif hasattr(agent, "run"):
                # Fallback to generic run method
                result = await agent.run(task=task.task_name, **task.parameters)
            else:
                raise AttributeError(f"Agent has no method: {method_name}")

            task.result = result
            task.status = TaskStatus.COMPLETED
            self._stats.total_tasks_completed += 1

            # Update agent stats
            agent_key = task.agent_type.value
            self._stats.tasks_by_agent[agent_key] = (
                self._stats.tasks_by_agent.get(agent_key, 0) + 1
            )

            # Handle triggers based on result
            await self._handle_task_result(task)

        except Exception as e:
            task.status = TaskStatus.FAILED
            task.error = str(e)
            self._stats.total_tasks_failed += 1
            logger.error(f"Task {task.task_id} failed: {e}")

            # Trigger failure handlers
            await self._handle_task_failure(task)

        finally:
            task.completed_at = datetime.now(timezone.utc)

            # Calculate duration
            if task.started_at:
                duration_ms = (task.completed_at - task.started_at).total_seconds() * 1000
                # Update rolling average
                total = self._stats.total_tasks_completed + self._stats.total_tasks_failed
                if total > 0:
                    self._stats.average_duration_ms = (
                        (self._stats.average_duration_ms * (total - 1) + duration_ms) / total
                    )

            # Move from running to completed
            if task.task_id in self._running_tasks:
                del self._running_tasks[task.task_id]
            self._completed_tasks.append(task)

            # Keep only last 100 completed tasks
            if len(self._completed_tasks) > 100:
                self._completed_tasks = self._completed_tasks[-100:]

            # Log execution to monitoring database
            await self._log_execution_to_db(task)

    async def _handle_task_result(self, task: ScheduledAgentTask) -> None:
        """Handle triggers based on task result."""
        if not task.result:
            return

        schedule = self._schedules.get(task.agent_type)
        if not schedule:
            return

        # Check if result indicates issues
        result_data = task.result.result_data or {}
        has_issues = (
            not task.result.success or
            result_data.get("issues_count", 0) > 0 or
            result_data.get("critical_count", 0) > 0 or
            result_data.get("degraded_count", 0) > 0
        )

        if has_issues and schedule.triggers_on_issues:
            for trigger_agent in schedule.triggers_on_issues:
                if trigger_agent in self._agents:
                    trigger_schedule = self._schedules.get(trigger_agent)
                    await self.schedule_task(
                        agent_type=trigger_agent,
                        task_name=trigger_schedule.default_task if trigger_schedule else "health_check",
                        priority=AgentPriority.HIGH,
                        trigger_type=TriggerType.EVENT,
                        triggered_by=task.task_id,
                    )
                    logger.info(
                        f"Triggered {trigger_agent.value} due to issues in {task.task_id}"
                    )

    async def _handle_task_failure(self, task: ScheduledAgentTask) -> None:
        """Handle triggers on task failure."""
        schedule = self._schedules.get(task.agent_type)
        if not schedule:
            return

        for trigger_agent in schedule.triggers_on_failure:
            if trigger_agent in self._agents:
                trigger_schedule = self._schedules.get(trigger_agent)
                await self.schedule_task(
                    agent_type=trigger_agent,
                    task_name=trigger_schedule.default_task if trigger_schedule else "health_check",
                    priority=AgentPriority.CRITICAL,
                    trigger_type=TriggerType.ALERT,
                    parameters={"failed_task_id": task.task_id, "error": task.error},
                    triggered_by=task.task_id,
                )
                logger.info(
                    f"Triggered {trigger_agent.value} due to failure in {task.task_id}"
                )

    # =========================================================================
    # Scheduler Loop
    # =========================================================================

    async def _schedule_loop(self, agent_type: AgentType) -> None:
        """Scheduling loop for a single agent."""
        schedule = self._schedules.get(agent_type)
        if not schedule or not schedule.enabled:
            return

        interval_seconds = schedule.interval_hours * 3600

        while self._running:
            try:
                # Schedule the default task
                await self.schedule_task(
                    agent_type=agent_type,
                    task_name=schedule.default_task,
                    priority=schedule.priority,
                    trigger_type=TriggerType.SCHEDULED,
                )

                # Wait for next interval
                await asyncio.sleep(interval_seconds)

            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.error(f"Schedule loop error for {agent_type.value}: {e}")
                await asyncio.sleep(60)  # Brief pause before retry

    async def _execution_loop(self) -> None:
        """Main execution loop - processes task queue."""
        while self._running:
            try:
                task = await self._get_next_task()

                if task:
                    # Execute in background
                    asyncio.create_task(self._execute_task(task))
                else:
                    # No tasks, brief pause
                    await asyncio.sleep(1)

            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.error(f"Execution loop error: {e}")
                await asyncio.sleep(1)

    async def start(self) -> None:
        """Start the multi-agent scheduler."""
        if self._running:
            logger.warning("Scheduler already running")
            return

        self._running = True
        self._started_at = datetime.now(timezone.utc)

        logger.info(f"Starting multi-agent scheduler with {len(self._agents)} agents")

        # Start execution loop
        execution_task = asyncio.create_task(self._execution_loop())

        # Start schedule loops for each agent
        schedule_tasks = []
        for agent_type in self._agents:
            task = asyncio.create_task(self._schedule_loop(agent_type))
            schedule_tasks.append(task)

        # Wait for all loops
        try:
            await asyncio.gather(execution_task, *schedule_tasks)
        except asyncio.CancelledError:
            pass
        finally:
            self._running = False

    def stop(self) -> None:
        """Stop the scheduler."""
        logger.info("Stopping multi-agent scheduler")
        self._running = False

    async def run_once(self) -> list[ScheduledAgentTask]:
        """
        Run all agents once (useful for testing).

        Returns:
            List of completed tasks
        """
        tasks = []

        for agent_type, schedule in self._schedules.items():
            if schedule.enabled:
                task = await self.schedule_task(
                    agent_type=agent_type,
                    task_name=schedule.default_task,
                    priority=schedule.priority,
                    trigger_type=TriggerType.MANUAL,
                )
                tasks.append(task)

        # Execute all tasks
        await asyncio.gather(*[
            self._execute_task(task) for task in tasks
        ])

        return tasks

    # =========================================================================
    # Status and Statistics
    # =========================================================================

    def get_stats(self) -> dict[str, Any]:
        """Get scheduler statistics."""
        if self._started_at:
            self._stats.uptime_seconds = (
                datetime.now(timezone.utc) - self._started_at
            ).total_seconds()

        return {
            "running": self._running,
            "registered_agents": list(self._agents.keys()),
            "total_tasks_scheduled": self._stats.total_tasks_scheduled,
            "total_tasks_completed": self._stats.total_tasks_completed,
            "total_tasks_failed": self._stats.total_tasks_failed,
            "tasks_by_agent": self._stats.tasks_by_agent,
            "average_duration_ms": self._stats.average_duration_ms,
            "pending_tasks": len(self._task_queue),
            "running_tasks": len(self._running_tasks),
            "uptime_seconds": self._stats.uptime_seconds,
        }

    def get_queue_status(self) -> list[dict[str, Any]]:
        """Get current task queue status."""
        return [
            {
                "task_id": task.task_id,
                "agent_type": task.agent_type.value,
                "task_name": task.task_name,
                "priority": task.priority.value,
                "status": task.status.value,
                "created_at": task.created_at.isoformat(),
            }
            for task in self._task_queue
        ]

    def get_running_tasks(self) -> list[dict[str, Any]]:
        """Get currently running tasks."""
        return [
            {
                "task_id": task.task_id,
                "agent_type": task.agent_type.value,
                "task_name": task.task_name,
                "started_at": task.started_at.isoformat() if task.started_at else None,
            }
            for task in self._running_tasks.values()
        ]

    def get_recent_completions(self, limit: int = 10) -> list[dict[str, Any]]:
        """Get recently completed tasks."""
        return [
            {
                "task_id": task.task_id,
                "agent_type": task.agent_type.value,
                "task_name": task.task_name,
                "status": task.status.value,
                "started_at": task.started_at.isoformat() if task.started_at else None,
                "completed_at": task.completed_at.isoformat() if task.completed_at else None,
                "success": task.result.success if task.result else False,
                "error": task.error,
            }
            for task in self._completed_tasks[-limit:][::-1]
        ]


# ============================================================================
# Exports
# ============================================================================

__all__ = [
    "AgentType",
    "TaskStatus",
    "ScheduledAgentTask",
    "AgentScheduleConfig",
    "SchedulerStats",
    "MultiAgentScheduler",
]
