"""
Nexus High-Level API — build multi-agent systems in 5 lines of code.

This module provides a simple, batteries-included interface that wires up
all Nexus subsystems (engine, orchestration, models, observability) behind
a single `Team` object.

Quick Start:

    from nexus.api import Team, Agent

    team = Team(model="gpt-4o")

    researcher = team.add_agent("researcher", instructions="You do research.")
    writer     = team.add_agent("writer",     instructions="You write articles.")

    result = await team.run("Write a blog post about AI agents")
    print(result)
"""

from __future__ import annotations

import asyncio
import logging
from collections.abc import AsyncIterator, Callable
from typing import Any

from nexus.core.agent import Agent, AgentHooks, AgentTool
from nexus.core.handoff import HandoffMode
from nexus.core.message import Message
from nexus.core.task import Task, TaskResult
from nexus.models.base import ModelProvider
from nexus.observability.metrics import MetricsCollector
from nexus.observability.tracing import Tracer
from nexus.orchestration.topology import TopologyEngine, TopologyMode
from nexus.protocol.bus import MessageBus
from nexus.protocol.discovery import DiscoveryService
from nexus.runtime.budget import BudgetConfig, TokenBudgetManager
from nexus.runtime.checkpoint import CheckpointStore, FileCheckpointStore, InMemoryCheckpointStore
from nexus.runtime.engine import ExecutionEngine
from nexus.runtime.policies import (
    PolicyChain,
    RetryPolicy,
)
from nexus.types import AgentCapability, TaskConstraints

logger = logging.getLogger("nexus.api")


def _create_provider(model: str, **kwargs: Any) -> ModelProvider:
    """
    Auto-detect and create a model provider from a model name string.

    Naming conventions:
    - "gpt-*", "o1-*", "o3-*"           → OpenAIProvider
    - "claude-*"                         → AnthropicProvider
    - "groq/...", "bedrock/...", etc.    → LiteLLMProvider (100+ providers)
    - explicit provider="litellm"        → LiteLLMProvider
    - everything else                    → OllamaProvider (local)
    """
    model_lower = model.lower()

    # Explicit provider override
    explicit_provider = kwargs.pop("provider", None)
    if explicit_provider == "litellm":
        from nexus.models.litellm import LiteLLMProvider
        return LiteLLMProvider(model=model, **kwargs)

    # OpenAI models
    if any(model_lower.startswith(prefix) for prefix in ("gpt-", "o1", "o3", "chatgpt")):
        from nexus.models.openai import OpenAIProvider
        return OpenAIProvider(model=model, **kwargs)

    # Anthropic models
    if model_lower.startswith("claude"):
        from nexus.models.anthropic import AnthropicProvider
        return AnthropicProvider(model=model, **kwargs)

    # LiteLLM auto-detection: if model contains "/" and prefix is a known provider
    if "/" in model:
        prefix = model_lower.split("/", 1)[0]
        from nexus.models.litellm import LITELLM_PROVIDER_PREFIXES
        if prefix in LITELLM_PROVIDER_PREFIXES:
            from nexus.models.litellm import LiteLLMProvider
            return LiteLLMProvider(model=model, **kwargs)

    # Default: treat as Ollama local model
    from nexus.models.ollama import OllamaProvider
    return OllamaProvider(model=model, **kwargs)


