import asyncio
from typing import List, Dict, Optional, Union
from datetime import datetime
from .agent import Agent
from .pool import AgentPool


class TraceStep:
    def __init__(self, agent: str, role: str, content: str):
        self.agent = agent
        self.role = role
        self.content = content
        self.timestamp = datetime.utcnow().isoformat()

    def to_dict(self):
        return {"agent": self.agent, "role": self.role, "content": self.content, "timestamp": self.timestamp}


class BridgeResult:
    """Feature 9: Full conversation trace."""
    def __init__(self, final: str, trace: List[TraceStep], agents_used: List[str], tokens_used: int = 0):
        self.final = final
        self.trace = trace
        self.agents_used = agents_used
        self.tokens_used = tokens_used
        self.timestamp = datetime.utcnow().isoformat()

    def summary(self) -> str:
        lines = ["=== Synapse Bridge Result ===", f"Agents: {' → '.join(self.agents_used)}", ""]
        for step in self.trace:
            lines.append(f"[{step.agent}] {step.role.upper()}")
            lines.append(f"  {step.content[:300]}{'...' if len(step.content) > 300 else ''}")
            lines.append("")
        lines.append(f"FINAL:\n{self.final}")
        return "\n".join(lines)

    def to_dict(self):
        return {
            "final": self.final,
            "trace": [s.to_dict() for s in self.trace],
            "agents_used": self.agents_used,
            "tokens_used": self.tokens_used,
            "timestamp": self.timestamp,
        }


class Bridge:
    """
    Feature 2:  Local agent-to-agent communication
    Feature 8:  Shared memory across agents
    Feature 9:  Full conversation trace
    Feature 16: Multi-agent orchestration
    Feature 20: Agent chaining
    """

    def __init__(
        self,
        agents: List[Union[Agent, AgentPool]],
        memory: bool = True,
        verbose: bool = False,
    ):
        if not agents:
            raise ValueError("Bridge requires at least one agent.")
        self.agents = agents
        self.use_memory = memory
        self.verbose = verbose
        self._history: List[Dict] = []

    def run(self, task: str) -> BridgeResult:
        """
        Feature 20: Sequential pipeline — each agent builds on previous.
        """
        self._history = []
        trace = []
        agents_used = []

        self._history.append({"role": "user", "content": task})
        trace.append(TraceStep("user", "input", task))

        if self.verbose:
            print(f"\n🧠 Synapse Bridge")
            print(f"Task: {task[:100]}")
            print(f"Pipeline: {' → '.join(self._agent_name(a) for a in self.agents)}\n")

        current = task
        for agent in self.agents:
            name = self._agent_name(agent)
            if self.verbose:
                print(f"🤖 [{name}] thinking...")

            messages = self._history.copy() if self.use_memory else [{"role": "user", "content": current}]

            if isinstance(agent, AgentPool):
                response = agent.complete(messages)
            else:
                response = agent.complete(messages)

            self._history.append({"role": "assistant", "content": response})
            self._history.append({"role": "user", "content": f"Continue based on {name}'s response above."})

            trace.append(TraceStep(name, "output", response))
            agents_used.append(name)
            current = response

            if self.verbose:
                print(f"✅ [{name}] done.\n")

        return BridgeResult(final=current, trace=trace, agents_used=agents_used)

    async def run_async(self, task: str) -> BridgeResult:
        """Async version of run."""
        loop = asyncio.get_event_loop()
        return await loop.run_in_executor(None, self.run, task)

    async def run_parallel(self, task: str) -> List[BridgeResult]:
        """Run all agents in parallel and return all results."""
        async def run_single(agent):
            name = self._agent_name(agent)
            messages = [{"role": "user", "content": task}]
            if isinstance(agent, AgentPool):
                response = agent.complete(messages)
            else:
                response = agent.complete(messages)
            return BridgeResult(
                final=response,
                trace=[TraceStep("user", "input", task), TraceStep(name, "output", response)],
                agents_used=[name]
            )

        tasks = [run_single(a) for a in self.agents]
        return await asyncio.gather(*tasks)

    def send(self, from_agent: str, to_agent: str, message: str) -> str:
        """Feature 2: Direct agent-to-agent message."""
        target = self._find_agent(to_agent)
        if not target:
            raise ValueError(f"Agent '{to_agent}' not found. Available: {[self._agent_name(a) for a in self.agents]}")

        if self.verbose:
            print(f"📨 [{from_agent}] → [{to_agent}]")

        messages = [{"role": "user", "content": f"Message from agent '{from_agent}':\n\n{message}"}]

        if isinstance(target, AgentPool):
            return target.complete(messages)
        return target.complete(messages)

    def _find_agent(self, name: str):
        for agent in self.agents:
            if self._agent_name(agent) == name:
                return agent
        return None

    def _agent_name(self, agent) -> str:
        if isinstance(agent, AgentPool):
            return f"pool({agent.agents[0].name})"
        return agent.name

    def add_agent(self, agent):
        self.agents.append(agent)

    def remove_agent(self, name: str):
        self.agents = [a for a in self.agents if self._agent_name(a) != name]

    def __repr__(self):
        return f"Bridge(agents={[self._agent_name(a) for a in self.agents]}, memory={self.use_memory})"
