"""
Orchestration module for multi-expert agent coordination.

Provides infrastructure for Tatlock to orchestrate expert agents
with streaming think updates to keep users informed of progress.

Key pattern: Stream user-facing interactions, use run() internally
to avoid Ollama streaming+tool call bugs.

Supports:
- Single expert delegation with think updates
- Sequential multi-expert execution (task A → task B → task C)
- Parallel multi-expert execution (tasks A, B, C concurrently)
- Result aggregation from multiple experts
- Partial failure handling
"""
import asyncio
from dataclasses import dataclass, field
from enum import Enum
from typing import AsyncGenerator, Optional, Callable, Any

from src.agents.delegation import DelegationTask, DelegationResult, delegate_to_librarian
from src.core.logging_config import get_logger

logger = get_logger(__name__)


class ExecutionMode(str, Enum):
    """Execution mode for multi-expert coordination."""
    SEQUENTIAL = "sequential"  # One at a time, in order
    PARALLEL = "parallel"      # All at once, concurrently


@dataclass
class OrchestrationContext:
    """
    Context for an orchestration session.

    Tracks the user's request, delegation tasks, and results.
    """
    user_message: str
    steward_note: str
    conversation_id: Optional[str] = None


def parse_delegation_from_steward_note(steward_note: str) -> Optional[DelegationTask]:
    """
    Parse a delegation task from Steward's note.

    Looks for the DELEGATE: pattern in the Steward's recommendation.

    Args:
        steward_note: Formatted note from Steward

    Returns:
        DelegationTask if delegation found, None otherwise

    Example:
        >>> note = "DELEGATE: librarian to create a wiki page about CI/CD"
        >>> task = parse_delegation_from_steward_note(note)
        >>> task.expert_name
        'librarian'
        >>> task.task
        'create a wiki page about CI/CD'
    """
    import re

    # Look for DELEGATE: pattern
    # Match: "DELEGATE: expert_name to action description"
    match = re.search(
        r'DELEGATE:\s*(\w+)\s+to\s+(.+?)(?:\n|REASON:|COMPLEXITY:|CONTEXT:|$)',
        steward_note,
        re.IGNORECASE | re.MULTILINE
    )

    if match:
        expert_name = match.group(1).lower()
        task_description = match.group(2).strip()

        # Handle "none" case
        if expert_name == "none":
            return None

        return DelegationTask(
            expert_name=expert_name,
            task=task_description,
        )

    return None


async def execute_delegation(
    task: DelegationTask,
) -> DelegationResult:
    """
    Execute a delegation task.

    Routes to the appropriate expert agent based on expert_name.

    Args:
        task: Delegation task to execute

    Returns:
        DelegationResult from the expert agent
    """
    logger.info(
        "executing_delegation",
        expert=task.expert_name,
        task=task.task[:50],
    )

    if task.expert_name == "librarian":
        return await delegate_to_librarian(
            task=task.task,
            context=task.context,
        )

    # Future experts would be added here:
    # elif task.expert_name == "memory":
    #     return await delegate_to_memory(task.task, task.context)
    # elif task.expert_name == "home_automation":
    #     return await delegate_to_home_automation(task.task, task.context)

    # Unknown expert - return error result
    logger.warning("unknown_expert", expert=task.expert_name)
    return DelegationResult(
        expert_name=task.expert_name,
        task=task.task,
        success=False,
        output="",
        error=f"Unknown expert: {task.expert_name}",
    )


async def orchestrate_with_think_updates(
    user_message: str,
    steward_note: str,
    delegation_task: Optional[DelegationTask] = None,
) -> AsyncGenerator[str, None]:
    """
    Orchestrate expert delegation with streaming think updates.

    Emits <think> updates before and after delegation calls to
    keep the user informed of progress. Expert calls use run()
    internally to avoid Ollama streaming bugs.

    Args:
        user_message: Original user message
        steward_note: Steward's analysis and instructions
        delegation_task: Optional pre-parsed delegation task

    Yields:
        Think update strings and final expert output

    Example:
        >>> async for update in orchestrate_with_think_updates(
        ...     "Create a wiki page about CI/CD",
        ...     "DELEGATE: librarian to create wiki page",
        ... ):
        ...     print(update)
        <think>Consulting The Librarian...</think>
        <think>Delegation complete.</think>
        [Wiki page created successfully...]
    """
    # Parse delegation if not provided
    if delegation_task is None:
        delegation_task = parse_delegation_from_steward_note(steward_note)

    if delegation_task is None:
        # No delegation needed - nothing to orchestrate
        logger.debug("no_delegation_needed")
        return

    # Stream: About to delegate
    expert_display_name = delegation_task.expert_name.title()
    if delegation_task.expert_name == "librarian":
        expert_display_name = "The Librarian"

    yield f"<think>🤝 Consulting {expert_display_name}...</think>\n"

    # Execute delegation (uses run() internally)
    result = await execute_delegation(delegation_task)

    if result.success:
        yield f"<think>✅ {expert_display_name} completed research.</think>\n"

        # Yield the expert's findings
        if result.output:
            yield f"\n{result.output}"
    else:
        yield f"<think>⚠️ {expert_display_name} encountered an issue: {result.error}</think>\n"

    logger.info(
        "orchestration_complete",
        expert=delegation_task.expert_name,
        success=result.success,
    )


