"""Sub-agent spawn module for Oclawma.

This module provides thin sub-agent spawning capabilities for parallel task execution.
Each sub-agent receives only task-specific context, enabling efficient parallel processing
with resource limits and result aggregation.
"""

from __future__ import annotations

import asyncio
import contextlib
import uuid
from collections.abc import Callable
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from typing import Any

from oclawma.providers.base import (
    BaseProvider,
    CompletionRequest,
    Message,
)
from oclawma.tools import ToolRegistry, create_default_tools


class SubAgentStatus(Enum):
    """Status of a sub-agent."""

    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELLED = "cancelled"


@dataclass
class SubAgentResult:
    """Result from a sub-agent execution."""

    agent_id: str
    task: str
    status: SubAgentStatus
    output: str = ""
    error: str | None = None
    tokens_used: int = 0
    start_time: datetime | None = None
    end_time: datetime | None = None
    metadata: dict[str, Any] = field(default_factory=dict)

    @property
    def duration_seconds(self) -> float:
        """Calculate execution duration in seconds."""
        if self.start_time and self.end_time:
            return (self.end_time - self.start_time).total_seconds()
        return 0.0

    @property
    def success(self) -> bool:
        """Check if the sub-agent completed successfully."""
        return self.status == SubAgentStatus.COMPLETED and not self.error

    def to_dict(self) -> dict[str, Any]:
        """Convert result to dictionary."""
        return {
            "agent_id": self.agent_id,
            "task": self.task,
            "status": self.status.value,
            "output": self.output,
            "error": self.error,
            "tokens_used": self.tokens_used,
            "duration_seconds": self.duration_seconds,
            "metadata": self.metadata,
        }


@dataclass
class SubAgentConfig:
    """Configuration for a sub-agent."""

    model: str | None = None
    temperature: float = 0.7
    max_tokens: int = 2000
    max_iterations: int = 5
    timeout_seconds: float = 300.0
    system_prompt: str | None = None
    tools_enabled: bool = True
    context_budget: int = 4000  # Token budget for this sub-agent


class SubAgent:
    """A lightweight sub-agent for executing a specific task.

    Each sub-agent receives only task-specific context, making it
    efficient for parallel execution with minimal memory overhead.
    """

    def __init__(
        self,
        provider: BaseProvider,
        task: str,
        config: SubAgentConfig | None = None,
        agent_id: str | None = None,
        parent_context: dict[str, Any] | None = None,
    ):
        """Initialize a sub-agent.

        Args:
            provider: The LLM provider to use
            task: The specific task to execute
            config: Sub-agent configuration
            agent_id: Optional unique ID (generated if not provided)
            parent_context: Optional minimal context from parent
        """
        self.agent_id = agent_id or str(uuid.uuid4())[:8]
        self.provider = provider
        self.task = task
        self.config = config or SubAgentConfig()
        self.parent_context = parent_context or {}

        self.status = SubAgentStatus.PENDING
        self.result: SubAgentResult | None = None
        self._task: asyncio.Task[SubAgentResult] | None = None

        # Initialize tools if enabled
        if self.config.tools_enabled:
            self.tool_registry = create_default_tools()
        else:
            self.tool_registry = ToolRegistry()

    def _create_system_prompt(self) -> str:
        """Create a minimal system prompt for this sub-agent.

        Sub-agents get only task-specific context, not full parent context.
        """
        if self.config.system_prompt:
            return self.config.system_prompt

        tool_descriptions = []
        if self.config.tools_enabled:
            for schema in self.tool_registry.get_all_schemas():
                params = ", ".join(p["name"] for p in schema.get("parameters", []))
                tool_descriptions.append(f"- {schema['name']}({params}): {schema['description']}")

        tools_text = "\n".join(tool_descriptions) if tool_descriptions else "No tools available."

        prompt = f"""You are a specialized sub-agent working on a specific task.

Your Task: {self.task}

Guidelines:
1. Focus ONLY on the assigned task
2. Be concise and direct
3. Return your result in a clear, structured format
4. Use tools when necessary to complete the task

Available tools:
{tools_text}

When using tools, respond with:
TOOL_CALL: tool_name(param1=value1, param2=value2)

Your response will be aggregated with other sub-agents working in parallel."""

        # Add any minimal parent context if provided
        if self.parent_context.get("context_hint"):
            prompt += f"\n\nContext: {self.parent_context['context_hint']}"

        return prompt

    async def run(self) -> SubAgentResult:
        """Execute the sub-agent task.

        Returns:
            The result of the task execution
        """
        self.status = SubAgentStatus.RUNNING
        start_time = datetime.now()

        try:
            # Create minimal message context
            messages = [
                Message(role="system", content=self._create_system_prompt()),
                Message(role="user", content=f"Execute this task: {self.task}"),
            ]

            request = CompletionRequest(
                messages=messages,
                model=self.config.model or "qwen2.5:3b",
                temperature=self.config.temperature,
                max_tokens=self.config.max_tokens,
            )

            # Execute with timeout
            response = await asyncio.wait_for(
                self.provider.complete(request),
                timeout=self.config.timeout_seconds,
            )

            self.status = SubAgentStatus.COMPLETED

            self.result = SubAgentResult(
                agent_id=self.agent_id,
                task=self.task,
                status=self.status,
                output=response.content,
                tokens_used=response.usage.total_tokens if response.usage else 0,
                start_time=start_time,
                end_time=datetime.now(),
            )

        except asyncio.TimeoutError:
            self.status = SubAgentStatus.FAILED
            self.result = SubAgentResult(
                agent_id=self.agent_id,
                task=self.task,
                status=self.status,
                error=f"Task timed out after {self.config.timeout_seconds}s",
                start_time=start_time,
                end_time=datetime.now(),
            )

        except Exception as e:
            self.status = SubAgentStatus.FAILED
            self.result = SubAgentResult(
                agent_id=self.agent_id,
                task=self.task,
                status=self.status,
                error=str(e),
                start_time=start_time,
                end_time=datetime.now(),
            )

        return self.result

    async def start(self) -> asyncio.Task[SubAgentResult]:
        """Start the sub-agent asynchronously.

        Returns:
            The async task for this sub-agent
        """
        self._task = asyncio.create_task(self.run())
        return self._task

    async def cancel(self) -> None:
        """Cancel the sub-agent if running."""
        if self._task and not self._task.done():
            self._task.cancel()
            with contextlib.suppress(asyncio.CancelledError):
                await self._task
            self.status = SubAgentStatus.CANCELLED


