# Copyright (c) 2023-2026, HaiyangLi <quantocean.li at gmail dot com>
# SPDX-License-Identifier: Apache-2.0

"""AG2 GroupChat endpoint: delegates to build_group_chat/stream_group_chat and converts events to StreamChunk."""

from __future__ import annotations

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

from pydantic import BaseModel, Field

from lionagi.ln.concurrency.utils import maybe_await
from lionagi.service.connections import AgenticEndpoint, EndpointConfig
from lionagi.service.types import StreamChunk
from lionagi.utils import to_dict

from ._config import AG2Configs

logger = logging.getLogger(__name__)

__all__ = [
    "HandoffCondition",
    "AgentSpec",
    "GroupChatSpec",
    "ResearchPlan",
    "AG2GroupChatRequest",
    "build_group_chat",
    "stream_group_chat",
    "AG2GroupChatEndpoint",
]


# ---------------------------------------------------------------------------
# Models (structured output targets)
# ---------------------------------------------------------------------------


class HandoffCondition(BaseModel):
    """Condition for handing off to another agent in AG2 GroupChat."""

    target: str = Field(description="Target agent name to hand off to")
    condition: str = Field(
        description="Natural language condition for handoff trigger",
    )


class AgentSpec(BaseModel):
    """Specification for a single AG2 ConversableAgent in a GroupChat."""

    name: str = Field(description="Agent name (e.g. 'Researcher', 'Analyst')")
    role: str = Field(description="One-line role description")
    system_message: str = Field(default="", description="Full system prompt for the agent")
    tools: list[str] = Field(
        default_factory=list,
        description="Tool names from the registry this agent can invoke",
    )
    handoffs: list[HandoffCondition] = Field(
        default_factory=list,
        description="Conditions for handing off to other agents",
    )
    nlip_url: str | None = Field(
        default=None,
        description="If set, this agent is a remote NLIP agent at the given URL. "
        "system_message and llm_config are ignored — the remote server handles them.",
    )
    state_template: str | None = Field(
        default=None,
        description=(
            "Dynamic system message template using context variables. "
            "Maps to AG2 v0.12 UpdateSystemMessage. "
            "Example: 'Insights stored: {insights_stored}'"
        ),
    )


class GroupChatSpec(BaseModel):
    """Specification for a single AG2 GroupChat — one node in a DAG."""

    name: str = Field(description="GroupChat identifier")
    objective: str = Field(description="What this GroupChat should accomplish")
    agents: list[AgentSpec] = Field(description="Participating agents")
    max_round: int = Field(default=15, gt=0, description="Maximum conversation rounds")
    context: dict[str, Any] = Field(
        default_factory=dict,
        description="Initial context variables shared across agents",
    )


class ResearchPlan(BaseModel):
    """Full research plan generated by LLM via structured output."""

    topic: str = Field(description="The research topic")
    hypothesis: str = Field(description="Central hypothesis to test")
    group_chats: list[GroupChatSpec] = Field(
        description="Independent GroupChat teams running in parallel",
    )
    synthesis_instruction: str = Field(
        description="How to combine all GroupChat outputs",
    )
    expected_output: str = Field(
        description="What the final deliverable should look like",
    )
    safeguard_policy: list[str] = Field(
        default_factory=list,
        description="Research governance guidelines enforced at conversation level",
    )


class AG2GroupChatRequest(BaseModel):
    """Request for AG2 GroupChat endpoint."""

    messages: list[dict[str, Any]] = Field(default_factory=list)
    prompt: str = ""
    max_round: int = Field(default=15, gt=0)
    context_variables: dict[str, Any] = Field(default_factory=dict)


# ---------------------------------------------------------------------------
# Callback types (parallel to claude_code.py on_* handlers)
# ---------------------------------------------------------------------------

AG2_HANDLER_PARAMS = (
    "on_text",
    "on_tool_use",
    "on_tool_result",
    "on_speaker",
    "on_complete",
)


# ---------------------------------------------------------------------------
# Runner (in-process — no subprocess, unlike claude_code.py)
# ---------------------------------------------------------------------------


