# Familiar - Self-Hosted AI Companion Platform
# Copyright (c) 2026 George Scott Foley
# ORCID: 0009-0006-4957-0540
# Email: Georgescottfoley@proton.me
# Licensed under the MIT License - see LICENSE file for details

"""
Multi-Agent Orchestration (v1.4.0 Phase 4 - Enhancement #19)

Coordinate multiple agents working together on complex tasks.
"""

import asyncio
import logging
import threading
import time
import uuid
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from datetime import datetime, timezone
from enum import Enum
from typing import Any, Callable, Dict, List, Optional, Set

logger = logging.getLogger(__name__)


def _utcnow() -> datetime:
    return datetime.now(timezone.utc)


DEFAULT_TIMEOUT = 300.0
DEFAULT_MAX_RETRIES = 3
DEFAULT_MAX_CONCURRENT = 5


class AgentRole(str, Enum):
    SUPERVISOR = "supervisor"
    COORDINATOR = "coordinator"
    SPECIALIST = "specialist"
    GENERALIST = "generalist"
    REVIEWER = "reviewer"
    CRITIC = "critic"
    ROUTER = "router"
    AGGREGATOR = "aggregator"


class AgentStatus(str, Enum):
    IDLE = "idle"
    BUSY = "busy"
    PAUSED = "paused"
    ERROR = "error"
    OFFLINE = "offline"


class MessageType(str, Enum):
    TASK_ASSIGN = "task_assign"
    TASK_RESULT = "task_result"
    TASK_PROGRESS = "task_progress"
    TASK_ERROR = "task_error"
    PING = "ping"
    PONG = "pong"
    PAUSE = "pause"
    RESUME = "resume"
    SHUTDOWN = "shutdown"
    HANDOFF = "handoff"
    QUERY = "query"
    RESPONSE = "response"
    BROADCAST = "broadcast"
    FEEDBACK = "feedback"
    APPROVAL = "approval"
    REJECTION = "rejection"


class OrchestrationStrategy(str, Enum):
    SEQUENTIAL = "sequential"
    PARALLEL = "parallel"
    PIPELINE = "pipeline"
    HIERARCHICAL = "hierarchical"
    CONSENSUS = "consensus"
    DEBATE = "debate"
    ROUND_ROBIN = "round_robin"


class TaskStatus(str, Enum):
    PENDING = "pending"
    ASSIGNED = "assigned"
    IN_PROGRESS = "in_progress"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELLED = "cancelled"


@dataclass
class Message:
    id: str = field(default_factory=lambda: str(uuid.uuid4()))
    type: MessageType = MessageType.QUERY
    sender_id: Optional[str] = None
    recipient_id: Optional[str] = None
    content: Any = None
    metadata: Dict[str, Any] = field(default_factory=dict)
    timestamp: datetime = field(default_factory=_utcnow)
    expires_at: Optional[datetime] = None
    correlation_id: Optional[str] = None
    reply_to: Optional[str] = None

    def to_dict(self) -> Dict[str, Any]:
        return {
            "id": self.id,
            "type": self.type.value,
            "sender_id": self.sender_id,
            "recipient_id": self.recipient_id,
            "content": self.content,
            "metadata": self.metadata,
            "timestamp": self.timestamp.isoformat(),
            "expires_at": self.expires_at.isoformat() if self.expires_at else None,
            "correlation_id": self.correlation_id,
            "reply_to": self.reply_to,
        }

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "Message":
        data = data.copy()
        data["type"] = MessageType(data["type"])
        data["timestamp"] = datetime.fromisoformat(data["timestamp"])
        if data.get("expires_at"):
            data["expires_at"] = datetime.fromisoformat(data["expires_at"])
        return cls(**data)


