# createsonline/agents/orchestrator.py
"""
Multi-agent orchestration for CREATESONLINE.

Provides Pipeline, ParallelPipeline, and AgentOrchestrator for
coordinating multiple agents on complex tasks.
"""

import time
import asyncio
import json
from typing import Any, Callable, Dict, List, Optional, Union

from .base import Agent, SubAgent, AgentResponse


class Pipeline:
    """
    Sequential pipeline — passes output from one agent to the next.
    
    Usage::
    
        pipe = Pipeline([
            researcher,   # First: research the topic
            summarizer,   # Then: summarize findings
            writer,       # Finally: write the article
        ])
        result = await pipe.run("Write an article about quantum computing")
    """
    
    def __init__(self, agents: List[Agent], name: str = 'pipeline',
                 transform_fn: Callable = None):
        self.agents = agents
        self.name = name
        # Optional transform between steps: fn(prev_output, next_agent) -> str
        self.transform_fn = transform_fn
        self._history: List[Dict[str, Any]] = []
    
    async def run(self, initial_input: str, context: dict = None) -> AgentResponse:
        """Run the pipeline sequentially."""
        t0 = time.time()
        current_input = initial_input
        all_responses: List[AgentResponse] = []
        
        for i, agent in enumerate(self.agents):
            step = {
                'step': i + 1,
                'agent': agent.name,
                'input': current_input[:200],
                'start_time': time.time(),
            }
            
            # Transform input for next agent if needed
            if i > 0 and self.transform_fn:
                current_input = self.transform_fn(current_input, agent)
            
            # Run agent
            response = await agent.run(current_input, context)
            all_responses.append(response)
            
            step['output'] = response.content[:200]
            step['success'] = response.success
            step['duration'] = response.duration
            self._history.append(step)
            
            if not response.success:
                # Pipeline failed at this step
                final = AgentResponse(
                    content=f"Pipeline failed at step {i + 1} ({agent.name}): {response.error}",
                    success=False,
                    error=response.error,
                    agent_name=self.name,
                    duration=time.time() - t0,
                )
                final.metadata['steps'] = self._history[-len(self.agents):]
                return final
            
            # Next agent receives this agent's output
            current_input = response.content
        
        # Final result is the last agent's output
        final = AgentResponse(
            content=current_input,
            success=True,
            agent_name=self.name,
            duration=time.time() - t0,
        )
        final.metadata['steps'] = self._history[-len(self.agents):]
        final.tool_calls = [tc for r in all_responses for tc in r.tool_calls]
        final.thoughts = [t for r in all_responses for t in r.thoughts]
        return final
    
    @property
    def history(self) -> List[dict]:
        return list(self._history)
    
    def __repr__(self):
        agents = ' → '.join(a.name for a in self.agents)
        return f"Pipeline({agents})"


class ParallelPipeline:
    """
    Parallel pipeline — runs multiple agents concurrently and merges results.
    
    Usage::
    
        parallel = ParallelPipeline(
            agents=[web_researcher, doc_researcher, code_researcher],
            merge_fn=lambda results: "\\n---\\n".join(r.content for r in results),
        )
        result = await parallel.run("Find info about Python type hints")
    """
    
    def __init__(self, agents: List[Agent], name: str = 'parallel',
                 merge_fn: Callable = None, timeout: float = 60.0):
        self.agents = agents
        self.name = name
        self.timeout = timeout
        self.merge_fn = merge_fn or self._default_merge
    
    @staticmethod
    def _default_merge(responses: List[AgentResponse]) -> str:
        """Default merge: concatenate all successful responses."""
        parts = []
        for r in responses:
            if r.success:
                parts.append(f"[{r.agent_name}]:\n{r.content}")
        return '\n\n---\n\n'.join(parts)
    
    async def run(self, user_input: str, context: dict = None) -> AgentResponse:
        """Run all agents in parallel."""
        t0 = time.time()
        
        # Create tasks
        tasks = [agent.run(user_input, context) for agent in self.agents]
        
        try:
            responses = await asyncio.wait_for(
                asyncio.gather(*tasks, return_exceptions=True),
                timeout=self.timeout,
            )
        except asyncio.TimeoutError:
            return AgentResponse(
                content="Parallel pipeline timed out",
                success=False,
                error="Timeout",
                agent_name=self.name,
                duration=time.time() - t0,
            )
        
        # Filter valid responses
        valid_responses = []
        for r in responses:
            if isinstance(r, AgentResponse):
                valid_responses.append(r)
            elif isinstance(r, Exception):
                valid_responses.append(AgentResponse(
                    content='',
                    success=False,
                    error=str(r),
                    agent_name='unknown',
                ))
        
        # Merge
        merged_content = self.merge_fn(valid_responses)
        
        final = AgentResponse(
            content=merged_content,
            success=any(r.success for r in valid_responses),
            agent_name=self.name,
            duration=time.time() - t0,
        )
        final.metadata['individual_results'] = [
            {'agent': r.agent_name, 'success': r.success, 'length': len(r.content)}
            for r in valid_responses
        ]
        final.tool_calls = [tc for r in valid_responses for tc in r.tool_calls]
        return final
    
    def __repr__(self):
        agents = ', '.join(a.name for a in self.agents)
        return f"ParallelPipeline([{agents}])"