def build_group_chat(
    spec: GroupChatSpec,
    llm_config: Any,
    tool_registry: dict[str, Callable] | None = None,
    code_executor: Any | None = None,
):
    """Build AG2 agents and DefaultPattern from a GroupChatSpec.

    Returns (user_proxy, pattern, agents_by_name).
    """
    from autogen import ConversableAgent, register_function
    from autogen.agentchat.conversable_agent import UpdateSystemMessage
    from autogen.agentchat.group import (
        AgentTarget,
        ContextVariables,
        OnCondition,
        StringLLMCondition,
        TerminateTarget,
    )
    from autogen.agentchat.group.patterns import DefaultPattern

    tool_registry = tool_registry or {}
    user_configs = {
        "name": "User",
        "human_input_mode": "NEVER",
        "llm_config": False,
    }
    if code_executor:
        user_configs["code_execution_config"] = {"executor": code_executor}

    user = ConversableAgent(**user_configs)

    agents_by_name: dict[str, ConversableAgent] = {}
    ordered: list[ConversableAgent] = []

    for agent_spec in spec.agents:
        if agent_spec.nlip_url:
            # SSRF guard: validate the caller-supplied URL before handing it
            # to NlipRemoteAgent, which opens its own HTTP connection and
            # bypasses the Endpoint transport guards.
            from lionagi.providers.ag2.nlip import _assert_nlip_url_safe

            _assert_nlip_url_safe(agent_spec.nlip_url)
            try:
                from autogen.agentchat.contrib.nlip_agent import NlipRemoteAgent

                agent = NlipRemoteAgent(
                    url=agent_spec.nlip_url,
                    name=agent_spec.name,
                )
            except ImportError:
                logger.warning(
                    "nlip_sdk not installed — creating %r as local agent instead",
                    agent_spec.name,
                )
                agent = ConversableAgent(
                    name=agent_spec.name,
                    system_message=agent_spec.system_message or f"You are {agent_spec.name}.",
                    llm_config=llm_config,
                    human_input_mode="NEVER",
                )
        else:
            kwargs: dict[str, Any] = dict(
                name=agent_spec.name,
                system_message=agent_spec.system_message,
                llm_config=llm_config,
                human_input_mode="NEVER",
            )
            if agent_spec.state_template:
                kwargs["update_agent_state_before_reply"] = [
                    UpdateSystemMessage(agent_spec.state_template),
                ]

            agent = ConversableAgent(**kwargs)

            for tool_name in agent_spec.tools:
                if tool_name in tool_registry:
                    fn = tool_registry[tool_name]
                    desc = getattr(fn, "__doc__", "") or tool_name
                    register_function(
                        fn,
                        caller=agent,
                        executor=user,
                        description=desc,
                    )

        agents_by_name[agent_spec.name] = agent
        ordered.append(agent)

    for idx, agent_spec in enumerate(spec.agents):
        agent = agents_by_name[agent_spec.name]
        for hc in agent_spec.handoffs:
            if hc.target in agents_by_name:
                agent.handoffs.add_llm_conditions(
                    [
                        OnCondition(
                            target=AgentTarget(agents_by_name[hc.target]),
                            condition=StringLLMCondition(prompt=hc.condition),
                        )
                    ]
                )
            else:
                logger.warning(
                    "Handoff target %r from agent %r not found in spec — skipped",
                    hc.target,
                    agent_spec.name,
                )

        # Fallback: flow to next agent in spec order if no LLM-condition handoff fires.
        if idx < len(spec.agents) - 1:
            next_agent = agents_by_name[spec.agents[idx + 1].name]
            agent.handoffs.set_after_work(AgentTarget(next_agent))

    if ordered:
        ordered[-1].handoffs.set_after_work(TerminateTarget())

    pattern = DefaultPattern(
        initial_agent=ordered[0] if ordered else user,
        agents=ordered,
        user_agent=user,
        context_variables=ContextVariables(data=spec.context),
    )

    return user, pattern, agents_by_name