@dataclass
class Task:
    id: str = field(default_factory=lambda: str(uuid.uuid4()))
    description: str = ""
    assigned_to: Optional[str] = None
    created_by: Optional[str] = None
    input_data: Any = None
    output_data: Any = None
    context: Dict[str, Any] = field(default_factory=dict)
    status: TaskStatus = TaskStatus.PENDING
    error: Optional[str] = None
    created_at: datetime = field(default_factory=_utcnow)
    started_at: Optional[datetime] = None
    completed_at: Optional[datetime] = None
    timeout: float = DEFAULT_TIMEOUT
    depends_on: List[str] = field(default_factory=list)
    subtasks: List[str] = field(default_factory=list)
    parent_task: Optional[str] = None
    priority: int = 0
    retries: int = 0
    max_retries: int = DEFAULT_MAX_RETRIES
    metadata: Dict[str, Any] = field(default_factory=dict)

    @property
    def elapsed_ms(self) -> float:
        if not self.started_at:
            return 0.0
        end = self.completed_at or _utcnow()
        return (end - self.started_at).total_seconds() * 1000

    @property
    def is_complete(self) -> bool:
        return self.status in (TaskStatus.COMPLETED, TaskStatus.FAILED, TaskStatus.CANCELLED)

    def to_dict(self) -> Dict[str, Any]:
        return {
            "id": self.id,
            "description": self.description,
            "assigned_to": self.assigned_to,
            "created_by": self.created_by,
            "input_data": self.input_data,
            "output_data": self.output_data,
            "context": self.context,
            "status": self.status.value,
            "error": self.error,
            "created_at": self.created_at.isoformat(),
            "started_at": self.started_at.isoformat() if self.started_at else None,
            "completed_at": self.completed_at.isoformat() if self.completed_at else None,
            "timeout": self.timeout,
            "depends_on": self.depends_on,
            "subtasks": self.subtasks,
            "parent_task": self.parent_task,
            "priority": self.priority,
            "retries": self.retries,
            "max_retries": self.max_retries,
            "metadata": self.metadata,
        }


@dataclass
class TaskResult:
    task_id: str
    success: bool
    output: Any = None
    error: Optional[str] = None
    agent_id: Optional[str] = None
    elapsed_ms: float = 0.0
    metadata: Dict[str, Any] = field(default_factory=dict)


class AgentBase(ABC):
    def __init__(
        self,
        agent_id: Optional[str] = None,
        name: str = "Agent",
        role: AgentRole = AgentRole.GENERALIST,
        capabilities: Optional[List[str]] = None,
    ):
        self.id = agent_id or str(uuid.uuid4())
        self.name = name
        self.role = role
        self.capabilities = capabilities or []
        self.status = AgentStatus.IDLE
        self._inbox: asyncio.Queue = asyncio.Queue()
        self._outbox: asyncio.Queue = asyncio.Queue()
        self._current_task: Optional[Task] = None
        self._task_history: List[str] = []
        self._metadata: Dict[str, Any] = {}

    @abstractmethod
    async def execute(self, task: Task) -> TaskResult:
        pass

    async def receive_message(self, message: Message):
        await self._inbox.put(message)

    async def send_message(self, message: Message):
        message.sender_id = self.id
        await self._outbox.put(message)

    async def get_outgoing_messages(self) -> List[Message]:
        messages = []
        while not self._outbox.empty():
            try:
                msg = self._outbox.get_nowait()
                messages.append(msg)
            except asyncio.QueueEmpty:
                break
        return messages

    def can_handle(self, task: Task) -> bool:
        required = task.metadata.get("required_capabilities", [])
        if not required:
            return True
        return any(cap in self.capabilities for cap in required)

    def get_info(self) -> Dict[str, Any]:
        return {
            "id": self.id,
            "name": self.name,
            "role": self.role.value,
            "status": self.status.value,
            "capabilities": self.capabilities,
            "current_task": self._current_task.id if self._current_task else None,
        }


class FunctionAgent(AgentBase):
    def __init__(self, fn: Callable[[Task], Any], **kwargs):
        super().__init__(**kwargs)
        self._fn = fn
        self._is_async = asyncio.iscoroutinefunction(fn)

    async def execute(self, task: Task) -> TaskResult:
        self.status = AgentStatus.BUSY
        self._current_task = task
        start_time = time.time()
        try:
            if self._is_async:
                output = await self._fn(task)
            else:
                output = self._fn(task)
            elapsed = (time.time() - start_time) * 1000
            return TaskResult(
                task_id=task.id, success=True, output=output, agent_id=self.id, elapsed_ms=elapsed
            )
        except Exception as e:
            elapsed = (time.time() - start_time) * 1000
            return TaskResult(
                task_id=task.id, success=False, error=str(e), agent_id=self.id, elapsed_ms=elapsed
            )
        finally:
            self.status = AgentStatus.IDLE
            self._current_task = None