def extract_delegation_context(
    steward_note: str,
) -> dict[str, str]:
    """
    Extract context fields from Steward's note.

    Args:
        steward_note: Formatted note from Steward

    Returns:
        Dict with reason, complexity, and context
    """
    import re

    result = {
        "reason": "",
        "complexity": "",
        "context": "",
    }

    # Extract REASON:
    reason_match = re.search(r'REASON:\s*(.+?)(?:\n|COMPLEXITY:|CONTEXT:|$)', steward_note, re.IGNORECASE)
    if reason_match:
        result["reason"] = reason_match.group(1).strip()

    # Extract COMPLEXITY:
    complexity_match = re.search(r'COMPLEXITY:\s*(.+?)(?:\n|CONTEXT:|$)', steward_note, re.IGNORECASE)
    if complexity_match:
        result["complexity"] = complexity_match.group(1).strip()

    # Extract CONTEXT:
    context_match = re.search(r'CONTEXT:\s*(.+?)$', steward_note, re.IGNORECASE | re.MULTILINE)
    if context_match:
        result["context"] = context_match.group(1).strip()

    return result


# ============================================================================
# Multi-Expert Coordination
# ============================================================================

@dataclass
class MultiExpertResult:
    """
    Aggregated result from multiple expert delegations.

    Attributes:
        results: Dict mapping expert name to their result
        all_succeeded: True if all delegations succeeded
        failed_experts: List of expert names that failed
        combined_output: Aggregated output from all successful experts
    """
    results: dict[str, DelegationResult] = field(default_factory=dict)
    all_succeeded: bool = True
    failed_experts: list[str] = field(default_factory=list)
    combined_output: str = ""

    def add_result(self, result: DelegationResult) -> None:
        """Add a result and update aggregation state."""
        self.results[result.expert_name] = result
        if not result.success:
            self.all_succeeded = False
            self.failed_experts.append(result.expert_name)

    def aggregate_outputs(self, separator: str = "\n\n---\n\n") -> str:
        """Combine all successful outputs into one string."""
        outputs = []
        for expert_name, result in self.results.items():
            if result.success and result.output:
                outputs.append(f"**{expert_name.title()}**: {result.output}")

        self.combined_output = separator.join(outputs)
        return self.combined_output


async def execute_sequential(
    tasks: list[DelegationTask],
    stop_on_failure: bool = False,
) -> MultiExpertResult:
    """
    Execute multiple delegation tasks sequentially.

    Tasks run one after another in order. Later tasks can depend on
    earlier results (though this function doesn't handle passing
    results between tasks - that's the orchestrator's job).

    Args:
        tasks: List of delegation tasks to execute in order
        stop_on_failure: If True, stop execution if any task fails

    Returns:
        MultiExpertResult with all task results

    Example:
        >>> tasks = [
        ...     DelegationTask(expert_name="memory", task="get user location"),
        ...     DelegationTask(expert_name="librarian", task="search weather"),
        ... ]
        >>> result = await execute_sequential(tasks)
        >>> result.all_succeeded
        True
    """
    multi_result = MultiExpertResult()

    logger.info(
        "sequential_execution_started",
        task_count=len(tasks),
        experts=[t.expert_name for t in tasks],
    )

    for i, task in enumerate(tasks):
        logger.debug(
            "sequential_task_executing",
            index=i,
            expert=task.expert_name,
            task=task.task[:50],
        )

        result = await execute_delegation(task)
        multi_result.add_result(result)

        if not result.success and stop_on_failure:
            logger.warning(
                "sequential_execution_stopped",
                failed_at=i,
                expert=task.expert_name,
                error=result.error,
            )
            break

    multi_result.aggregate_outputs()

    logger.info(
        "sequential_execution_complete",
        total_tasks=len(tasks),
        succeeded=len(tasks) - len(multi_result.failed_experts),
        failed=len(multi_result.failed_experts),
    )

    return multi_result


