#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# ================================================== #
# This file is a part of PYGPT package               #
# Website: https://pygpt.net                         #
# GitHub:  https://github.com/szczyglis-dev/py-gpt   #
# MIT License                                        #
# Created By  : Marcin Szczygliński                  #
# Updated Date: 2025.09.26 15:00:00                  #
# ================================================== #

from __future__ import annotations

from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Tuple
import inspect
from pydantic import ValidationError

# LlamaIndex workflow / agent
from llama_index.core.workflow import (
    Workflow,
    Context,
    StartEvent,
    StopEvent,
    Event,
    step,
)
from llama_index.core.llms.llm import LLM
from llama_index.core.tools.types import BaseTool

# v12/v13 agent workflow events + agent
from llama_index.core.agent.workflow import (
    FunctionAgent,
    ToolCall,
    ToolCallResult,
    AgentStream,
    AgentOutput,
    AgentInput,  # ensure AgentInput propagation includes agent name
)

# v12/v13 compatibility imports
try:
    # v13+
    from llama_index.core.memory import ChatMemoryBuffer
except Exception:  # pragma: no cover
    try:
        # v12
        from llama_index.memory import ChatMemoryBuffer
    except Exception:
        ChatMemoryBuffer = None  # type: ignore

try:
    from llama_index.core.objects.base import ObjectRetriever
except Exception:  # pragma: no cover
    ObjectRetriever = None  # type: ignore

try:
    from llama_index.core.settings import Settings
except Exception:  # pragma: no cover
    Settings = None  # type: ignore

# optional: OpenAI type for hints only
try:
    from llama_index.llms.openai import OpenAI  # noqa: F401
except Exception:
    pass

try:
    from .events import StepEvent  # local helper, same as in your Planner
except Exception:  # pragma: no cover
    StepEvent = None  # type: ignore


DEFAULT_MAX_FUNCTION_CALLS = 5
DEFAULT_SYSTEM_PROMPT = (
    "You are an OpenAI function-calling agent. "
    "Use tools when helpful, reason step-by-step, and produce concise, correct answers."
)


class QueryEvent(StartEvent):
    query: str


class FinalEvent(StopEvent):
    pass


def _safe_tool_name(t: BaseTool) -> str:
    """
    Get a safe tool name from the BaseTool instance.

    :param t: BaseTool instance
    :return: str: Tool name or class name if name is not available
    """
    try:
        # v13 BaseTool.metadata.name
        n = (getattr(t, "metadata", None) or {}).get("name") if isinstance(getattr(t, "metadata", None), dict) else None
        if not n and hasattr(t, "metadata") and hasattr(t.metadata, "name"):
            n = t.metadata.name  # pydantic model
        if not n:
            n = getattr(t, "name", None)
        if not n:
            n = t.__class__.__name__
        return str(n)
    except Exception:
        return t.__class__.__name__


def _list_tool_names(tools: Sequence[BaseTool]) -> List[str]:
    """
    Get a list of safe tool names from a sequence of BaseTool instances.

    :param tools: Sequence of BaseTool instances
    :return: List[str]: List of tool names
    """
    return [_safe_tool_name(t) for t in tools]