class LLMAgent(AgentBase):
    def __init__(self, llm_provider: Callable, system_prompt: Optional[str] = None, **kwargs):
        super().__init__(**kwargs)
        self._llm = llm_provider
        self._system_prompt = system_prompt or f"You are {self.name}, a {self.role.value} agent."
        self._is_async = asyncio.iscoroutinefunction(llm_provider)

    async def execute(self, task: Task) -> TaskResult:
        self.status = AgentStatus.BUSY
        self._current_task = task
        start_time = time.time()
        try:
            messages = [{"role": "system", "content": self._system_prompt}]
            if task.context:
                context_str = "\n".join(f"{k}: {v}" for k, v in task.context.items())
                messages.append({"role": "system", "content": f"Context:\n{context_str}"})
            messages.append(
                {
                    "role": "user",
                    "content": task.description
                    + (f"\n\nInput: {task.input_data}" if task.input_data else ""),
                }
            )
            if self._is_async:
                response = await self._llm(messages=messages)
            else:
                response = self._llm(messages=messages)
            if isinstance(response, str):
                output = response
            elif hasattr(response, "text"):
                output = response.text
            elif hasattr(response, "content"):
                output = response.content
            else:
                output = str(response)
            elapsed = (time.time() - start_time) * 1000
            return TaskResult(
                task_id=task.id, success=True, output=output, agent_id=self.id, elapsed_ms=elapsed
            )
        except Exception as e:
            elapsed = (time.time() - start_time) * 1000
            return TaskResult(
                task_id=task.id, success=False, error=str(e), agent_id=self.id, elapsed_ms=elapsed
            )
        finally:
            self.status = AgentStatus.IDLE
            self._current_task = None


class AgentRegistry:
    def __init__(self):
        self._agents: Dict[str, AgentBase] = {}
        self._by_role: Dict[AgentRole, Set[str]] = {}
        self._by_capability: Dict[str, Set[str]] = {}
        self._lock = threading.RLock()

    def register(self, agent: AgentBase) -> str:
        with self._lock:
            self._agents[agent.id] = agent
            if agent.role not in self._by_role:
                self._by_role[agent.role] = set()
            self._by_role[agent.role].add(agent.id)
            for cap in agent.capabilities:
                if cap not in self._by_capability:
                    self._by_capability[cap] = set()
                self._by_capability[cap].add(agent.id)
            return agent.id

    def unregister(self, agent_id: str) -> bool:
        with self._lock:
            agent = self._agents.get(agent_id)
            if not agent:
                return False
            del self._agents[agent_id]
            if agent.role in self._by_role:
                self._by_role[agent.role].discard(agent_id)
            for cap in agent.capabilities:
                if cap in self._by_capability:
                    self._by_capability[cap].discard(agent_id)
            return True

    def get(self, agent_id: str) -> Optional[AgentBase]:
        return self._agents.get(agent_id)

    def get_by_role(self, role: AgentRole) -> List[AgentBase]:
        with self._lock:
            ids = self._by_role.get(role, set())
            return [self._agents[aid] for aid in ids if aid in self._agents]

    def get_by_capability(self, capability: str) -> List[AgentBase]:
        with self._lock:
            ids = self._by_capability.get(capability, set())
            return [self._agents[aid] for aid in ids if aid in self._agents]

    def get_available(self) -> List[AgentBase]:
        with self._lock:
            return [a for a in self._agents.values() if a.status == AgentStatus.IDLE]

    def get_all(self) -> List[AgentBase]:
        with self._lock:
            return list(self._agents.values())

    def find_agent_for_task(self, task: Task) -> Optional[AgentBase]:
        with self._lock:
            available = self.get_available()
            suitable = [a for a in available if a.can_handle(task)]
            if not suitable:
                return None
            specialists = [a for a in suitable if a.role == AgentRole.SPECIALIST]
            if specialists:
                return specialists[0]
            return suitable[0]

    def __len__(self) -> int:
        return len(self._agents)

    def __contains__(self, agent_id: str) -> bool:
        return agent_id in self._agents