@dataclass
class AggregationConfig:
    """Configuration for result aggregation."""

    strategy: str = "concatenate"  # concatenate, summarize, vote, merge
    separator: str = "\n\n---\n\n"
    include_metadata: bool = False
    max_output_length: int = 10000


class SubAgentPool:
    """Pool for managing concurrent sub-agents with resource limits.

    Provides:
    - Resource limiting (max concurrent agents)
    - Async/concurrent execution
    - Result aggregation
    - Progress tracking
    """

    def __init__(
        self,
        provider: BaseProvider,
        max_concurrent: int = 4,
        aggregation_config: AggregationConfig | None = None,
    ):
        """Initialize the sub-agent pool.

        Args:
            provider: The LLM provider for sub-agents
            max_concurrent: Maximum number of concurrent sub-agents
            aggregation_config: Configuration for result aggregation
        """
        self.provider = provider
        self.max_concurrent = max_concurrent
        self.aggregation_config = aggregation_config or AggregationConfig()

        self._semaphore = asyncio.Semaphore(max_concurrent)
        self._agents: list[SubAgent] = []
        self._results: list[SubAgentResult] = []

    async def spawn(
        self,
        task: str,
        config: SubAgentConfig | None = None,
        parent_context: dict[str, Any] | None = None,
    ) -> SubAgent:
        """Create and track a new sub-agent (without starting it).

        Args:
            task: The task for the sub-agent
            config: Optional configuration
            parent_context: Optional minimal context from parent

        Returns:
            The created sub-agent
        """
        agent = SubAgent(
            provider=self.provider,
            task=task,
            config=config,
            parent_context=parent_context,
        )
        self._agents.append(agent)
        return agent

    async def spawn_many(
        self,
        tasks: list[str],
        config: SubAgentConfig | None = None,
        parent_context: dict[str, Any] | None = None,
    ) -> list[SubAgent]:
        """Create multiple sub-agents for parallel execution.

        Args:
            tasks: List of tasks to execute
            config: Optional configuration (applied to all)
            parent_context: Optional minimal context from parent

        Returns:
            List of created sub-agents
        """
        agents = []
        for task in tasks:
            agent = await self.spawn(task, config, parent_context)
            agents.append(agent)
        return agents

    async def execute_concurrent(
        self,
        agents: list[SubAgent] | None = None,
        progress_callback: Callable[[int, int], None] | None = None,
    ) -> list[SubAgentResult]:
        """Execute sub-agents with resource limiting.

        Args:
            agents: List of agents to execute (uses tracked agents if None)
            progress_callback: Optional callback(completed, total) for progress

        Returns:
            List of results from all agents
        """
        agents_to_run = agents or self._agents
        results: list[SubAgentResult] = []
        completed = 0
        total = len(agents_to_run)

        async def run_with_semaphore(agent: SubAgent) -> SubAgentResult:
            """Run agent with semaphore-controlled concurrency."""
            async with self._semaphore:
                result = await agent.run()
                nonlocal completed
                completed += 1
                if progress_callback:
                    progress_callback(completed, total)
                return result

        # Execute all agents concurrently with semaphore limiting
        tasks = [run_with_semaphore(agent) for agent in agents_to_run]
        results = await asyncio.gather(*tasks, return_exceptions=True)

        # Process results
        processed_results: list[SubAgentResult] = []
        for i, result in enumerate(results):
            if isinstance(result, Exception):
                # Convert exception to failed result
                processed_results.append(
                    SubAgentResult(
                        agent_id=agents_to_run[i].agent_id,
                        task=agents_to_run[i].task,
                        status=SubAgentStatus.FAILED,
                        error=str(result),
                        start_time=datetime.now(),
                        end_time=datetime.now(),
                    )
                )
            else:
                processed_results.append(result)

        self._results = processed_results
        return processed_results

    def aggregate_results(self, results: list[SubAgentResult] | None = None) -> str:
        """Aggregate results from multiple sub-agents.

        Args:
            results: List of results to aggregate (uses stored results if None)

        Returns:
            Aggregated output string
        """
        results_to_aggregate = results or self._results

        if not results_to_aggregate:
            return "No results to aggregate."

        strategy = self.aggregation_config.strategy

        if strategy == "concatenate":
            return self._concatenate_results(results_to_aggregate)
        elif strategy == "summarize":
            return self._summarize_results(results_to_aggregate)
        elif strategy == "vote":
            return self._vote_results(results_to_aggregate)
        elif strategy == "merge":
            return self._merge_results(results_to_aggregate)
        else:
            return self._concatenate_results(results_to_aggregate)

    def _concatenate_results(self, results: list[SubAgentResult]) -> str:
        """Concatenate results with separators."""
        parts = []
        for i, result in enumerate(results, 1):
            header = f"### Result {i} (Agent: {result.agent_id})"
            if self.aggregation_config.include_metadata:
                header += f"\nStatus: {result.status.value}"
                header += f"\nTokens: {result.tokens_used}"
                header += f"\nDuration: {result.duration_seconds:.1f}s"

            content = result.output if result.success else f"ERROR: {result.error}"
            parts.append(f"{header}\n\n{content}")

        output = self.aggregation_config.separator.join(parts)

        # Truncate if too long
        if len(output) > self.aggregation_config.max_output_length:
            output = output[: self.aggregation_config.max_output_length]
            output += f"\n\n[...truncated, total results: {len(results)}]"

        return output

    def _summarize_results(self, results: list[SubAgentResult]) -> str:
        """Summarize results (placeholder for LLM-based summarization)."""
        successful = [r for r in results if r.success]
        failed = [r for r in results if not r.success]

        summary = f"""Aggregation Summary:
- Total agents: {len(results)}
- Successful: {len(successful)}
- Failed: {len(failed)}
- Total tokens used: {sum(r.tokens_used for r in results)}

Results:
"""
        for result in successful:
            # Extract first line or first 100 chars as summary
            preview = result.output.split("\n")[0][:100]
            summary += f"\n• Agent {result.agent_id}: {preview}..."

        for result in failed:
            summary += f"\n• Agent {result.agent_id}: FAILED - {result.error}"

        return summary

    def _vote_results(self, results: list[SubAgentResult]) -> str:
        """Vote-based aggregation for categorical results."""
        successful = [r for r in results if r.success]

        if not successful:
            return "No successful results to vote on."

        # Simple voting: count occurrences of each output
        from collections import Counter

        outputs = [r.output.strip() for r in successful]
        vote_counts = Counter(outputs)

        # Get the most common result
        most_common = vote_counts.most_common(1)[0]

        output = f"""Voting Results:
- Total votes: {len(successful)}
- Winning result: "{most_common[0]}"
- Confidence: {most_common[1]}/{len(successful)} ({most_common[1]/len(successful)*100:.1f}%)

All votes:
"""
        for output_text, count in vote_counts.most_common():
            output += f"\n  {count}x: {output_text[:50]}..."

        return output

    def _merge_results(self, results: list[SubAgentResult]) -> str:
        """Merge results by combining unique outputs."""
        successful = [r for r in results if r.success]

        if not successful:
            return "No successful results to merge."

        # Deduplicate outputs
        seen = set()
        unique_outputs = []
        for result in successful:
            # Use a simple hash for deduplication
            key = result.output.strip().lower()[:100]
            if key not in seen:
                seen.add(key)
                unique_outputs.append(result.output)

        output = f"Merged Results ({len(unique_outputs)} unique from {len(successful)} agents):\n\n"
        output += self.aggregation_config.separator.join(unique_outputs)

        return output

    def get_stats(self) -> dict[str, Any]:
        """Get statistics about the pool execution."""
        if not self._results:
            return {"status": "no_results"}

        successful = sum(1 for r in self._results if r.success)
        failed = len(self._results) - successful

        return {
            "total_agents": len(self._results),
            "successful": successful,
            "failed": failed,
            "success_rate": successful / len(self._results) if self._results else 0,
            "total_tokens": sum(r.tokens_used for r in self._results),
            "total_duration": sum(r.duration_seconds for r in self._results),
            "avg_duration": (
                sum(r.duration_seconds for r in self._results) / len(self._results)
                if self._results
                else 0
            ),
        }

    async def cancel_all(self) -> None:
        """Cancel all running sub-agents."""
        for agent in self._agents:
            await agent.cancel()


