"""Chain that does self-ask with search."""

from __future__ import annotations

from collections.abc import Sequence
from typing import TYPE_CHECKING, Any

from langchain_core._api import deprecated
from langchain_core.language_models import BaseLanguageModel
from langchain_core.prompts import BasePromptTemplate
from langchain_core.runnables import Runnable, RunnablePassthrough
from langchain_core.tools import BaseTool, Tool
from pydantic import Field
from typing_extensions import override

from langchain_classic.agents.agent import Agent, AgentExecutor, AgentOutputParser
from langchain_classic.agents.agent_types import AgentType
from langchain_classic.agents.format_scratchpad import format_log_to_str
from langchain_classic.agents.self_ask_with_search.output_parser import (
    SelfAskOutputParser,
)
from langchain_classic.agents.self_ask_with_search.prompt import PROMPT
from langchain_classic.agents.utils import validate_tools_single_input

if TYPE_CHECKING:
    from langchain_community.utilities.google_serper import GoogleSerperAPIWrapper
    from langchain_community.utilities.searchapi import SearchApiAPIWrapper
    from langchain_community.utilities.serpapi import SerpAPIWrapper


@deprecated("0.1.0", alternative="create_self_ask_with_search", removal="1.0")
class SelfAskWithSearchAgent(Agent):
    """Agent for the self-ask-with-search paper."""

    output_parser: AgentOutputParser = Field(default_factory=SelfAskOutputParser)

    @classmethod
    @override
    def _get_default_output_parser(cls, **kwargs: Any) -> AgentOutputParser:
        return SelfAskOutputParser()

    @property
    def _agent_type(self) -> str:
        """Return Identifier of an agent type."""
        return AgentType.SELF_ASK_WITH_SEARCH

    @classmethod
    @override
    def create_prompt(cls, tools: Sequence[BaseTool]) -> BasePromptTemplate:
        """Prompt does not depend on tools."""
        return PROMPT

    @classmethod
    def _validate_tools(cls, tools: Sequence[BaseTool]) -> None:
        validate_tools_single_input(cls.__name__, tools)
        super()._validate_tools(tools)
        if len(tools) != 1:
            msg = f"Exactly one tool must be specified, but got {tools}"
            raise ValueError(msg)
        tool_names = {tool.name for tool in tools}
        if tool_names != {"Intermediate Answer"}:
            msg = f"Tool name should be Intermediate Answer, got {tool_names}"
            raise ValueError(msg)

    @property
    def observation_prefix(self) -> str:
        """Prefix to append the observation with."""
        return "Intermediate answer: "

    @property
    def llm_prefix(self) -> str:
        """Prefix to append the LLM call with."""
        return ""


@deprecated("0.1.0", removal="1.0")
class SelfAskWithSearchChain(AgentExecutor):
    """[Deprecated] Chain that does self-ask with search."""

    def __init__(
        self,
        llm: BaseLanguageModel,
        search_chain: GoogleSerperAPIWrapper | SearchApiAPIWrapper | SerpAPIWrapper,
        **kwargs: Any,
    ):
        """Initialize only with an LLM and a search chain."""
        search_tool = Tool(
            name="Intermediate Answer",
            func=search_chain.run,
            coroutine=search_chain.arun,
            description="Search",
        )
        agent = SelfAskWithSearchAgent.from_llm_and_tools(llm, [search_tool])
        super().__init__(agent=agent, tools=[search_tool], **kwargs)