class MessageBus:
    def __init__(self, registry: AgentRegistry):
        self._registry = registry
        self._subscribers: Dict[str, Set[str]] = {}
        self._message_history: List[Message] = []
        self._max_history = 1000
        self._lock = threading.RLock()

    async def send(self, message: Message):
        with self._lock:
            self._message_history.append(message)
            if len(self._message_history) > self._max_history:
                self._message_history = self._message_history[-self._max_history :]
        if message.recipient_id:
            agent = self._registry.get(message.recipient_id)
            if agent:
                await agent.receive_message(message)
        else:
            for agent in self._registry.get_all():
                if agent.id != message.sender_id:
                    await agent.receive_message(message)

    async def publish(self, topic: str, message: Message):
        with self._lock:
            subscriber_ids = self._subscribers.get(topic, set())
        for agent_id in subscriber_ids:
            agent = self._registry.get(agent_id)
            if agent:
                await agent.receive_message(message)

    def subscribe(self, agent_id: str, topic: str):
        with self._lock:
            if topic not in self._subscribers:
                self._subscribers[topic] = set()
            self._subscribers[topic].add(agent_id)

    def unsubscribe(self, agent_id: str, topic: str):
        with self._lock:
            if topic in self._subscribers:
                self._subscribers[topic].discard(agent_id)

    def get_history(
        self, limit: int = 100, message_type: Optional[MessageType] = None
    ) -> List[Message]:
        with self._lock:
            messages = self._message_history
            if message_type:
                messages = [m for m in messages if m.type == message_type]
            return messages[-limit:]


@dataclass
class OrchestrationResult:
    success: bool
    output: Any = None
    tasks_completed: int = 0
    tasks_failed: int = 0
    task_results: List[TaskResult] = field(default_factory=list)
    elapsed_ms: float = 0.0
    errors: List[str] = field(default_factory=list)
    strategy: Optional[OrchestrationStrategy] = None
    agents_used: List[str] = field(default_factory=list)
    metadata: Dict[str, Any] = field(default_factory=dict)