class Team:
    """
    High-level multi-agent team — the primary user-facing API.

    Wraps the entire Nexus stack behind a simple interface:
    1. Add agents with `add_agent()`
    2. Define flows with `add_handoff()`, `set_graph()`, or let the router decide
    3. Run with `run()`

    Usage:
        team = Team(model="gpt-4o")
        coder = team.add_agent("coder", instructions="Write Python code")
        reviewer = team.add_agent("reviewer", instructions="Review code")
        team.add_handoff(coder, reviewer)  # coder → reviewer pipeline
        result = await team.run("Build a fibonacci function")
    """

    def __init__(
        self,
        model: str | ModelProvider = "gpt-4o",
        mode: TopologyMode = TopologyMode.ROUTER,
        budget: BudgetConfig | None = None,
        checkpoint_dir: str | None = None,
        enable_tracing: bool = True,
        enable_metrics: bool = True,
        verbose: bool = False,
        embedding_model: str | None = None,
        embedding_api_key: str | None = None,
        embedding_base_url: str | None = None,
        **model_kwargs: Any,
    ) -> None:
        # Model provider
        if isinstance(model, str):
            self._default_provider = _create_provider(model, **model_kwargs)
        else:
            self._default_provider = model
            model_kwargs = {}

        # Core infrastructure
        self.bus = MessageBus()
        self.discovery = DiscoveryService()
        self.budget_manager = TokenBudgetManager(budget)
        self.tracer = Tracer("nexus") if enable_tracing else None
        self.metrics = MetricsCollector() if enable_metrics else None

        # Checkpoint
        if checkpoint_dir:
            self._checkpoint_store: CheckpointStore = FileCheckpointStore(checkpoint_dir)
        else:
            self._checkpoint_store = InMemoryCheckpointStore()

        # Policies
        self._policies = PolicyChain(
            retry=RetryPolicy(max_attempts=2),
        )

        # Engine
        self.engine = ExecutionEngine(
            checkpoint_store=self._checkpoint_store,
            budget_config=budget,
            policy_chain=self._policies,
            message_bus=self.bus,
            discovery=self.discovery,
        )

        # Orchestration
        self.topology = TopologyEngine(mode=mode, discovery=self.discovery)

        # In Router mode, create a triage agent that picks the right specialist
        if mode == TopologyMode.ROUTER:
            self._router_agent = Agent(
                name="_router",
                instructions=(
                    "You are a routing agent. Your ONLY job is to decide which "
                    "specialist agent should handle a user's request.\n"
                    "You will be given the task description and a list of available agents.\n"
                    "Respond with ONLY the agent name. No explanation."
                ),
                model_provider=self._default_provider,
            )
            self.topology.router.router_agent = self._router_agent
        else:
            self._router_agent = None

        # Embedding-based capability inference (Adaptive mode)
        # NOTE: Embedding uses its OWN endpoint — NOT the team's base_url/api_key.
        # By default it relies on env vars (e.g. OPENAI_API_KEY for OpenAI embeddings).
        # Pass embedding_api_key / embedding_base_url only if the embedding
        # provider needs separate auth.
        self._capability_inference = None
        if embedding_model and mode == TopologyMode.ADAPTIVE:
            from nexus.orchestration.capability_inference import CapabilityInferenceService
            self._capability_inference = CapabilityInferenceService(
                embedding_model=embedding_model,
                api_key=embedding_api_key,
                api_base=embedding_base_url,
            )
            logger.info(
                f"Capability inference enabled (embedding_model={embedding_model})"
            )

        # State
        self._agents: dict[str, Agent] = {}
        self._mode = mode
        self._verbose = verbose
        self._conversation_history: list[dict[str, str]] = []  # Multi-turn history

        # Configure logging
        if verbose:
            logging.basicConfig(
                level=logging.DEBUG,
                format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
                datefmt="%H:%M:%S",
            )
            # Suppress noisy third-party loggers
            for noisy in (
                "httpx", "httpcore", "openai", "anthropic",
                "litellm", "LiteLLM", "LiteLLM Router", "LiteLLM Proxy",
                "litellm.litellm", "litellm.utils", "litellm.main",
                "litellm.cost_calculator",
            ):
                logging.getLogger(noisy).setLevel(logging.WARNING)
        else:
            # At least show INFO for nexus namespace
            nexus_logger = logging.getLogger("nexus")
            if not nexus_logger.handlers:
                handler = logging.StreamHandler()
                handler.setFormatter(
                    logging.Formatter("%(asctime)s [%(levelname)s] %(name)s: %(message)s", "%H:%M:%S")
                )
                nexus_logger.addHandler(handler)
                nexus_logger.setLevel(logging.INFO)

    def add_agent(
        self,
        name: str,
        instructions: str | Callable[..., str] = "",
        tools: list[AgentTool | Callable[..., Any]] | None = None,
        capabilities: AgentCapability | None = None,
        model: str | ModelProvider | None = None,
        hooks: AgentHooks | None = None,
        **kwargs: Any,
    ) -> Agent:
        """
        Create and register an agent in the team.

        Args:
            name: Unique human-readable name
            instructions: System prompt or callable that returns one
            tools: List of tools (callables or AgentTool)
            capabilities: Agent capabilities vector for matching
            model: Override model for this agent (default: team's model)
            hooks: Lifecycle hooks

        Returns:
            The created Agent instance
        """
        # Resolve provider
        if model is not None:
            provider = (
                _create_provider(model)
                if isinstance(model, str)
                else model
            )
        else:
            provider = self._default_provider

        agent = Agent(
            name=name,
            instructions=instructions,
            tools=tools or [],
            capabilities=capabilities or AgentCapability(),
            model_provider=provider,
            hooks=hooks,
            **kwargs,
        )

        self._agents[agent.id] = agent
        self.engine.register_agent(agent)
        self.topology.add_agent(agent)

        logger.info(f"Agent '{name}' registered (id={agent.id})")
        return agent

    def add_handoff(
        self,
        from_agent: Agent | str,
        to_agent: Agent | str,
        mode: HandoffMode = HandoffMode.DELEGATE,
    ) -> None:
        """
        Declare that from_agent can hand off work to to_agent.

        In Graph mode, this also creates an edge in the graph.
        """
        from_id = from_agent.id if isinstance(from_agent, Agent) else from_agent
        to_id = to_agent.id if isinstance(to_agent, Agent) else to_agent

        # Register in graph orchestrator if in graph mode
        if self._mode == TopologyMode.GRAPH:
            self.topology.graph.add_edge(from_id, to_id)

    # ---- Main execution entry points ----

    async def run(
        self,
        prompt: str,
        agent: Agent | str | None = None,
        constraints: TaskConstraints | None = None,
        context: dict[str, Any] | None = None,
    ) -> str:
        """
        Run a task through the team and return the result.

        This is the simplest way to use Nexus:
            result = await team.run("Summarize this article...")

        Args:
            prompt: The user's request
            agent: Optionally specify which agent handles it
            constraints: Budget / quality / capability constraints
            context: Additional context data

        Returns:
            The final text response
        """
        task = Task(
            name="user_request",
            description=prompt,
            input_data=context or {},
            constraints=constraints or TaskConstraints(),
        )
        task.conversation.append(Message.user(prompt))

        # In Adaptive mode, use embedding service to infer task capabilities
        if (
            self._mode == TopologyMode.ADAPTIVE
            and task.required_capabilities is None
            and self._capability_inference is not None
        ):
            try:
                inferred = await self._capability_inference.infer(prompt)
                if not inferred.is_zero():
                    task.required_capabilities = inferred
                else:
                    logger.debug("Embedding inference returned zero vector")
            except Exception as e:
                logger.warning(f"Capability inference failed: {e}")

        # Guard: at least one agent must be registered
        if not self._agents:
            raise RuntimeError(
                "No agents registered. Call team.add_agent() before team.run().\n"
                "Example:\n"
                "  team.add_agent('coder', instructions='You are a programmer')\n"
                "  team.add_agent('writer', instructions='You are a writer')"
            )

        # Resolve target agent
        agent_id: str | None = None
        if agent:
            agent_id = agent.id if isinstance(agent, Agent) else agent

        with self._trace_context("team.run", prompt=prompt):
            # Use the appropriate orchestration path based on mode
            if agent_id:
                # Explicit agent — bypass routing
                result = await self.engine.run_task(task, agent_id=agent_id)
            elif self._mode == TopologyMode.ROUTER:
                # Router mode — let RouterOrchestrator pick the best agent
                result = await self.topology.router.route_and_execute(task)
            elif self._mode == TopologyMode.ADAPTIVE:
                # Adaptive mode — capability matching + ELO
                from nexus.core.task import TaskGraph
                graph = TaskGraph()
                graph.add_task(task)
                results = await self.topology.adaptive.run(graph)
                if task.id not in results:
                    raise RuntimeError(f"Adaptive mode produced no result for task")
                result = results[task.id]
            else:
                # Graph mode or fallback — use engine directly
                result = await self.engine.run_task(task, agent_id=agent_id)

            if self.metrics:
                self.metrics.task_completed(
                    task_type=task.task_type,
                    agent_id=task.assigned_agent_id or "",
                    latency=task.elapsed_seconds or 0.0,
                    tokens=task.token_usage.total_tokens,
                    cost=task.token_usage.cost_usd,
                )

            return result.content

    async def run_tasks(
        self,
        tasks: list[Task] | list[str],
        parallel: bool = True,
    ) -> list[TaskResult]:
        """
        Run multiple independent tasks.

        Args:
            tasks: Either Task objects or simple prompt strings
            parallel: Whether to run them concurrently (default: True)
        """
        # Normalize strings to Tasks
        task_objects = []
        for t in tasks:
            if isinstance(t, str):
                task = Task(name="task", description=t)
                task.conversation.append(Message.user(t))
                task_objects.append(task)
            else:
                task_objects.append(t)

        with self._trace_context("team.run_tasks", count=len(task_objects)):
            results = await self.topology.run_tasks(task_objects)
            return results

    async def run_pipeline(
        self,
        prompt: str,
        pipeline: list[Agent | str],
    ) -> str:
        """
        Run a prompt through a sequence of agents (pipeline pattern).

        Each agent's output becomes the next agent's input.

        Usage:
            result = await team.run_pipeline(
                "Write a blog post",
                [researcher, writer, editor]
            )
        """
        current_text = prompt

        with self._trace_context("team.run_pipeline", steps=len(pipeline)):
            for step_agent in pipeline:
                current_text = await self.run(current_text, agent=step_agent)

        return current_text

    # ---- Multi-turn & streaming entry points ----

    async def _resolve_agent_for_task(
        self,
        task: Task,
        agent: Agent | str | None = None,
    ) -> Agent:
        """
        Determine which agent should handle *task*.

        Resolution order:
        1. Explicit agent parameter
        2. Router / Adaptive / Graph orchestration
        3. First registered agent (fallback)
        """
        if not self._agents:
            raise RuntimeError(
                "No agents registered. Call team.add_agent() before running.\n"
                "Example:\n"
                "  team.add_agent('coder', instructions='You are a programmer')"
            )

        # Explicit agent
        if agent:
            agent_id = agent.id if isinstance(agent, Agent) else agent
            resolved = self._agents.get(agent_id)
            if resolved:
                logger.info(f"🎯 Agent resolved (explicit): '{resolved.name}'")
                return resolved
            # Try to match by name
            for a in self._agents.values():
                if a.name == agent_id:
                    logger.info(f"🎯 Agent resolved (explicit by name): '{a.name}'")
                    return a
            raise KeyError(f"Agent '{agent_id}' not found in team")

        # Adaptive mode — use capability matching
        if self._mode == TopologyMode.ADAPTIVE:
            selected_id = self.topology.adaptive._select_best_agent(task, current_load={})
            if selected_id:
                resolved = self._agents.get(selected_id)
                if resolved:
                    logger.info(f"🎯 Agent resolved (adaptive): '{resolved.name}'")
                    return resolved

        # Router mode — ask the router LLM to pick
        if self._mode == TopologyMode.ROUTER and self._router_agent:
            agent_names = [a.name for a in self._agents.values()]
            pick_prompt = (
                f"Available agents: {agent_names}\n"
                f"Task: {task.description}\n"
                "Which agent should handle this? Reply with the agent name only."
            )
            pick_task = Task(name="_pick", description=pick_prompt)
            pick_task.conversation.append(Message.user(pick_prompt))
            result = await self._router_agent.run(pick_task)
            chosen_name = result.content.strip()
            for a in self._agents.values():
                if a.name.lower() == chosen_name.lower():
                    logger.info(f"🎯 Agent resolved (router): '{a.name}'")
                    return a

        # Fallback: first registered agent
        fallback = next(iter(self._agents.values()))
        logger.info(f"🎯 Agent resolved (fallback): '{fallback.name}'")
        return fallback

    async def chat(
        self,
        prompt: str,
        agent: Agent | str | None = None,
        constraints: TaskConstraints | None = None,
        context: dict[str, Any] | None = None,
    ) -> str:
        """
        Multi-turn conversation (non-streaming).

        Each call appends the user prompt and assistant response to
        ``self._conversation_history`` so subsequent calls carry the full
        dialogue context.

        Usage::

            answer1 = await team.chat("你好")
            answer2 = await team.chat("继续上面的话题")  # remembers answer1
        """
        task = Task(
            name="chat",
            description=prompt,
            input_data=context or {},
            constraints=constraints or TaskConstraints(),
        )

        # Inject conversation history into the task
        for msg in self._conversation_history:
            if msg["role"] == "user":
                task.conversation.append(Message.user(msg["content"]))
            else:
                task.conversation.append(Message.agent(content=msg["content"], agent_name="assistant"))

        # Append the new user message
        task.conversation.append(Message.user(prompt))

        # Capability inference (adaptive mode)
        if (
            self._mode == TopologyMode.ADAPTIVE
            and self._capability_inference is not None
        ):
            try:
                inferred = await self._capability_inference.infer(prompt)
                if not inferred.is_zero():
                    task.required_capabilities = inferred
            except Exception as e:
                logger.warning(f"Capability inference failed: {e}")

        resolved_agent = await self._resolve_agent_for_task(task, agent)

        with self._trace_context("team.chat", prompt=prompt):
            result = await resolved_agent.run(task)

        # Persist to history
        self._conversation_history.append({"role": "user", "content": prompt})
        self._conversation_history.append({"role": "assistant", "content": result.content})

        return result.content

    async def stream(
        self,
        prompt: str,
        agent: Agent | str | None = None,
        constraints: TaskConstraints | None = None,
        context: dict[str, Any] | None = None,
    ) -> AsyncIterator[str]:
        """
        Single-turn streaming — yields text chunks as they arrive.

        Usage::

            async for chunk in team.stream("写一首诗"):
                print(chunk, end="", flush=True)
        """
        task = Task(
            name="stream",
            description=prompt,
            input_data=context or {},
            constraints=constraints or TaskConstraints(),
        )
        task.conversation.append(Message.user(prompt))

        # Capability inference (adaptive mode)
        if (
            self._mode == TopologyMode.ADAPTIVE
            and self._capability_inference is not None
        ):
            try:
                inferred = await self._capability_inference.infer(prompt)
                if not inferred.is_zero():
                    task.required_capabilities = inferred
            except Exception as e:
                logger.warning(f"Capability inference failed: {e}")

        resolved_agent = await self._resolve_agent_for_task(task, agent)

        with self._trace_context("team.stream", prompt=prompt):
            async for chunk in resolved_agent.run_stream(task):
                yield chunk

    async def chat_stream(
        self,
        prompt: str,
        agent: Agent | str | None = None,
        constraints: TaskConstraints | None = None,
        context: dict[str, Any] | None = None,
    ) -> AsyncIterator[str]:
        """
        Multi-turn conversation **with** streaming — the best of both worlds.

        Keeps conversation history like ``chat()`` but yields tokens like
        ``stream()``.

        Usage::

            async for chunk in team.chat_stream("你好"):
                print(chunk, end="", flush=True)
            print()
            async for chunk in team.chat_stream("继续"):
                print(chunk, end="", flush=True)
        """
        task = Task(
            name="chat_stream",
            description=prompt,
            input_data=context or {},
            constraints=constraints or TaskConstraints(),
        )

        # Inject conversation history
        for msg in self._conversation_history:
            if msg["role"] == "user":
                task.conversation.append(Message.user(msg["content"]))
            else:
                task.conversation.append(Message.agent(content=msg["content"], agent_name="assistant"))

        task.conversation.append(Message.user(prompt))

        # Capability inference (adaptive mode)
        if (
            self._mode == TopologyMode.ADAPTIVE
            and self._capability_inference is not None
        ):
            try:
                inferred = await self._capability_inference.infer(prompt)
                if not inferred.is_zero():
                    task.required_capabilities = inferred
            except Exception as e:
                logger.warning(f"Capability inference failed: {e}")

        resolved_agent = await self._resolve_agent_for_task(task, agent)

        collected = ""
        with self._trace_context("team.chat_stream", prompt=prompt):
            async for chunk in resolved_agent.run_stream(task):
                collected += chunk
                yield chunk

        # Persist to history
        self._conversation_history.append({"role": "user", "content": prompt})
        self._conversation_history.append({"role": "assistant", "content": collected})

    def reset_conversation(self) -> None:
        """Clear conversation history to start a fresh dialogue."""
        self._conversation_history.clear()

    async def close(self) -> None:
        """Clean up all resources."""
        await self._default_provider.close()
        for agent in self._agents.values():
            if agent.model_provider and agent.model_provider is not self._default_provider:
                await agent.model_provider.close()

    # ---- Context managers ----

    async def __aenter__(self) -> Team:
        return self

    async def __aexit__(self, *exc: Any) -> None:
        await self.close()

    # ---- Helpers ----

    def _trace_context(self, name: str, **attrs: Any):
        """Return a tracing span if tracing is enabled, otherwise a no-op."""
        if self.tracer:
            return self.tracer.span(name, attributes=attrs)

        # Return a no-op context manager
        import contextlib
        return contextlib.nullcontext()


# ---- Convenience sync wrapper ----

def run(
    prompt: str,
    model: str = "gpt-4o",
    agents: list[dict[str, Any]] | None = None,
    **kwargs: Any,
) -> str:
    """
    One-liner sync entry point for quick scripts.

    Usage:
        from nexus.api import run
        result = run("Summarize this article", model="gpt-4o")
    """
    async def _inner() -> str:
        async with Team(model=model, **kwargs) as team:
            if agents:
                for agent_def in agents:
                    team.add_agent(**agent_def)
            return await team.run(prompt)

    return asyncio.run(_inner())
