"""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 dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from typing import Any, Callable

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