class Orchestrator:
    def __init__(
        self, max_concurrent: int = DEFAULT_MAX_CONCURRENT, default_timeout: float = DEFAULT_TIMEOUT
    ):
        self.registry = AgentRegistry()
        self.message_bus = MessageBus(self.registry)
        self.max_concurrent = max_concurrent
        self.default_timeout = default_timeout
        self._tasks: Dict[str, Task] = {}
        self._task_queue: asyncio.Queue = asyncio.Queue()
        self._results: Dict[str, TaskResult] = {}
        self._stats = {
            "total_tasks": 0,
            "completed_tasks": 0,
            "failed_tasks": 0,
            "total_orchestrations": 0,
        }
        self._lock = threading.RLock()

    def register(self, agent: AgentBase, role: Optional[AgentRole] = None) -> str:
        if role:
            agent.role = role
        return self.registry.register(agent)

    def unregister(self, agent_id: str) -> bool:
        return self.registry.unregister(agent_id)

    async def execute(
        self,
        task,
        strategy: OrchestrationStrategy = OrchestrationStrategy.SEQUENTIAL,
        agents: Optional[List[str]] = None,
        context: Optional[Dict[str, Any]] = None,
        timeout: Optional[float] = None,
    ) -> OrchestrationResult:
        start_time = time.time()
        timeout = timeout or self.default_timeout
        if isinstance(task, str):
            task = Task(description=task, context=context or {})
        elif context:
            task.context.update(context)
        self._stats["total_orchestrations"] += 1
        try:
            if strategy == OrchestrationStrategy.SEQUENTIAL:
                result = await self._execute_sequential(task, agents, timeout)
            elif strategy == OrchestrationStrategy.PARALLEL:
                result = await self._execute_parallel(task, agents, timeout)
            elif strategy == OrchestrationStrategy.PIPELINE:
                result = await self._execute_sequential(task, agents, timeout)
            elif strategy == OrchestrationStrategy.HIERARCHICAL:
                result = await self._execute_hierarchical(task, agents, timeout)
            elif strategy == OrchestrationStrategy.CONSENSUS:
                result = await self._execute_consensus(task, agents, timeout)
            elif strategy == OrchestrationStrategy.DEBATE:
                result = await self._execute_debate(task, agents, timeout)
            else:
                result = await self._execute_sequential(task, agents, timeout)
            result.strategy = strategy
            result.elapsed_ms = (time.time() - start_time) * 1000
            return result
        except asyncio.TimeoutError:
            return OrchestrationResult(
                success=False,
                errors=["Orchestration timed out"],
                strategy=strategy,
                elapsed_ms=(time.time() - start_time) * 1000,
            )
        except Exception as e:
            return OrchestrationResult(
                success=False,
                errors=[str(e)],
                strategy=strategy,
                elapsed_ms=(time.time() - start_time) * 1000,
            )

    async def delegate(self, task: Task, agent_id: Optional[str] = None) -> TaskResult:
        if agent_id:
            agent = self.registry.get(agent_id)
        else:
            agent = self.registry.find_agent_for_task(task)
        if not agent:
            return TaskResult(task_id=task.id, success=False, error="No suitable agent available")
        task.assigned_to = agent.id
        task.status = TaskStatus.ASSIGNED
        task.started_at = _utcnow()
        with self._lock:
            self._tasks[task.id] = task
            self._stats["total_tasks"] += 1
        try:
            result = await asyncio.wait_for(agent.execute(task), timeout=task.timeout)
            task.status = TaskStatus.COMPLETED if result.success else TaskStatus.FAILED
            task.completed_at = _utcnow()
            task.output_data = result.output
            task.error = result.error
            with self._lock:
                self._results[task.id] = result
                if result.success:
                    self._stats["completed_tasks"] += 1
                else:
                    self._stats["failed_tasks"] += 1
            return result
        except asyncio.TimeoutError:
            task.status = TaskStatus.FAILED
            task.error = "Task timed out"
            return TaskResult(
                task_id=task.id, success=False, error="Task timed out", agent_id=agent.id
            )

    async def handoff(
        self,
        from_agent_id: str,
        to_agent_id: str,
        task: Task,
        context: Optional[Dict[str, Any]] = None,
    ) -> bool:
        to_agent = self.registry.get(to_agent_id)
        if not to_agent:
            return False
        task.context["handoff_from"] = from_agent_id
        if context:
            task.context.update(context)
        message = Message(
            type=MessageType.HANDOFF,
            sender_id=from_agent_id,
            recipient_id=to_agent_id,
            content={
                "task_id": task.id,
                "task_description": task.description,
                "context": task.context,
            },
        )
        await self.message_bus.send(message)
        return True

    def get_stats(self) -> Dict[str, Any]:
        with self._lock:
            return {
                **self._stats,
                "registered_agents": len(self.registry),
                "active_tasks": sum(
                    1
                    for t in self._tasks.values()
                    if t.status in (TaskStatus.ASSIGNED, TaskStatus.IN_PROGRESS)
                ),
            }

    async def _execute_sequential(
        self, task: Task, agent_ids: Optional[List[str]], timeout: float
    ) -> OrchestrationResult:
        agents = self._get_agents(agent_ids)
        if not agents:
            return OrchestrationResult(success=False, errors=["No agents available"])
        results = []
        agents_used = []
        output = task.input_data
        for agent in agents:
            subtask = Task(
                description=task.description,
                input_data=output,
                context=task.context.copy(),
                timeout=timeout / len(agents),
            )
            result = await self.delegate(subtask, agent.id)
            results.append(result)
            agents_used.append(agent.id)
            if not result.success:
                return OrchestrationResult(
                    success=False,
                    output=output,
                    task_results=results,
                    tasks_completed=len([r for r in results if r.success]),
                    tasks_failed=len([r for r in results if not r.success]),
                    errors=[result.error or "Task failed"],
                    agents_used=agents_used,
                )
            output = result.output
        return OrchestrationResult(
            success=True,
            output=output,
            task_results=results,
            tasks_completed=len(results),
            tasks_failed=0,
            agents_used=agents_used,
        )

    async def _execute_parallel(
        self, task: Task, agent_ids: Optional[List[str]], timeout: float
    ) -> OrchestrationResult:
        agents = self._get_agents(agent_ids)
        if not agents:
            return OrchestrationResult(success=False, errors=["No agents available"])
        subtasks = [
            (
                Task(
                    description=task.description,
                    input_data=task.input_data,
                    context=task.context.copy(),
                    timeout=timeout,
                ),
                agent,
            )
            for agent in agents
        ]

        async def run_agent(subtask: Task, agent: AgentBase) -> TaskResult:
            return await self.delegate(subtask, agent.id)

        results = await asyncio.gather(
            *[run_agent(st, ag) for st, ag in subtasks], return_exceptions=True
        )
        task_results = []
        agents_used = []
        outputs = []
        errors = []
        for (subtask, agent), result in zip(subtasks, results):
            agents_used.append(agent.id)
            if isinstance(result, Exception):
                task_results.append(
                    TaskResult(
                        task_id=subtask.id, success=False, error=str(result), agent_id=agent.id
                    )
                )
                errors.append(str(result))
            else:
                task_results.append(result)
                if result.success:
                    outputs.append(result.output)
                else:
                    errors.append(result.error or "Task failed")
        success = len(outputs) > 0
        return OrchestrationResult(
            success=success,
            output=outputs if success else None,
            task_results=task_results,
            tasks_completed=len([r for r in task_results if r.success]),
            tasks_failed=len([r for r in task_results if not r.success]),
            errors=errors,
            agents_used=agents_used,
        )

    async def _execute_hierarchical(
        self, task: Task, agent_ids: Optional[List[str]], timeout: float
    ) -> OrchestrationResult:
        agents = self._get_agents(agent_ids)
        supervisors = [a for a in agents if a.role == AgentRole.SUPERVISOR]
        workers = [a for a in agents if a.role != AgentRole.SUPERVISOR]
        if not supervisors:
            return await self._execute_parallel(task, agent_ids, timeout)
        supervisor = supervisors[0]
        if not workers:
            result = await self.delegate(task, supervisor.id)
            return OrchestrationResult(
                success=result.success,
                output=result.output,
                task_results=[result],
                tasks_completed=1 if result.success else 0,
                tasks_failed=0 if result.success else 1,
                errors=[result.error] if result.error else [],
                agents_used=[supervisor.id],
            )
        planning_task = Task(
            description=f"Break down this task into subtasks for {len(workers)} workers: {task.description}",
            input_data=task.input_data,
            context={"worker_count": len(workers)},
            timeout=timeout / 4,
        )
        plan_result = await self.delegate(planning_task, supervisor.id)
        if not plan_result.success:
            return OrchestrationResult(
                success=False,
                errors=["Supervisor failed to create plan"],
                task_results=[plan_result],
                agents_used=[supervisor.id],
            )
        worker_results = await self._execute_parallel(
            Task(description=task.description, input_data=plan_result.output, context=task.context),
            [w.id for w in workers],
            timeout * 0.5,
        )
        aggregate_task = Task(
            description=f"Aggregate these results: {worker_results.output}",
            input_data=worker_results.output,
            timeout=timeout / 4,
        )
        final_result = await self.delegate(aggregate_task, supervisor.id)
        all_results = [plan_result] + worker_results.task_results + [final_result]
        all_agents = [supervisor.id] + worker_results.agents_used + [supervisor.id]
        return OrchestrationResult(
            success=final_result.success,
            output=final_result.output,
            task_results=all_results,
            tasks_completed=len([r for r in all_results if r.success]),
            tasks_failed=len([r for r in all_results if not r.success]),
            agents_used=list(set(all_agents)),
        )

    async def _execute_consensus(
        self, task: Task, agent_ids: Optional[List[str]], timeout: float
    ) -> OrchestrationResult:
        parallel_result = await self._execute_parallel(task, agent_ids, timeout * 0.7)
        if not parallel_result.success:
            return parallel_result
        outputs = parallel_result.output
        if not outputs:
            return parallel_result
        parallel_result.output = outputs[0]
        parallel_result.metadata["all_outputs"] = outputs
        parallel_result.metadata["consensus_method"] = "first"
        return parallel_result

    async def _execute_debate(
        self, task: Task, agent_ids: Optional[List[str]], timeout: float
    ) -> OrchestrationResult:
        agents = self._get_agents(agent_ids)
        if len(agents) < 2:
            return await self._execute_sequential(task, agent_ids, timeout)
        agent_a, agent_b = agents[0], agents[1]
        task_a = Task(
            description=f"Argue FOR: {task.description}",
            input_data=task.input_data,
            context=task.context,
            timeout=timeout / 4,
        )
        task_b = Task(
            description=f"Argue AGAINST: {task.description}",
            input_data=task.input_data,
            context=task.context,
            timeout=timeout / 4,
        )
        result_a, result_b = await asyncio.gather(
            self.delegate(task_a, agent_a.id), self.delegate(task_b, agent_b.id)
        )
        rebuttal_a = Task(
            description=f"Respond to this counterargument: {result_b.output}",
            input_data=result_a.output,
            context={"opponent_argument": result_b.output},
            timeout=timeout / 4,
        )
        rebuttal_b = Task(
            description=f"Respond to this argument: {result_a.output}",
            input_data=result_b.output,
            context={"opponent_argument": result_a.output},
            timeout=timeout / 4,
        )
        final_a, final_b = await asyncio.gather(
            self.delegate(rebuttal_a, agent_a.id), self.delegate(rebuttal_b, agent_b.id)
        )
        output = {
            "position_for": final_a.output,
            "position_against": final_b.output,
            "debate_rounds": 2,
        }
        all_results = [result_a, result_b, final_a, final_b]
        return OrchestrationResult(
            success=True,
            output=output,
            task_results=all_results,
            tasks_completed=len([r for r in all_results if r.success]),
            tasks_failed=len([r for r in all_results if not r.success]),
            agents_used=[agent_a.id, agent_b.id],
            metadata={"debate_format": "two_rounds"},
        )

    def _get_agents(self, agent_ids: Optional[List[str]]) -> List[AgentBase]:
        if agent_ids:
            agents = [self.registry.get(aid) for aid in agent_ids]
            return [a for a in agents if a is not None]
        return self.registry.get_available() or self.registry.get_all()


