"""
Agent — the primary actor in the Nexus framework.

Design notes:
  - Agents are defined by: identity + capabilities + instructions + tools + lifecycle hooks
  - Each agent auto-generates an AgentCard (A2A-compatible) for discovery
  - Tools can be plain functions, MCP servers, or other agents (via AgentTool)
  - Agents run an async event loop: receive task → plan → execute → reflect → respond
  - The agent itself is model-agnostic; the ModelProvider is injected
"""

from __future__ import annotations

import inspect
import json
import logging
from collections.abc import AsyncIterator, Callable, Sequence
from typing import Any

from pydantic import BaseModel, Field

from nexus.core.handoff import Handoff
from nexus.core.message import Conversation, Message
from nexus.core.task import Task, TaskResult
from nexus.types import AgentCapability, AgentCard, TokenUsage, new_id

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


# ---------------------------------------------------------------------------
# Tool definition
# ---------------------------------------------------------------------------


class ToolDefinition(BaseModel):
    """Schema for a tool that an agent can call."""

    name: str
    description: str = ""
    parameters: dict[str, Any] = Field(default_factory=dict)
    # The actual callable is stored separately (not serializable)


class AgentTool:
    """
    Wraps a callable (sync or async) as a tool the agent can invoke.

    Automatically extracts the JSON schema from type hints and docstrings.
    """

    def __init__(self, func: Callable[..., Any], name: str | None = None, description: str | None = None) -> None:
        self.func = func
        self.name = name or func.__name__
        self.description = description or (func.__doc__ or "").strip()
        self.is_async = inspect.iscoroutinefunction(func)
        self._schema = self._extract_schema()

    def _extract_schema(self) -> dict[str, Any]:
        """Extract a JSON-schema-like parameter spec from the function signature."""
        sig = inspect.signature(self.func)
        properties: dict[str, Any] = {}
        required: list[str] = []

        for param_name, param in sig.parameters.items():
            if param_name in ("self", "cls"):
                continue

            prop: dict[str, Any] = {}
            annotation = param.annotation

            if annotation is str or annotation is inspect.Parameter.empty:
                prop["type"] = "string"
            elif annotation is int:
                prop["type"] = "integer"
            elif annotation is float:
                prop["type"] = "number"
            elif annotation is bool:
                prop["type"] = "boolean"
            elif annotation is list:
                prop["type"] = "array"
            elif annotation is dict:
                prop["type"] = "object"
            else:
                prop["type"] = "string"

            if param.default == inspect.Parameter.empty:
                required.append(param_name)
            else:
                prop["default"] = param.default

            properties[param_name] = prop

        return {
            "type": "object",
            "properties": properties,
            "required": required,
        }

    @property
    def definition(self) -> ToolDefinition:
        return ToolDefinition(
            name=self.name,
            description=self.description,
            parameters=self._schema,
        )

    def to_openai_tool(self) -> dict[str, Any]:
        """Convert to OpenAI function calling format."""
        return {
            "type": "function",
            "function": {
                "name": self.name,
                "description": self.description,
                "parameters": self._schema,
            },
        }

    async def __call__(self, **kwargs: Any) -> Any:
        """Execute the tool."""
        if self.is_async:
            return await self.func(**kwargs)
        return self.func(**kwargs)


def tool(
    name: str | None = None,
    description: str | None = None,
) -> Callable[[Callable[..., Any]], AgentTool]:
    """Decorator to convert a function into an AgentTool."""

    def decorator(func: Callable[..., Any]) -> AgentTool:
        return AgentTool(func, name=name, description=description)

    return decorator


# ---------------------------------------------------------------------------
# Lifecycle hooks
# ---------------------------------------------------------------------------


class AgentHooks(BaseModel):
    """
    Optional lifecycle hooks for customizing agent behavior.

    All hooks are agent_id-aware so the same hook function
    can be shared across agents.
    """

    class Config:
        arbitrary_types_allowed = True

    # Called before the agent processes a task
    on_task_start: Callable[[str, Task], Any] | None = None
    # Called after a task completes
    on_task_complete: Callable[[str, Task, TaskResult], Any] | None = None
    # Called when a task fails
    on_task_error: Callable[[str, Task, Exception], Any] | None = None
    # Called before each LLM call
    on_llm_call: Callable[[str, list[Message]], Any] | None = None
    # Called after each LLM call
    on_llm_response: Callable[[str, Message, TokenUsage], Any] | None = None
    # Called when a handoff is initiated
    on_handoff: Callable[[str, Handoff], Any] | None = None