class SubAgentManager:
    """High-level manager for sub-agent operations.

    Provides a simple interface for common sub-agent patterns.
    """

    def __init__(
        self,
        provider: BaseProvider,
        default_max_concurrent: int = 4,
    ):
        """Initialize the sub-agent manager.

        Args:
            provider: The LLM provider for sub-agents
            default_max_concurrent: Default max concurrent agents
        """
        self.provider = provider
        self.default_max_concurrent = default_max_concurrent

    async def parallel_map(
        self,
        tasks: list[str],
        config: SubAgentConfig | None = None,
        max_concurrent: int | None = None,
        parent_context: dict[str, Any] | None = None,
        aggregation_strategy: str = "concatenate",
        progress_callback: Callable[[int, int], None] | None = None,
    ) -> tuple[list[SubAgentResult], str]:
        """Map tasks to sub-agents in parallel and aggregate results.

        This is the main entry point for parallel task execution.

        Args:
            tasks: List of tasks to execute
            config: Configuration for sub-agents
            max_concurrent: Override default max concurrent
            parent_context: Optional minimal context from parent
            aggregation_strategy: How to aggregate results
            progress_callback: Optional progress callback

        Returns:
            Tuple of (individual results, aggregated output)
        """
        pool = SubAgentPool(
            provider=self.provider,
            max_concurrent=max_concurrent or self.default_max_concurrent,
            aggregation_config=AggregationConfig(strategy=aggregation_strategy),
        )

        # Spawn agents for all tasks
        await pool.spawn_many(tasks, config, parent_context)

        # Execute concurrently
        results = await pool.execute_concurrent(progress_callback=progress_callback)

        # Aggregate results
        aggregated = pool.aggregate_results(results)

        return results, aggregated

    async def parallel_vote(
        self,
        task: str,
        num_agents: int = 3,
        config: SubAgentConfig | None = None,
        max_concurrent: int | None = None,
    ) -> tuple[list[SubAgentResult], str]:
        """Run the same task with multiple agents and vote on results.

        Useful for increasing confidence in answers through redundancy.

        Args:
            task: The task to execute
            num_agents: Number of agents to run
            config: Configuration for sub-agents
            max_concurrent: Override default max concurrent

        Returns:
            Tuple of (individual results, voted output)
        """
        tasks = [task] * num_agents
        return await self.parallel_map(
            tasks=tasks,
            config=config,
            max_concurrent=max_concurrent,
            aggregation_strategy="vote",
        )

    async def divide_and_conquer(
        self,
        subtasks: list[str],
        merge_prompt: str | None = None,
        config: SubAgentConfig | None = None,
        max_concurrent: int | None = None,
    ) -> tuple[list[SubAgentResult], str]:
        """Divide work into subtasks and merge results.

        Args:
            subtasks: List of subtasks to execute
            merge_prompt: Optional prompt for merging results
            config: Configuration for sub-agents
            max_concurrent: Override default max concurrent

        Returns:
            Tuple of (individual results, merged output)
        """
        results, aggregated = await self.parallel_map(
            tasks=subtasks,
            config=config,
            max_concurrent=max_concurrent,
            aggregation_strategy="merge",
        )

        # If a merge prompt is provided, do an additional LLM-based merge
        if merge_prompt:
            # This would use the provider to do an intelligent merge
            # For now, we just return the merge-strategy aggregation
            pass

        return results, aggregated