async def stream_group_chat(
    pattern,
    prompt: str,
    max_rounds: int = 15,
    safeguard_policy: dict[str, Any] | None = None,
    yield_on: list | None = None,
    on_text: Callable[[str, str], None] | None = None,
    on_tool_use: Callable[[str, str, Any], None] | None = None,
    on_tool_result: Callable[[str, Any], None] | None = None,
    on_speaker: Callable[[str], None] | None = None,
    on_complete: Callable[[Any], None] | None = None,
) -> AsyncIterator[Any]:
    """Stream AG2 GroupChat events. In-process, no subprocess.

    Yields AG2 BaseEvent objects. Optionally calls on_* handlers.
    """
    from autogen.agentchat.group.multi_agent_chat import a_run_group_chat_iter
    from autogen.events.agent_events import (
        GroupChatRunChatEvent,
        SelectSpeakerEvent,
        TextEvent,
        ToolCallEvent,
        ToolResponseEvent,
    )

    if yield_on is None:
        yield_on = [
            TextEvent,
            ToolCallEvent,
            ToolResponseEvent,
            SelectSpeakerEvent,
            GroupChatRunChatEvent,
        ]

    kwargs: dict[str, Any] = dict(
        pattern=pattern,
        messages=prompt,
        max_rounds=max_rounds,
        yield_on=yield_on,
    )
    if safeguard_policy:
        kwargs["safeguard_policy"] = safeguard_policy

    try:
        async for event in a_run_group_chat_iter(**kwargs):
            inner = getattr(event, "content", None)

            if isinstance(event, TextEvent) and on_text:
                text = getattr(inner, "content", "") if inner is not None else ""
                sender = getattr(inner, "sender", "unknown") if inner is not None else "unknown"
                await maybe_await(on_text(text, sender))
            elif isinstance(event, ToolCallEvent) and on_tool_use:
                tool_calls = getattr(inner, "tool_calls", []) if inner is not None else []
                first = tool_calls[0] if tool_calls else None
                tool_name = getattr(getattr(first, "function", None), "name", "") if first else ""
                tool_args = (
                    getattr(getattr(first, "function", None), "arguments", None) if first else None
                )
                sender = getattr(inner, "sender", "unknown") if inner is not None else "unknown"
                await maybe_await(on_tool_use(tool_name, sender, tool_args))
            elif isinstance(event, ToolResponseEvent) and on_tool_result:
                tool_responses = getattr(inner, "tool_responses", []) if inner is not None else []
                first = tool_responses[0] if tool_responses else None
                tool_output = getattr(first, "content", None) if first else None
                sender = getattr(inner, "sender", "unknown") if inner is not None else "unknown"
                await maybe_await(on_tool_result(sender, tool_output))
            elif isinstance(event, GroupChatRunChatEvent) and on_speaker:
                speaker = getattr(inner, "speaker", "unknown") if inner is not None else "unknown"
                await maybe_await(on_speaker(speaker))
            elif isinstance(event, SelectSpeakerEvent) and on_speaker:
                agents = getattr(inner, "agents", []) if inner is not None else []
                first_name = getattr(agents[0], "name", str(agents[0])) if agents else "unknown"
                await maybe_await(on_speaker(first_name))
            yield event
    except Exception:
        logger.exception("AG2 GroupChat streaming failed")
        raise

    if on_complete:
        await maybe_await(on_complete(None))


logger = logging.getLogger(__name__)