# ---------------------------------------------------------------------------
# Agent
# ---------------------------------------------------------------------------


class Agent:
    """
    The primary actor in the Nexus framework.

    An Agent is defined by:
    - Identity: name, description, capabilities
    - Instructions: system prompt guiding behavior
    - Tools: callable functions the agent can invoke
    - Model: the LLM provider (injected, not owned)
    - Hooks: optional lifecycle callbacks

    Usage:
        agent = Agent(
            name="Researcher",
            instructions="You are a senior research analyst...",
            capabilities=AgentCapability(research=0.9, writing=0.7),
            tools=[search_tool, scrape_tool],
        )
        result = await agent.run(task)
    """

    def __init__(
        self,
        name: str,
        instructions: str | Callable[[], str] = "You are a helpful assistant.",
        capabilities: AgentCapability | None = None,
        tools: Sequence[AgentTool | Callable[..., Any]] = (),
        model_provider: Any | None = None,  # Type: ModelProvider (avoid circular import)
        hooks: AgentHooks | None = None,
        max_iterations: int = 10,
        metadata: dict[str, Any] | None = None,
    ) -> None:
        self.id = new_id()
        self.name = name
        self._instructions = instructions
        self.capabilities = capabilities or AgentCapability()
        self.model_provider = model_provider
        self.hooks = hooks or AgentHooks()
        self.max_iterations = max_iterations
        self.metadata = metadata or {}

        # Register tools
        self._tools: dict[str, AgentTool] = {}
        for t in tools:
            if isinstance(t, AgentTool):
                self._tools[t.name] = t
            elif callable(t):
                wrapped = AgentTool(t)
                self._tools[wrapped.name] = wrapped
            else:
                raise TypeError(f"Expected AgentTool or callable, got {type(t)}")

        # Runtime state
        self._total_usage = TokenUsage()
        self._task_count = 0

    @property
    def instructions(self) -> str:
        """Get the current instructions (supports dynamic instructions via callable)."""
        if callable(self._instructions):
            return self._instructions()
        return self._instructions

    @property
    def tools(self) -> list[AgentTool]:
        return list(self._tools.values())

    @property
    def tool_names(self) -> list[str]:
        return list(self._tools.keys())

    @property
    def total_usage(self) -> TokenUsage:
        return self._total_usage

    def add_tool(self, tool: AgentTool | Callable[..., Any]) -> None:
        """Register a new tool."""
        if not isinstance(tool, AgentTool):
            tool = AgentTool(tool)
        self._tools[tool.name] = tool

    def remove_tool(self, name: str) -> None:
        """Unregister a tool."""
        self._tools.pop(name, None)

    def get_tool(self, name: str) -> AgentTool | None:
        """Look up a tool by name."""
        return self._tools.get(name)

    # ---- Core execution loop ----

    async def run(self, task: Task) -> TaskResult:
        """
        Execute a task through the agent's planning-execution loop.

        Loop: receive → think (LLM) → act (tools/handoff) → observe → repeat
        Terminates when: the agent produces a final answer, max iterations reached,
        or constraints are violated.
        """
        if not self.model_provider:
            raise RuntimeError(
                f"Agent '{self.name}' has no model_provider. "
                "Assign one via agent.model_provider = provider"
            )

        task.start(self.id)
        self._task_count += 1

        # Notify hook
        if self.hooks.on_task_start:
            self.hooks.on_task_start(self.id, task)

        # Build conversation
        conversation = task.conversation
        conversation.append(Message.system(self.instructions))

        # Add task description as user message
        task_prompt = self._build_task_prompt(task)
        conversation.append(Message.user(task_prompt))

        try:
            result = await self._execution_loop(task, conversation)
            task.complete(result)

            if self.hooks.on_task_complete:
                self.hooks.on_task_complete(self.id, task, result)

            return result

        except Exception as exc:
            task.fail(str(exc))
            if self.hooks.on_task_error:
                self.hooks.on_task_error(self.id, task, exc)
            raise
    async def run_stream(self, task: Task) -> AsyncIterator[str]:
        """
        Execute a task and stream the response token-by-token.

        Yields text chunks as they arrive from the LLM.  This is a
        simplified path: no tool calling, no ReAct loop — just
        streaming the direct LLM response.

        For tasks that require tool use, fall back to ``run()``.
        """
        if not self.model_provider:
            raise RuntimeError(
                f"Agent '{self.name}' has no model_provider. "
                "Assign one via agent.model_provider = provider"
            )

        task.start(self.id)
        self._task_count += 1

        # Build conversation
        conversation = task.conversation
        conversation.append(Message.system(self.instructions))

        task_prompt = self._build_task_prompt(task)
        conversation.append(Message.user(task_prompt))

        messages = conversation.to_openai_messages()

        # Stream from the provider
        collected_content = ""
        async for chunk in self.model_provider.chat_stream(messages=messages):
            collected_content += chunk
            yield chunk

        # Record the assistant message for conversation history
        conversation.append(Message.agent(content=collected_content, agent_name=self.name))

        # Mark task as complete
        result = TaskResult(
            content=collected_content,
            metadata={"streamed": True},
        )
        task.complete(result)
    async def _execution_loop(self, task: Task, conversation: Conversation) -> TaskResult:
        """The core ReAct loop: Think → Act → Observe → Repeat."""
        tool_names = ", ".join(self.tool_names) if self.tools else "(none)"
        logger.info(
            f"🔄 [{self.name}] Starting ReAct loop "
            f"(max_iterations={self.max_iterations}, tools=[{tool_names}])"
        )

        for iteration in range(self.max_iterations):
            logger.info(
                f"💭 [{self.name}] Iteration {iteration + 1}/{self.max_iterations} — calling LLM..."
            )

            # Check budget constraints
            if task.constraints.max_tokens and task.token_usage.total_tokens >= task.constraints.max_tokens:
                logger.warning(f"⚠️  [{self.name}] Token budget exceeded ({task.token_usage.total_tokens} tokens), stopping")
                break

            # ---- THINK: call the LLM ----
            if self.hooks.on_llm_call:
                self.hooks.on_llm_call(self.id, conversation.messages)

            openai_tools = [t.to_openai_tool() for t in self.tools] if self.tools else None

            response = await self.model_provider.chat(
                messages=conversation.to_openai_messages(),
                tools=openai_tools,
            )

            # Track usage
            usage = response.get("usage", TokenUsage())
            if isinstance(usage, dict):
                usage = TokenUsage(**usage)
            task.token_usage = task.token_usage + usage
            self._total_usage = self._total_usage + usage

            assistant_msg = response["message"]
            finish_reason = response.get("finish_reason", "unknown")

            logger.info(
                f"📨 [{self.name}] LLM responded "
                f"(finish_reason={finish_reason}, "
                f"tokens={usage.prompt_tokens}+{usage.completion_tokens}={usage.total_tokens})"
            )

            if self.hooks.on_llm_response:
                self.hooks.on_llm_response(self.id, assistant_msg, usage)

            conversation.append(assistant_msg)

            # ---- ACT: handle tool calls or produce final answer ----
            if assistant_msg.tool_calls:
                n_calls = len(assistant_msg.tool_calls)
                call_summary = ", ".join(
                    tc.get("function", {}).get("name", "?") for tc in assistant_msg.tool_calls
                )
                logger.info(
                    f"🔧 [{self.name}] LLM requested {n_calls} tool call(s): [{call_summary}]"
                )

                # Execute tool calls
                for idx, tool_call in enumerate(assistant_msg.tool_calls, 1):
                    func_info = tool_call.get("function", {})
                    tc_name = func_info.get("name", "?")
                    tc_args = func_info.get("arguments", "{}")
                    logger.info(
                        f"  🛠️  [{self.name}] Executing tool [{idx}/{n_calls}]: "
                        f"{tc_name}({tc_args})"
                    )

                    tool_result = await self._execute_tool_call(tool_call)

                    # Check if the tool returned a Handoff
                    if isinstance(tool_result, Handoff):
                        logger.info(
                            f"  🤝 [{self.name}] Tool returned Handoff → "
                            f"{tool_result.target_agent_id or tool_result.target_agent_ids}"
                        )
                        if self.hooks.on_handoff:
                            self.hooks.on_handoff(self.id, tool_result)
                        # Return handoff as a special result
                        return TaskResult(
                            content=f"[HANDOFF:{tool_result.mode.value}] → {tool_result.target_agent_id or tool_result.target_agent_ids}",
                            data={"handoff": tool_result.model_dump()},
                            metadata={"handoff_triggered": True},
                        )

                    # Log tool result (truncate if too long)
                    result_str = str(tool_result) if not isinstance(tool_result, str) else tool_result
                    display_result = result_str[:200] + "..." if len(result_str) > 200 else result_str
                    logger.info(
                        f"  ✅ [{self.name}] Tool {tc_name} returned: {display_result}"
                    )

                    # Add tool result to conversation
                    tool_call_id = tool_call.get("id", "")
                    conversation.append(
                        Message.tool_result(result_str, tool_call_id=tool_call_id)
                    )

                logger.info(
                    f"🔁 [{self.name}] All tool calls done, looping back to LLM for next step..."
                )
            else:
                # No tool calls — this is the final answer
                content_preview = (assistant_msg.content or "")[:150]
                logger.info(
                    f"✨ [{self.name}] Final answer (iteration {iteration + 1}): "
                    f"{content_preview}{'...' if len(assistant_msg.content or '') > 150 else ''}"
                )
                return TaskResult(
                    content=assistant_msg.content,
                    metadata={"iterations": iteration + 1},
                )

        # Max iterations reached
        logger.warning(
            f"⚠️  [{self.name}] Max iterations ({self.max_iterations}) reached without final answer"
        )
        last_msg = conversation.last
        return TaskResult(
            content=last_msg.content if last_msg else "",
            metadata={"max_iterations_reached": True},
        )

    async def _execute_tool_call(self, tool_call: dict[str, Any]) -> Any:
        """Execute a single tool call and return the result."""
        func_info = tool_call.get("function", {})
        tool_name = func_info.get("name", "")
        arguments_str = func_info.get("arguments", "{}")

        tool_obj = self.get_tool(tool_name)
        if not tool_obj:
            logger.warning(f"❌ [{self.name}] Tool '{tool_name}' not found. Available: {self.tool_names}")
            return f"Error: Tool '{tool_name}' not found. Available tools: {self.tool_names}"

        try:
            arguments = json.loads(arguments_str) if isinstance(arguments_str, str) else arguments_str
            result = await tool_obj(**arguments)
            return result
        except Exception as exc:
            logger.error(f"❌ [{self.name}] Tool '{tool_name}' execution error: {exc}")
            return f"Error executing tool '{tool_name}': {exc}"

    def _build_task_prompt(self, task: Task) -> str:
        """Construct the user prompt from a task."""
        parts = []
        if task.description:
            parts.append(task.description)
        if task.input_data:
            parts.append(f"\n**Input data:**\n```json\n{json.dumps(task.input_data, ensure_ascii=False, indent=2)}\n```")
        if task.input_messages:
            for msg in task.input_messages:
                parts.append(f"\n[{msg.role.value}]: {msg.content}")
        return "\n".join(parts) if parts else "Please complete the assigned task."

    # ---- Agent Card (A2A-compatible) ----

    def to_agent_card(self, endpoint: str | None = None) -> AgentCard:
        """Generate an A2A-compatible AgentCard for this agent."""
        return AgentCard(
            agent_id=self.id,
            name=self.name,
            description=self.instructions[:200] if isinstance(self._instructions, str) else "",
            capabilities=self.capabilities,
            skills=[t.name for t in self.tools],
            tools=[t.name for t in self.tools],
            endpoint=endpoint,
            metadata=self.metadata,
        )

    # ---- Serialization ----

    def to_dict(self) -> dict[str, Any]:
        """Serialize agent configuration (not the runtime state)."""
        return {
            "id": self.id,
            "name": self.name,
            "instructions": self.instructions,
            "capabilities": self.capabilities.model_dump(),
            "tools": [t.definition.model_dump() for t in self.tools],
            "max_iterations": self.max_iterations,
            "metadata": self.metadata,
        }

    def __repr__(self) -> str:
        return (
            f"Agent(id={self.id!r}, name={self.name!r}, "
            f"tools={self.tool_names}, capabilities={self.capabilities})"
        )