async def execute_parallel(
    tasks: list[DelegationTask],
) -> MultiExpertResult:
    """
    Execute multiple delegation tasks in parallel.

    All tasks run concurrently using asyncio.gather. Use this when
    tasks are independent and don't depend on each other's results.

    Args:
        tasks: List of delegation tasks to execute concurrently

    Returns:
        MultiExpertResult with all task results

    Example:
        >>> tasks = [
        ...     DelegationTask(expert_name="librarian", task="search wiki"),
        ...     DelegationTask(expert_name="memory", task="get preferences"),
        ... ]
        >>> result = await execute_parallel(tasks)
        >>> len(result.results)
        2
    """
    multi_result = MultiExpertResult()

    logger.info(
        "parallel_execution_started",
        task_count=len(tasks),
        experts=[t.expert_name for t in tasks],
    )

    # Execute all tasks concurrently
    results = await asyncio.gather(
        *[execute_delegation(task) for task in tasks],
        return_exceptions=True,
    )

    # Process results
    for i, result in enumerate(results):
        if isinstance(result, Exception):
            # Handle exceptions as failed delegations
            error_result = DelegationResult(
                expert_name=tasks[i].expert_name,
                task=tasks[i].task,
                success=False,
                output="",
                error=str(result),
            )
            multi_result.add_result(error_result)
            logger.error(
                "parallel_task_exception",
                expert=tasks[i].expert_name,
                error=str(result),
            )
        else:
            multi_result.add_result(result)

    multi_result.aggregate_outputs()

    logger.info(
        "parallel_execution_complete",
        total_tasks=len(tasks),
        succeeded=len(tasks) - len(multi_result.failed_experts),
        failed=len(multi_result.failed_experts),
    )

    return multi_result


async def orchestrate_multi_expert(
    tasks: list[DelegationTask],
    mode: ExecutionMode = ExecutionMode.SEQUENTIAL,
    stop_on_failure: bool = False,
) -> AsyncGenerator[str, None]:
    """
    Orchestrate multiple expert delegations with streaming think updates.

    Emits <think> updates for each delegation phase and yields
    combined results at the end.

    Args:
        tasks: List of delegation tasks
        mode: SEQUENTIAL or PARALLEL execution
        stop_on_failure: For sequential mode, stop if a task fails

    Yields:
        Think updates and combined expert output

    Example:
        >>> tasks = [
        ...     DelegationTask(expert_name="memory", task="get location"),
        ...     DelegationTask(expert_name="librarian", task="search weather"),
        ... ]
        >>> async for update in orchestrate_multi_expert(tasks):
        ...     print(update)
        <think>Starting multi-expert coordination (2 tasks)...</think>
        <think>Consulting Memory...</think>
        <think>Memory completed.</think>
        <think>Consulting The Librarian...</think>
        <think>The Librarian completed.</think>
        <think>All experts completed successfully.</think>
        [Combined output from all experts...]
    """
    if not tasks:
        logger.debug("no_tasks_to_orchestrate")
        return

    # Stream: Starting multi-expert coordination
    yield f"<think>🎯 Starting multi-expert coordination ({len(tasks)} tasks, {mode.value})...</think>\n"

    if mode == ExecutionMode.PARALLEL:
        # Parallel execution - emit one update then run all at once
        expert_names = ", ".join(_get_display_name(t.expert_name) for t in tasks)
        yield f"<think>🔄 Consulting in parallel: {expert_names}...</think>\n"

        result = await execute_parallel(tasks)

        # Emit completion updates for each
        for expert_name, expert_result in result.results.items():
            display_name = _get_display_name(expert_name)
            if expert_result.success:
                yield f"<think>✅ {display_name} completed.</think>\n"
            else:
                yield f"<think>⚠️ {display_name} failed: {expert_result.error}</think>\n"

    else:
        # Sequential execution - emit updates for each task
        result = MultiExpertResult()

        for task in tasks:
            display_name = _get_display_name(task.expert_name)
            yield f"<think>🤝 Consulting {display_name}...</think>\n"

            task_result = await execute_delegation(task)
            result.add_result(task_result)

            if task_result.success:
                yield f"<think>✅ {display_name} completed.</think>\n"
            else:
                yield f"<think>⚠️ {display_name} failed: {task_result.error}</think>\n"
                if stop_on_failure:
                    yield "<think>🛑 Stopping due to failure.</think>\n"
                    break

        result.aggregate_outputs()

    # Stream: Summary
    if result.all_succeeded:
        yield "<think>🎉 All experts completed successfully.</think>\n"
    else:
        failed_names = ", ".join(_get_display_name(e) for e in result.failed_experts)
        yield f"<think>⚠️ Some experts failed: {failed_names}</think>\n"

    # Yield combined output
    if result.combined_output:
        yield f"\n{result.combined_output}"

    logger.info(
        "multi_expert_orchestration_complete",
        task_count=len(tasks),
        mode=mode.value,
        all_succeeded=result.all_succeeded,
    )


def _get_display_name(expert_name: str) -> str:
    """Get user-friendly display name for an expert."""
    display_names = {
        "librarian": "The Librarian",
        "memory": "Memory",
        "home_automation": "Home Automation",
        "tatlock_core": "Core Tools",
    }
    return display_names.get(expert_name, expert_name.title())