class AgentOrchestrator:
    """
    High-level orchestrator for managing a team of agents.
    
    The orchestrator dispatches tasks to the most appropriate agent
    based on routing rules, and handles complex multi-step workflows.
    
    Usage::
    
        orch = AgentOrchestrator(name="team_lead")
        
        orch.register_agent(researcher, skills=["research", "search"])
        orch.register_agent(coder, skills=["code", "debug", "python"])
        orch.register_agent(writer, skills=["write", "summarize", "explain"])
        
        # Auto-route based on skills
        result = await orch.run("Write a Python function to sort a list")
        # → Routes to 'coder' agent
        
        # Multi-step workflow
        result = await orch.workflow([
            ("research", "Find best sorting algorithms"),
            ("code", "Implement quicksort in Python"),
            ("write", "Document the implementation"),
        ])
    """
    
    def __init__(self, name: str = 'orchestrator', router: Callable = None,
                 verbose: bool = False):
        self.name = name
        self.verbose = verbose
        self._agents: Dict[str, Agent] = {}
        self._skills: Dict[str, List[str]] = {}      # agent_name -> [skills]
        self._skill_index: Dict[str, str] = {}        # skill -> agent_name
        self._router = router                          # custom routing function
        self._task_log: List[Dict[str, Any]] = []
        self._default_agent: Optional[str] = None
    
    # ------------------------------------------------------------------
    # Registration
    # ------------------------------------------------------------------
    def register_agent(self, agent: Agent, skills: List[str] = None,
                       is_default: bool = False):
        """Register an agent with optional skill tags."""
        self._agents[agent.name] = agent
        skills = skills or []
        self._skills[agent.name] = skills
        for skill in skills:
            self._skill_index[skill.lower()] = agent.name
        if is_default:
            self._default_agent = agent.name
    
    def unregister_agent(self, name: str):
        self._agents.pop(name, None)
        skills = self._skills.pop(name, [])
        for skill in skills:
            self._skill_index.pop(skill.lower(), None)
    
    # ------------------------------------------------------------------
    # Routing
    # ------------------------------------------------------------------
    def _auto_route(self, query: str) -> Optional[str]:
        """Route a query to the best agent based on skill keyword matching."""
        if self._router:
            agent_name = self._router(query, list(self._agents.keys()))
            if agent_name:
                return agent_name
        
        query_lower = query.lower()
        best_agent = None
        best_score = 0
        
        for agent_name, skills in self._skills.items():
            score = 0
            for skill in skills:
                if skill.lower() in query_lower:
                    score += 1
                # Also check words
                for word in skill.lower().split():
                    if word in query_lower and len(word) > 2:
                        score += 0.5
            if score > best_score:
                best_score = score
                best_agent = agent_name
        
        if best_agent:
            return best_agent
        
        return self._default_agent or (
            list(self._agents.keys())[0] if self._agents else None
        )
    
    # ------------------------------------------------------------------
    # Execution
    # ------------------------------------------------------------------
    async def run(self, query: str, agent_name: str = None,
                  context: dict = None) -> AgentResponse:
        """
        Run a query with automatic or explicit agent routing.
        
        Args:
            query: The user's request.
            agent_name: Explicitly specify which agent to use.
            context: Additional context.
        
        Returns:
            AgentResponse from the selected agent.
        """
        if not self._agents:
            return AgentResponse(
                content="No agents registered",
                success=False,
                error="No agents available",
                agent_name=self.name,
            )
        
        # Route
        target = agent_name or self._auto_route(query)
        agent = self._agents.get(target)
        
        if agent is None:
            return AgentResponse(
                content=f"Agent '{target}' not found",
                success=False,
                error=f"Unknown agent: {target}",
                agent_name=self.name,
            )
        
        if self.verbose:
            print(f"[{self.name}] Routing to '{target}': {query[:80]}")
        
        # Execute
        response = await agent.run(query, context)
        
        # Log
        self._task_log.append({
            'query': query[:200],
            'routed_to': target,
            'success': response.success,
            'duration': response.duration,
            'timestamp': time.time(),
        })
        
        return response
    
    async def workflow(self, steps: List[tuple],
                       context: dict = None) -> AgentResponse:
        """
        Execute a multi-step workflow.
        
        Args:
            steps: List of (agent_name_or_skill, task) tuples.
            context: Shared context dict.
        
        Returns:
            Final AgentResponse with accumulated results.
        """
        t0 = time.time()
        results = []
        current_context = context or {}
        
        for i, step in enumerate(steps):
            if len(step) == 2:
                target, task = step
            elif len(step) == 3:
                target, task, step_ctx = step
                current_context.update(step_ctx)
            else:
                continue
            
            # Resolve target — could be agent name or skill
            agent_name = target
            if target in self._skill_index:
                agent_name = self._skill_index[target]
            
            # Inject previous results into context
            if results:
                current_context['previous_step'] = results[-1].content
                task_with_context = (
                    f"{task}\n\n[Context from previous step]:\n"
                    f"{results[-1].content[:1000]}"
                )
            else:
                task_with_context = task
            
            if self.verbose:
                print(f"[{self.name}] Workflow step {i + 1}: {agent_name} → {task[:60]}")
            
            response = await self.run(
                task_with_context,
                agent_name=agent_name,
                context=current_context,
            )
            results.append(response)
            
            if not response.success:
                final = AgentResponse(
                    content=f"Workflow failed at step {i + 1}: {response.error}",
                    success=False,
                    error=response.error,
                    agent_name=self.name,
                    duration=time.time() - t0,
                )
                final.metadata['completed_steps'] = i
                return final
        
        # Merge all results
        final = AgentResponse(
            content=results[-1].content if results else '',
            success=True,
            agent_name=self.name,
            duration=time.time() - t0,
        )
        final.metadata['workflow_steps'] = [
            {'agent': r.agent_name, 'success': r.success, 'duration': r.duration}
            for r in results
        ]
        final.thoughts = [t for r in results for t in r.thoughts]
        final.tool_calls = [tc for r in results for tc in r.tool_calls]
        return final
    
    async def broadcast(self, query: str, context: dict = None) -> List[AgentResponse]:
        """Send a query to ALL registered agents in parallel."""
        tasks = [
            agent.run(query, context) for agent in self._agents.values()
        ]
        return await asyncio.gather(*tasks, return_exceptions=False)
    
    # ------------------------------------------------------------------
    # Info
    # ------------------------------------------------------------------
    @property
    def agents(self) -> Dict[str, Agent]:
        return dict(self._agents)
    
    @property
    def task_log(self) -> List[dict]:
        return list(self._task_log)
    
    def get_stats(self) -> dict:
        return {
            'name': self.name,
            'agents': {
                name: {
                    'skills': self._skills.get(name, []),
                    'stats': agent.stats,
                }
                for name, agent in self._agents.items()
            },
            'total_tasks': len(self._task_log),
            'success_rate': (
                sum(1 for t in self._task_log if t['success']) / len(self._task_log)
                if self._task_log else 0.0
            ),
        }
    
    def __repr__(self):
        return f"AgentOrchestrator({self.name!r}, agents={len(self._agents)})"