class OpenAIWorkflowAgent(Workflow):
    """
    Workflow-based replacement for the legacy OpenAIAgent (v12) using FunctionAgent (v12/v13).

    - memory: tries to set FunctionAgent.memory; falls back to injecting memory text into system prompt.
    - tools: accepts static list or dynamic tool_retriever (query-aware).
    - default_tool_choice: 'auto' | 'none' | '<tool_name>' -> filters visible tools for this run.
    - max_function_calls: mapped to FunctionAgent.max_steps.
    - streaming: forwards AgentStream/ToolCall/ToolCallResult/AgentOutput; emits StepEvent when available.
    """
    def __init__(
        self,
        tools: List[BaseTool],
        llm: LLM,
        memory: Optional[Any] = None,
        system_prompt: Optional[str] = None,
        prefix_messages: Optional[Sequence[Any]] = None,
        verbose: bool = False,
        max_function_calls: int = DEFAULT_MAX_FUNCTION_CALLS,
        default_tool_choice: str = "auto",
        tool_retriever: Optional[Any] = None,
        memory_char_limit: int = 8000,
        on_stop: Optional[Callable[[], bool]] = None,
    ):
        """
        Initialize the OpenAIWorkflowAgent.

        :param tools: List of BaseTool instances to use in the agent.
        :param llm: LLM instance to use for the agent.
        :param memory: Optional memory object to use for the agent. If provided, it will be set on FunctionAgent.memory.
        :param system_prompt: System prompt to use for the agent. If not provided, a default will be used.
        :param prefix_messages: List of messages to prepend to the system prompt for context.
        :param verbose: Verbosity flag for the agent.
        :param max_function_calls: Maximum number of function calls allowed in the agent run. This maps to FunctionAgent.max_steps.
        :param default_tool_choice: Default tool choice for the agent run. Can be 'auto', 'none', or a specific tool name.
        :param tool_retriever: Optional tool retriever to dynamically select tools based on the query.
        :param memory_char_limit: Optional character limit for the memory text representation. If set, will truncate memory text to this limit.
        :param on_stop: Optional callback function that returns a boolean indicating whether the agent should stop running.
        """
        super().__init__(timeout=None, verbose=verbose)
        self._llm = llm
        self._base_system_prompt = system_prompt or DEFAULT_SYSTEM_PROMPT
        self._prefix_messages = list(prefix_messages or [])
        self._tools = tools or []
        self._tool_retriever = tool_retriever
        self._memory = memory
        self._memory_char_limit = memory_char_limit
        self._default_tool_choice = (default_tool_choice or "auto").strip().lower()
        self._max_steps = int(max_function_calls or DEFAULT_MAX_FUNCTION_CALLS)
        self._on_stop = on_stop
        self.verbose = verbose

        # human-friendly display name propagated to UI via workflow events
        self._display_agent_name: str = "FunctionAgent"

        # construct FunctionAgent once, will override tools/system_prompt/memory per run
        self._agent = FunctionAgent(
            name="OpenAIWorkflowAgent",
            description="Workflow-based OpenAI function-calling agent.",
            tools=self._tools,
            llm=self._llm,
            system_prompt=self._base_system_prompt,
            max_steps=self._max_steps,
        )

        # try attach memory now (can be overridden in run())
        if self._memory is not None:
            self._set_agent_memory(self._memory)

    def run(
        self,
        query: str,
        ctx: Optional[Context] = None,
        memory: Optional[Any] = None,
        verbose: Optional[bool] = None,
        **kwargs: Any,
    ):
        """
        Start the workflow answering a single user query.

        :param query: user message
        :param ctx: workflow context
        :param memory: optional memory object to use for this run
        :param verbose: override verbosity
        :return: Workflow run handler (stream_events() supported)
        """
        if verbose is not None:
            self.verbose = bool(verbose)
        if memory is not None:
            self._memory = memory
        # system handles rest via steps
        return super().run(ctx=ctx, query=query)

    # ---------- steps ----------

    @step
    async def answer(self, ctx: Context, ev: QueryEvent) -> FinalEvent:
        """
        Single-step: select tools -> prepare prompt/memory -> run FunctionAgent with streaming.

        :param ctx: Context for the workflow
        :param ev: QueryEvent containing the user query
        :return: FinalEvent with the last answer from the agent
        """
        self._emit_step_event(ctx, name="run", meta={"query": ev.query})

        # prepare memory + prompt
        self._set_agent_memory(self._memory)
        effective_system_prompt = self._compose_system_prompt(self._base_system_prompt, self._prefix_messages, self._memory)

        # select tools for this query
        tools_for_run, selection_reason = await self._select_tools_for_query(ev.query)

        # apply default_tool_choice filter
        tools_for_run = self._apply_default_tool_choice_filter(tools_for_run)

        # update agent config for this run
        try:
            self._agent.system_prompt = effective_system_prompt  # type: ignore[attr-defined]
        except Exception:
            pass
        try:
            self._agent.tools = tools_for_run  # type: ignore[attr-defined]
        except Exception:
            pass

        # log tool selection
        self._emit_step_event(
            ctx,
            name="tools_selected",
            meta={
                "available": _list_tool_names(self._tools),
                "selected": _list_tool_names(tools_for_run),
                "reason": selection_reason,
                "default_tool_choice": self._default_tool_choice,
                "max_steps": self._max_steps,
            },
        )

        # run agent and stream
        last_answer = await self._run_agent_once(ctx, ev.query)
        return FinalEvent(result=last_answer or "")

    # ---------- internals ----------

    def _stopped(self) -> bool:
        """
        Check if the agent should stop running based on the provided callback.

        :return: bool: True if the agent should stop, False otherwise.
        """
        if self._on_stop:
            try:
                return bool(self._on_stop())
            except Exception:
                return False
        return False

    def _emit_step_event(
        self,
        ctx: Context,
        name: str,
        index: Optional[int] = None,
        total: Optional[int] = None,
        meta: Optional[dict] = None,
    ) -> None:
        """
        Emit a step event to the context stream.

        :param ctx: Context for the workflow
        :param name: Name of the step event
        :param index: Index of the step (optional)
        :param total: Total number of steps (optional)
        :param meta: Optional metadata dictionary for the step event
        """
        # Always include agent_name so UI can set it before first token arrives.
        m = dict(meta or {})
        m.setdefault("agent_name", self._display_agent_name)

        try:
            if StepEvent is not None:
                ctx.write_event_to_stream(StepEvent(name=name, index=index, total=total, meta=m))
                return
        except Exception:
            pass

        # Fallback: embed in AgentStream.raw for older AgentStream validators
        try:
            ctx.write_event_to_stream(
                AgentStream(
                    delta="",
                    response="",
                    current_agent_name=self._display_agent_name,
                    tool_calls=[],
                    raw={"StepEvent": {"name": name, "index": index, "total": total, "meta": m}},
                )
            )
        except Exception:
            pass

    def _set_agent_memory(self, memory: Optional[Any]) -> None:
        """
        Set the memory for the FunctionAgent instance.

        :param memory: Optional memory object to set on the agent.
        """
        if memory is None:
            return
        try:
            # Prefer native memory on the FunctionAgent if present
            if hasattr(self._agent, "memory"):
                self._agent.memory = memory  # type: ignore[attr-defined]
        except Exception:
            pass

    def _memory_to_text(self, memory: Any) -> str:
        """
        Convert memory to a text representation, handling various types and structures.

        :param memory: Memory object or content to convert
        :return: str: Text representation of the memory
        """
        if not memory:
            return ""
        try:
            if isinstance(memory, str):
                text = memory
            elif isinstance(memory, list):
                parts = []
                for m in memory:
                    if isinstance(m, str):
                        parts.append(m)
                    elif isinstance(m, dict) and ("content" in m or "text" in m):
                        role = m.get("role", "user")
                        content = m.get("content", m.get("text", ""))
                        parts.append(f"{role}: {content}")
                    else:
                        role = getattr(m, "role", None) or getattr(m, "sender", "user")
                        content = getattr(m, "content", None) or getattr(m, "text", "")
                        parts.append(f"{role}: {content}")
                text = "\n".join(parts)
            else:
                for attr in ("to_string", "to_str"):
                    fn = getattr(memory, attr, None)
                    if callable(fn):
                        text = fn()
                        break
                else:
                    for attr in ("get", "messages", "get_all", "dump"):
                        fn = getattr(memory, attr, None)
                        if callable(fn):
                            data = fn()
                            text = self._memory_to_text(data)
                            break
                    else:
                        text = str(memory)
        except Exception:
            text = str(memory)

        if self._memory_char_limit and len(text) > self._memory_char_limit:
            text = "...[truncated]...\n" + text[-self._memory_char_limit:]
        return text

    def _prefix_to_text(self, prefix_messages: Sequence[Any]) -> str:
        """
        Convert a sequence of prefix messages to a formatted text representation.

        :param prefix_messages: Sequence of messages to convert, can be strings or objects with 'role' and 'content' attributes.
        :return: str: Formatted text representation of the prefix messages.
        """
        if not prefix_messages:
            return ""
        parts: List[str] = []
        for m in prefix_messages:
            if isinstance(m, str):
                parts.append(m.strip())
                continue

            # chat-like
            role = getattr(m, "role", None) or getattr(m, "sender", "system")
            content = getattr(m, "content", None) or getattr(m, "text", "")
            if not content and isinstance(m, dict):
                content = m.get("content", m.get("text", ""))
            if content:
                parts.append(f"{role}: {content}")
        return "\n".join([p for p in parts if p])

    def _compose_system_prompt(
        self,
        base: str,
        prefix_messages: Sequence[Any],
        memory: Optional[Any],
    ) -> str:
        """
        Compose the system prompt for the FunctionAgent, including base prompt,

        :param base: Base system prompt text.
        :param prefix_messages: Sequence of messages to prepend to the system prompt.
        :param memory: Optional memory object to include in the system prompt.
        :return: str: Composed system prompt text.
        """
        out = [base.strip()]
        prefix_text = self._prefix_to_text(prefix_messages)
        if prefix_text:
            out += ["", "Additional preface:", prefix_text]
        mem_text = self._memory_to_text(memory)
        if mem_text:
            out += ["", "Relevant past memory/context:", mem_text]
        return "\n".join(out).strip()

    async def _select_tools_for_query(self, query: str) -> Tuple[List[BaseTool], str]:
        """
        Select tools for the given query, either from static tools or dynamically retrieved tools.

        :param query: User query to select tools for
        :return: Tuple containing the list of selected tools and the reason for selection.
        """
        # default: use provided static tools
        selected = list(self._tools)
        reason = "static tools"

        if not self._tool_retriever:
            return selected, reason

        retriever = self._tool_retriever
        candidates: Optional[Iterable[Any]] = None

        # try a range of method names for compatibility
        for name in ("aretrieve", "aget_retrieved_objects", "retrieve", "get_retrieved_objects"):
            fn = getattr(retriever, name, None)
            if not fn:
                continue
            try:
                if inspect.iscoroutinefunction(fn):
                    candidates = await fn(query)  # type: ignore[misc]
                else:
                    candidates = fn(query)  # type: ignore[misc]
                break
            except Exception:
                candidates = None

        if candidates is None:
            return selected, "retriever failed; using static tools"

        tools: List[BaseTool] = []
        for item in candidates:
            if isinstance(item, BaseTool):
                tools.append(item)
                continue
            # common wrappers
            for attr in ("obj", "object", "tool"):
                cand = getattr(item, attr, None)
                if isinstance(cand, BaseTool):
                    tools.append(cand)
                    break

        if tools:
            selected = tools
            reason = "retrieved tools"
        else:
            reason = "retriever returned no tools; using static tools"

        return selected, reason

    def _apply_default_tool_choice_filter(self, tools: List[BaseTool]) -> List[BaseTool]:
        """
        Apply the default tool choice filter to the list of tools.

        :param tools: List of BaseTool instances to filter.
        :return: List[BaseTool]: Filtered list of tools based on the default tool choice.
        """
        choice = (self._default_tool_choice or "auto").strip().lower()
        if choice in ("auto", "", "default"):
            return tools
        if choice in ("none", "no", "off"):
            return []

        # filter by name
        wanted = choice
        filtered: List[BaseTool] = []
        for t in tools:
            name = _safe_tool_name(t).strip().lower()
            if name == wanted:
                filtered = [t]
                break
        return filtered or tools  # if not found, keep original

    async def _emit_text(
            self,
            ctx: Context,
            text: str,
            agent_name: str = "FunctionAgent"
    ):
        """
        Emit text to the context stream, handling validation errors gracefully.

        :param ctx: Context for the workflow
        :param text: Text to emit to the stream
        :param agent_name: Name of the agent to set in the event (default: "FunctionAgent")
        """
        try:
            ctx.write_event_to_stream(AgentStream(delta=text))
        except ValidationError:
            ctx.write_event_to_stream(
                AgentStream(
                    delta=text,
                    response=text,
                    current_agent_name=agent_name,
                    tool_calls=[],
                    raw={},
                )
            )

    async def _run_agent_once(self, ctx: Context, prompt: str) -> str:
        """
        Run FunctionAgent for a single user message, streaming events.

        :param ctx: Context for the workflow
        :param prompt: User message to process
        :return: Last answer from the agent (text response)
        """
        sig = inspect.signature(self._agent.run)
        kwargs: Dict[str, Any] = {}
        if "user_msg" in sig.parameters:
            kwargs["user_msg"] = prompt
        elif "query" in sig.parameters:
            kwargs["query"] = prompt
        if "max_steps" in sig.parameters:
            kwargs["max_steps"] = self._max_steps

        handler = self._agent.run(**kwargs)
        last_answer = ""
        has_stream = False

        async def _stream():
            nonlocal last_answer, has_stream

            async for e in handler.stream_events():
                if isinstance(e, StopEvent):
                    continue

                # external stop callback
                if self._stopped():
                    try:
                        ctx.write_event_to_stream(StopEvent())
                    except Exception:
                        pass
                    try:
                        await handler.cancel_run()
                    except Exception:
                        pass
                    return last_answer

                if isinstance(e, AgentInput):
                    # Ensure the input event also carries the display name for UI
                    try:
                        e.current_agent_name = self._display_agent_name
                    except Exception:
                        pass
                    ctx.write_event_to_stream(e)
                    continue

                if isinstance(e, AgentStream):
                    if getattr(e, "delta", None):
                        has_stream = True
                    # Always enforce agent name for consistency in UI
                    try:
                        e.current_agent_name = self._display_agent_name
                    except Exception:
                        # If immutable, rebuild a compatible event object
                        try:
                            e = AgentStream(
                                delta=getattr(e, "delta", ""),
                                response=getattr(e, "response", ""),
                                current_agent_name=self._display_agent_name,
                                tool_calls=getattr(e, "tool_calls", []),
                                raw=getattr(e, "raw", {}),
                            )
                        except Exception:
                            pass
                    ctx.write_event_to_stream(e)
                    continue

                if isinstance(e, AgentOutput):
                    resp = getattr(e, "response", None)
                    content = self._to_text(resp).strip()
                    last_answer = content
                    if not has_stream and content:
                        ctx.write_event_to_stream(
                            AgentStream(
                                delta=content,
                                response=content,
                                current_agent_name=self._display_agent_name,
                                tool_calls=getattr(e, "tool_calls", []),
                                raw=getattr(e, "raw", {}),
                            )
                        )
                    continue

                if isinstance(e, (ToolCall, ToolCallResult)):
                    ctx.write_event_to_stream(e)
                    continue

                if isinstance(e, Event):
                    ctx.write_event_to_stream(e)

            try:
                await handler
            except Exception:
                pass

            return last_answer

        try:
            return await _stream()
        except Exception as ex:
            await self._emit_text(ctx, f"\n`Agent run failed: {ex}`")
            return last_answer

    def _to_text(self, resp: Any) -> str:
        """
        Convert response to text, handling various types and structures.

        :param resp: Response object or content to convert
        :return: str: Text representation of the response
        """
        try:
            if resp is None or str(resp) == "assistant: None":
                return ""
            msg = getattr(resp, "message", None)
            if msg is not None:
                return getattr(msg, "content", "") or ""
            c = getattr(resp, "content", None)
            if c is not None:
                if isinstance(c, list):
                    parts = []
                    for s in c:
                        parts.append(getattr(s, "text", s if isinstance(s, str) else str(s)))
                    return "".join(parts)
                return c if isinstance(c, str) else str(c)
            return str(resp)
        except Exception:
            return str(resp)