@AG2Configs.GROUP_CHAT.register
class AG2GroupChatEndpoint(AgenticEndpoint):
    """Wraps AG2 v0.12 GroupChat as a lionagi agentic endpoint; stream-only."""

    DEFAULT_CONCURRENCY_LIMIT = 1
    DEFAULT_QUEUE_CAPACITY = 3

    def __init__(self, config: EndpointConfig | None = None, **kwargs):
        super().__init__(config=config, **kwargs)
        self._agent_configs: list[dict[str, Any]] = kwargs.get("agent_configs", [])
        self._llm_config: dict[str, Any] = kwargs.get("llm_config", {})
        self._tool_registry: dict[str, Any] = kwargs.get("tool_registry", {})

    async def _call(self, payload, headers, **kwargs):
        raise NotImplementedError("AG2 GroupChat is stream-only. Use stream() to iterate events.")

    def create_payload(self, request: dict | BaseModel, **kwargs):

        req_dict = {**self.config.kwargs, **to_dict(request), **kwargs}
        messages = req_dict.pop("messages", [])
        prompt = req_dict.pop("prompt", "")
        max_round = req_dict.pop("max_round", 15)
        ctx = req_dict.pop("context_variables", {})
        return {
            "request": AG2GroupChatRequest(
                messages=messages,
                prompt=prompt,
                max_round=max_round,
                context_variables=ctx,
            )
        }, {}

    async def stream(self, request: dict | BaseModel, **kwargs) -> AsyncIterator[StreamChunk]:

        if isinstance(request, dict) and "request" in request:
            request_obj = request["request"]
        else:
            payload, _ = self.create_payload(request, **kwargs)
            request_obj = payload["request"]

        prompt = request_obj.prompt or (
            request_obj.messages[-1]["content"] if request_obj.messages else ""
        )
        if not prompt:
            raise ValueError(
                "AG2GroupChatEndpoint requires a non-empty prompt or at least one message."
            )

        agent_configs = kwargs.get("agent_configs", self._agent_configs)
        llm_config = kwargs.get("llm_config", self._llm_config)
        tool_registry = kwargs.get("tool_registry", self._tool_registry)
        code_executor = kwargs.get("code_executor")

        spec = GroupChatSpec(
            name="endpoint_chat",
            objective=prompt,
            agents=[
                {
                    "name": c["name"],
                    "role": c.get("role", ""),
                    "system_message": c.get("system_message", ""),
                    "tools": c.get("tools", []),
                    "handoffs": [
                        {"target": h["target"], "condition": h["condition"]}
                        for h in c.get("handoff_conditions", c.get("handoffs", []))
                    ],
                    "nlip_url": c.get("nlip_url"),
                    "state_template": c.get("state_template"),
                }
                for c in agent_configs
            ],
            context=request_obj.context_variables,
            max_round=request_obj.max_round,
        )

        user, pattern, agents_by_name = build_group_chat(
            spec, llm_config, tool_registry, code_executor
        )

        yield StreamChunk(
            type="system",
            metadata={
                "provider": "ag2",
                "api": "v0.12",
                "pattern": "DefaultPattern",
                "agent_count": len(agent_configs),
                "max_round": request_obj.max_round,
            },
        )

        async for event in stream_group_chat(
            pattern=pattern,
            prompt=prompt,
            max_rounds=request_obj.max_round,
        ):
            chunk = _event_to_chunk(event)
            if chunk:
                yield chunk

        yield StreamChunk(
            type="result",
            content="GroupChat complete",
            metadata={"agents": list(agents_by_name.keys())},
        )


def _event_to_chunk(event) -> StreamChunk | None:
    """Convert an AG2 wrapped event to a StreamChunk; returns None for unrecognized event types."""
    from autogen.events.agent_events import (
        GroupChatRunChatEvent,
        SelectSpeakerEvent,
        TextEvent,
        ToolCallEvent,
        ToolResponseEvent,
    )

    inner = getattr(event, "content", None)

    if isinstance(event, TextEvent):
        text = getattr(inner, "content", str(event)) if inner is not None else str(event)
        sender = getattr(inner, "sender", "unknown") if inner is not None else "unknown"
        return StreamChunk(
            type="text",
            content=text,
            metadata={"agent": sender},
        )
    if isinstance(event, GroupChatRunChatEvent):
        speaker = getattr(inner, "speaker", "unknown") if inner is not None else "unknown"
        return StreamChunk(
            type="system",
            content=f"Speaker: {speaker}",
            metadata={"event": "speaker_turn", "agent": speaker},
        )
    if isinstance(event, SelectSpeakerEvent):
        agents = getattr(inner, "agents", []) if inner is not None else []
        if agents:
            names = ", ".join(getattr(a, "name", str(a)) for a in agents)
        else:
            names = "?"
        return StreamChunk(
            type="system",
            content=f"Speaker candidates: {names}",
            metadata={"event": "speaker_candidates"},
        )
    if isinstance(event, ToolCallEvent):
        tool_calls = getattr(inner, "tool_calls", []) if inner is not None else []
        first = tool_calls[0] if tool_calls else None
        tool_name = getattr(getattr(first, "function", None), "name", None) if first else None
        tool_args = getattr(getattr(first, "function", None), "arguments", None) if first else None
        sender = getattr(inner, "sender", "unknown") if inner is not None else "unknown"
        return StreamChunk(
            type="tool_use",
            tool_name=tool_name,
            tool_id=None,
            tool_input=tool_args,
            metadata={"agent": sender},
        )
    if isinstance(event, ToolResponseEvent):
        tool_responses = getattr(inner, "tool_responses", []) if inner is not None else []
        first = tool_responses[0] if tool_responses else None
        tool_output = getattr(first, "content", None) if first else None
        tool_id = getattr(first, "tool_call_id", None) if first else None
        sender = getattr(inner, "sender", "unknown") if inner is not None else "unknown"
        return StreamChunk(
            type="tool_result",
            tool_output=tool_output,
            metadata={"agent": sender, "tool_call_id": tool_id},
        )
    return None
