"""BLCE Agent Runtime — ThreadPoolExecutor-based engine with dependency-aware scheduling.

Executes a directed acyclic graph of ``AgentTask`` objects, respecting
dependency edges.  Independent tasks within the same topological level
run in parallel via ``concurrent.futures.ThreadPoolExecutor``.

Usage::

    rt = AgentRuntime(max_workers=4, timeout=300)
    rt.add_task(AgentTask(task_id="a", name="parse", callable=parse_fn))
    rt.add_task(AgentTask(task_id="b", name="norm", callable=norm_fn, depends_on=["a"]))
    result = rt.run()
"""
from __future__ import annotations

import logging
import threading
import time
from collections import deque
from concurrent.futures import Future, ThreadPoolExecutor, as_completed
from enum import Enum
from typing import Any, Callable, Dict, List, Optional

from .contracts import RuntimeResult
from .message_bus import MessageBus

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Enums / data classes
# ---------------------------------------------------------------------------

class TaskStatus(Enum):
    """Lifecycle states for an agent task."""
    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"
    SKIPPED = "skipped"


class AgentTask:
    """Wraps a callable with metadata and dependency info.

    Parameters
    ----------
    task_id : str
        Unique identifier.
    name : str
        Human-readable name (e.g. the phase name).
    callable : Callable
        The function to execute.  Must accept ``**kwargs``.
    kwargs : dict, optional
        Keyword arguments forwarded to *callable*.
    depends_on : list[str], optional
        Task IDs that must complete before this task runs.
    """

    def __init__(
        self,
        task_id: str,
        name: str,
        callable: Callable,
        kwargs: Optional[Dict[str, Any]] = None,
        depends_on: Optional[List[str]] = None,
    ):
        self.task_id = task_id
        self.name = name
        self.callable = callable
        self.kwargs: Dict[str, Any] = kwargs or {}
        self.depends_on: List[str] = list(depends_on or [])
        self.status: TaskStatus = TaskStatus.PENDING
        self.result: Any = None
        self.error: Optional[str] = None
        self.duration: float = 0.0

    def to_dict(self) -> Dict[str, Any]:
        return {
            "task_id": self.task_id,
            "name": self.name,
            "status": self.status.value,
            "depends_on": self.depends_on,
            "duration": round(self.duration, 3),
            "error": self.error,
        }


# ---------------------------------------------------------------------------
# Runtime
# ---------------------------------------------------------------------------