_orchestrator: Optional[Orchestrator] = None
_orchestrator_lock = threading.Lock()


def get_orchestrator() -> Orchestrator:
    global _orchestrator
    with _orchestrator_lock:
        if _orchestrator is None:
            _orchestrator = Orchestrator()
        return _orchestrator


def set_orchestrator(orchestrator: Orchestrator):
    global _orchestrator
    with _orchestrator_lock:
        _orchestrator = orchestrator


def reset_orchestrator():
    global _orchestrator
    with _orchestrator_lock:
        _orchestrator = None


def create_orchestrator(max_concurrent: int = DEFAULT_MAX_CONCURRENT, **kwargs) -> Orchestrator:
    return Orchestrator(max_concurrent=max_concurrent, **kwargs)


def create_function_agent(
    fn: Callable[[Task], Any],
    name: str = "FunctionAgent",
    role: AgentRole = AgentRole.GENERALIST,
    capabilities: Optional[List[str]] = None,
) -> FunctionAgent:
    return FunctionAgent(fn=fn, name=name, role=role, capabilities=capabilities or [])


def create_llm_agent(
    llm_provider: Callable,
    name: str = "LLMAgent",
    role: AgentRole = AgentRole.GENERALIST,
    system_prompt: Optional[str] = None,
    capabilities: Optional[List[str]] = None,
) -> LLMAgent:
    return LLMAgent(
        llm_provider=llm_provider,
        name=name,
        role=role,
        system_prompt=system_prompt,
        capabilities=capabilities or [],
    )


__all__ = [
    "Orchestrator",
    "AgentBase",
    "FunctionAgent",
    "LLMAgent",
    "AgentRegistry",
    "MessageBus",
    "Task",
    "TaskResult",
    "Message",
    "OrchestrationResult",
    "AgentRole",
    "AgentStatus",
    "MessageType",
    "OrchestrationStrategy",
    "TaskStatus",
    "get_orchestrator",
    "set_orchestrator",
    "reset_orchestrator",
    "create_orchestrator",
    "create_function_agent",
    "create_llm_agent",
    "DEFAULT_TIMEOUT",
    "DEFAULT_MAX_RETRIES",
    "DEFAULT_MAX_CONCURRENT",
]