def create_self_ask_with_search_agent(
    llm: BaseLanguageModel,
    tools: Sequence[BaseTool],
    prompt: BasePromptTemplate,
) -> Runnable:
    """Create an agent that uses self-ask with search prompting.

    Args:
        llm: LLM to use as the agent.
        tools: List of tools. Should just be of length 1, with that tool having
            name `Intermediate Answer`
        prompt: The prompt to use, must have input key `agent_scratchpad` which will
            contain agent actions and tool outputs.

    Returns:
        A Runnable sequence representing an agent. It takes as input all the same input
        variables as the prompt passed in does. It returns as output either an
        AgentAction or AgentFinish.

    Examples:
        ```python
        from langchain_classic import hub
        from langchain_anthropic import ChatAnthropic
        from langchain_classic.agents import (
            AgentExecutor,
            create_self_ask_with_search_agent,
        )

        prompt = hub.pull("hwchase17/self-ask-with-search")
        model = ChatAnthropic(model="claude-3-haiku-20240307")
        tools = [...]  # Should just be one tool with name `Intermediate Answer`

        agent = create_self_ask_with_search_agent(model, tools, prompt)
        agent_executor = AgentExecutor(agent=agent, tools=tools)

        agent_executor.invoke({"input": "hi"})
        ```

    Prompt:

        The prompt must have input key `agent_scratchpad` which will
            contain agent actions and tool outputs as a string.

        Here's an example:

        ```python
        from langchain_core.prompts import PromptTemplate

        template = '''Question: Who lived longer, Muhammad Ali or Alan Turing?
        Are follow up questions needed here: Yes.
        Follow up: How old was Muhammad Ali when he died?
        Intermediate answer: Muhammad Ali was 74 years old when he died.
        Follow up: How old was Alan Turing when he died?
        Intermediate answer: Alan Turing was 41 years old when he died.
        So the final answer is: Muhammad Ali

        Question: When was the founder of craigslist born?
        Are follow up questions needed here: Yes.
        Follow up: Who was the founder of craigslist?
        Intermediate answer: Craigslist was founded by Craig Newmark.
        Follow up: When was Craig Newmark born?
        Intermediate answer: Craig Newmark was born on December 6, 1952.
        So the final answer is: December 6, 1952

        Question: Who was the maternal grandfather of George Washington?
        Are follow up questions needed here: Yes.
        Follow up: Who was the mother of George Washington?
        Intermediate answer: The mother of George Washington was Mary Ball Washington.
        Follow up: Who was the father of Mary Ball Washington?
        Intermediate answer: The father of Mary Ball Washington was Joseph Ball.
        So the final answer is: Joseph Ball

        Question: Are both the directors of Jaws and Casino Royale from the same country?
        Are follow up questions needed here: Yes.
        Follow up: Who is the director of Jaws?
        Intermediate answer: The director of Jaws is Steven Spielberg.
        Follow up: Where is Steven Spielberg from?
        Intermediate answer: The United States.
        Follow up: Who is the director of Casino Royale?
        Intermediate answer: The director of Casino Royale is Martin Campbell.
        Follow up: Where is Martin Campbell from?
        Intermediate answer: New Zealand.
        So the final answer is: No

        Question: {input}
        Are followup questions needed here:{agent_scratchpad}'''

        prompt = PromptTemplate.from_template(template)
        ```
    """  # noqa: E501
    missing_vars = {"agent_scratchpad"}.difference(
        prompt.input_variables + list(prompt.partial_variables),
    )
    if missing_vars:
        msg = f"Prompt missing required variables: {missing_vars}"
        raise ValueError(msg)

    if len(tools) != 1:
        msg = "This agent expects exactly one tool"
        raise ValueError(msg)
    tool = next(iter(tools))
    if tool.name != "Intermediate Answer":
        msg = "This agent expects the tool to be named `Intermediate Answer`"
        raise ValueError(msg)

    llm_with_stop = llm.bind(stop=["\nIntermediate answer:"])
    return (
        RunnablePassthrough.assign(
            agent_scratchpad=lambda x: format_log_to_str(
                x["intermediate_steps"],
                observation_prefix="\nIntermediate answer: ",
                llm_prefix="",
            ),
            # Give it a default
            chat_history=lambda x: x.get("chat_history", ""),
        )
        | prompt
        | llm_with_stop
        | SelfAskOutputParser()
    )