class AgentRuntime:
    """Dependency-aware task executor.

    Parameters
    ----------
    max_workers : int
        Thread pool size.
    timeout : int
        Per-task timeout in seconds.
    bus : MessageBus, optional
        If provided, task lifecycle events are published.
    """

    def __init__(
        self,
        max_workers: int = 4,
        timeout: int = 300,
        bus: Optional[MessageBus] = None,
    ):
        self.max_workers = max_workers
        self.timeout = timeout
        self.bus = bus or MessageBus()
        self._tasks: Dict[str, AgentTask] = {}
        self._lock = threading.Lock()

    # -- task registration --------------------------------------------------

    def add_task(self, task: AgentTask) -> None:
        """Register a task in the graph."""
        self._tasks[task.task_id] = task

    @property
    def tasks(self) -> Dict[str, AgentTask]:
        return dict(self._tasks)

    # -- execution ----------------------------------------------------------

    def run(self) -> RuntimeResult:
        """Execute all tasks respecting dependency order.

        Returns a ``RuntimeResult`` summarising the run.
        """
        t0 = time.time()

        if not self._tasks:
            return RuntimeResult(
                tasks_submitted=0,
                total_duration=0.0,
                level_count=0,
                messages_published=self.bus.message_count,
            )

        levels = self._resolve_levels()
        completed = 0
        failed = 0
        skipped = 0

        with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
            for level_tasks in levels:
                level_results = self._execute_level(level_tasks, executor)
                for tid, success in level_results.items():
                    task = self._tasks[tid]
                    if success:
                        completed += 1
                    elif task.status == TaskStatus.SKIPPED:
                        # Already counted via _skip_dependents
                        pass
                    else:
                        failed += 1
                        # Skip dependents of failed tasks
                        skipped += self._skip_dependents(tid)

        total_duration = round(time.time() - t0, 3)

        return RuntimeResult(
            tasks_submitted=len(self._tasks),
            tasks_completed=completed,
            tasks_failed=failed,
            tasks_skipped=skipped,
            total_duration=total_duration,
            level_count=len(levels),
            messages_published=self.bus.message_count,
            task_details=[t.to_dict() for t in self._tasks.values()],
        )

    # -- topology -----------------------------------------------------------

    def _resolve_levels(self) -> List[List[str]]:
        """Kahn's algorithm — topological sort into parallelisable levels.

        Returns a list of levels, each level a list of task_ids that can
        execute concurrently.

        Raises ``ValueError`` on circular dependencies.
        """
        in_degree: Dict[str, int] = {tid: 0 for tid in self._tasks}
        dependents: Dict[str, List[str]] = {tid: [] for tid in self._tasks}

        for tid, task in self._tasks.items():
            for dep in task.depends_on:
                if dep not in self._tasks:
                    continue  # Ignore missing dependencies
                in_degree[tid] += 1
                dependents[dep].append(tid)

        # Seed queue with zero-in-degree tasks
        queue: deque[str] = deque(
            tid for tid, deg in in_degree.items() if deg == 0
        )

        levels: List[List[str]] = []
        processed = 0

        while queue:
            current_level = list(queue)
            queue.clear()
            levels.append(current_level)
            processed += len(current_level)

            for tid in current_level:
                for dep_tid in dependents[tid]:
                    in_degree[dep_tid] -= 1
                    if in_degree[dep_tid] == 0:
                        queue.append(dep_tid)

        if processed != len(self._tasks):
            raise ValueError(
                "Circular dependency detected in task graph. "
                f"Processed {processed}/{len(self._tasks)} tasks."
            )

        return levels

    # -- level execution ----------------------------------------------------

    def _execute_level(
        self,
        task_ids: List[str],
        executor: ThreadPoolExecutor,
    ) -> Dict[str, bool]:
        """Submit all tasks in a level and wait for completion.

        Returns ``{task_id: success_bool}``
        """
        futures: Dict[Future, str] = {}
        results: Dict[str, bool] = {}

        for tid in task_ids:
            task = self._tasks[tid]
            if task.status == TaskStatus.SKIPPED:
                results[tid] = False
                continue
            fut = executor.submit(self._execute_task, task)
            futures[fut] = tid

        for fut in as_completed(futures):
            tid = futures[fut]
            try:
                success = fut.result(timeout=self.timeout)
                results[tid] = success
            except Exception as exc:
                task = self._tasks[tid]
                task.status = TaskStatus.FAILED
                task.error = str(exc)
                results[tid] = False
                self.bus.publish(
                    "agent.failed",
                    sender=task.name,
                    payload={"task_id": tid, "error": str(exc)},
                )

        return results

    def _execute_task(self, task: AgentTask) -> bool:
        """Run a single task, updating its status and publishing to the bus."""
        task.status = TaskStatus.RUNNING
        self.bus.publish(
            "agent.started",
            sender=task.name,
            payload={"task_id": task.task_id},
        )

        t0 = time.time()
        try:
            task.result = task.callable(**task.kwargs)
            task.status = TaskStatus.COMPLETED
            task.duration = round(time.time() - t0, 3)
            self.bus.publish(
                "agent.completed",
                sender=task.name,
                payload={
                    "task_id": task.task_id,
                    "duration": task.duration,
                },
            )
            return True
        except Exception as exc:
            task.status = TaskStatus.FAILED
            task.error = str(exc)
            task.duration = round(time.time() - t0, 3)
            self.bus.publish(
                "agent.failed",
                sender=task.name,
                payload={"task_id": task.task_id, "error": str(exc)},
            )
            return False

    # -- helpers ------------------------------------------------------------

    def _skip_dependents(self, failed_tid: str) -> int:
        """Mark all transitive dependents of *failed_tid* as SKIPPED."""
        skipped = 0
        queue = deque([failed_tid])
        visited = {failed_tid}
        while queue:
            current = queue.popleft()
            for tid, task in self._tasks.items():
                if tid in visited:
                    continue
                if current in task.depends_on:
                    if task.status == TaskStatus.PENDING:
                        task.status = TaskStatus.SKIPPED
                        task.error = f"Skipped — dependency '{current}' failed"
                        skipped += 1
                    visited.add(tid)
                    queue.append(tid)
        return skipped
