import torch
import os
langchain_install_guide = """pip install --upgrade langchain langchain-community"""
try:
    from langchain_core.agents import AgentAction, AgentFinish
    from typing import List, Optional, Any, Mapping, Union, Dict, Type
    from langchain_core.callbacks import CallbackManagerForLLMRun
    from langchain_core.language_models.llms import BaseLLM
    from langchain_core.outputs import Generation, LLMResult

    from langchain_community.llms import HuggingFaceHub
    from langchain_community.llms.huggingface_hub import HuggingFaceHub
    from langchain_community.llms.huggingface_pipeline import HuggingFacePipeline, VALID_TASKS
    from langchain_community.chat_models.huggingface import ChatHuggingFace
    from langchain_core.pydantic_v1 import root_validator

    # react style prompt
    from langchain import hub
    from langchain.agents import AgentExecutor, load_tools
    from langchain.agents.format_scratchpad import format_log_to_str
    from langchain.agents.output_parsers import (
        ReActJsonSingleInputOutputParser 
    )
    from langchain.tools.render import render_text_description, render_text_description_and_args
    from langchain_community.utilities import SerpAPIWrapper
    from langchain_core.prompts import ChatPromptTemplate
    

    from langchain_core.utils.function_calling import (
        convert_to_openai_function,
        convert_to_openai_tool,
    )
    from langchain_core.exceptions import OutputParserException
    from langchain_core.agents import AgentAction, AgentFinish

    from langchain_core.callbacks.manager import (
        AsyncCallbackManagerForLLMRun,
        CallbackManagerForLLMRun,
    )
    from langchain_core.language_models.chat_models import BaseChatModel
    from langchain_core.messages import (
        AIMessage,
        BaseMessage,
        HumanMessage,
        SystemMessage,
    )
    from langchain_core.outputs import ChatGeneration, ChatResult, LLMResult
    from langchain_core.pydantic_v1 import root_validator
    from langchain_core.pydantic_v1 import Extra
    # from langchain_community.tools.tavily_search import TavilySearchResults
    # from langchain_community.tools.tavily_search import (
    #     TavilySearchResults,
    #     TavilySearchAPIWrapper,
    #     Type,
    #     TavilyInput,
    #     CallbackManagerForToolRun,
    #     AsyncCallbackManagerForToolRun,
    # )
    # from langchain_core.pydantic_v1 import BaseModel, Field

    # from langchain_core.callbacks import (
    # AsyncCallbackManagerForToolRun,
    #     CallbackManagerForToolRun,
    # )
    # from langchain_core.pydantic_v1 import BaseModel, Field
    # from langchain_core.tools import BaseTool

    # ===
    from langchain_core.callbacks import (
        AsyncCallbackManagerForToolRun,
        CallbackManagerForToolRun,
    )
    from langchain_core.pydantic_v1 import BaseModel, Field
    from langchain_core.tools import BaseTool

    from langchain_community.utilities.tavily_search import TavilySearchAPIWrapper
    # from langchain_community.tools.tavily_search import (
    #     TavilySearchResults,
    # )

    LANGCHAIN_AVAILABLE = True

except Exception as e:
    print(f'{str(e)}\nNeed to install langchain: `{langchain_install_guide}`')

    LANGCHAIN_AVAILABLE = False


import logging
import importlib

logger = logging.getLogger(__name__)

DEFAULT_SYSTEM_PROMPT = """You are a helpful, respectful, and honest assistant."""


class AnyEnginePipeline(BaseLLM):
    engine: Any  #: :meta private:
    # model_id: str = DEFAULT_MODEL_ID
    """Model name to use."""
    model_kwargs: Optional[dict] = None
    """Keyword arguments passed to the model."""
    pipeline_kwargs: Optional[dict] = None
    """Keyword arguments passed to the pipeline."""
    batch_size: int = 1
    """Batch size to use when passing multiple documents to generate."""
    streaming: bool = False

    class Config:
        """Configuration for this pydantic object."""

        extra = Extra.forbid

    @classmethod
    def from_engine(
        cls,
        engine: Any,
        model_kwargs: Optional[dict] = None,
        **kwargs
    ):
        return cls(engine=engine, model_kwargs=model_kwargs, **kwargs)

    @property
    def _identifying_params(self) -> Mapping[str, Any]:
        """Get the identifying parameters."""
        return {
            # "model_id": self.model_id,
            "model_kwargs": self.model_kwargs,
            # "pipeline_kwargs": self.pipeline_kwargs,
        }

    @property
    def _llm_type(self) -> str:
        return "engine_pipeline"

    def _generate(
        self,
        prompts: List[str],
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> LLMResult:
        # List to hold all results
        text_generations: List[str] = []
        stop_strings = stop
        print(f'Pipeline run: {len(prompts)}')
        for i in range(0, len(prompts), self.batch_size):
            batch_prompts = prompts[i : i + self.batch_size]
            responses = []
            for p in batch_prompts:
                output = self.engine.generate_yield_string_final(p, stop_strings=stop_strings, **kwargs)
                responses.append(output[0])
            for j, (prompt, response) in enumerate(zip(batch_prompts, responses)):
                text = response
                if text.startswith(prompt):
                    text = text[len(prompt):]
                if stop is not None and any(x in text for x in stop):
                    text = text[:text.index(stop[0])]
                text_generations.append(text)
        return LLMResult(
            generations=[[Generation(text=text)] for text in text_generations]
        )


class ChatAnyEnginePipeline(BaseChatModel):
    """
    Wrapper for engine
    """
    llm: AnyEnginePipeline
    """LLM, must be of type HuggingFaceTextGenInference, HuggingFaceEndpoint, or 
        HuggingFaceHub."""
    system_message: SystemMessage = SystemMessage(content=DEFAULT_SYSTEM_PROMPT)
    tokenizer: Any = None
    model_id: Optional[str] = None

    def __init__(self, **kwargs: Any):
        super().__init__(**kwargs)
        self.tokenizer = self.llm.engine.tokenizer
    
    @root_validator()
    def validate_llm(cls, values: dict) -> dict:
        # if not isinstance(
        #     values["llm"],
        #     (HuggingFaceTextGenInference, HuggingFaceEndpoint, HuggingFaceHub),
        # ):
        #     raise TypeError(
        #         "Expected llm to be one of HuggingFaceTextGenInference, "
        #         f"HuggingFaceEndpoint, HuggingFaceHub, received {type(values['llm'])}"
        #     )
        return values
    
    def _generate(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> ChatResult:
        llm_input = self._to_chat_prompt(messages)
        llm_result = self.llm._generate(
            prompts=[llm_input], stop=stop, run_manager=run_manager, **kwargs
        )
        return self._to_chat_result(llm_result)
    
    async def _agenerate(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> ChatResult:
        llm_input = self._to_chat_prompt(messages)
        llm_result = await self.llm._agenerate(
            prompts=[llm_input], stop=stop, run_manager=run_manager, **kwargs
        )
        return self._to_chat_result(llm_result)

    def _to_chat_prompt(
        self,
        messages: List[BaseMessage],
    ) -> str:
        """Convert a list of messages into a prompt format expected by wrapped LLM."""
        if not messages:
            raise ValueError("At least one HumanMessage must be provided!")

        if not isinstance(messages[-1], HumanMessage):
            raise ValueError("Last message must be a HumanMessage!")

        messages_dicts = [self._to_chatml_format(m) for m in messages]

        return self.tokenizer.apply_chat_template(
            messages_dicts, tokenize=False, add_generation_prompt=True
        )

    def _to_chatml_format(self, message: BaseMessage) -> dict:
        """Convert LangChain message to ChatML format."""

        if isinstance(message, SystemMessage):
            role = "system"
        elif isinstance(message, AIMessage):
            role = "assistant"
        elif isinstance(message, HumanMessage):
            role = "user"
        else:
            raise ValueError(f"Unknown message type: {type(message)}")

        return {"role": role, "content": message.content}

    @staticmethod
    def _to_chat_result(llm_result: LLMResult) -> ChatResult:
        chat_generations = []

        for g in llm_result.generations[0]:
            chat_generation = ChatGeneration(
                message=AIMessage(content=g.text), generation_info=g.generation_info
            )
            chat_generations.append(chat_generation)

        return ChatResult(
            generations=chat_generations, llm_output=llm_result.llm_output
        )
    
    def _resolve_model_id(self) -> None:
        self.model_id = "debug"

    @property
    def _llm_type(self) -> str:
        return "engine-chat-wrapper"




class TavilyInput(BaseModel):
    """Input for the Tavily tool."""

    query: str = Field(description="search query to look up")


class NewTavilySearchAPIWrapper(TavilySearchAPIWrapper):
    def clean_results(self, results: List[Dict]) -> List[Dict]:
        """Clean results from Tavily Search API."""
        clean_results = []
        for result in results:
            clean_results.append(
                {
                    "url": result["url"],
                    "content": result.get("raw_content", result["content"]),
                }
            )
        return clean_results


class NewTavilySearchResults(BaseTool):
    """Tool that queries the Tavily Search API and gets back json."""

    name: str = "tavily_search_results_json"
    description: str = (
        "A search engine optimized for comprehensive, accurate, and trusted results. "
        "Useful for when you need to answer questions about current events. "
        "Input should be a search query."
    )
    api_wrapper: NewTavilySearchAPIWrapper = Field(default_factory=NewTavilySearchAPIWrapper)
    max_results: int = 5
    args_schema: Type[BaseModel] = TavilyInput

    def _run(
        self,
        query: str,
        run_manager: Optional[CallbackManagerForToolRun] = None,
    ) -> Union[List[Dict], str]:
        """Use the tool."""
        try:
            return self.api_wrapper.results(
                query,
                self.max_results,
                include_answer=True,
                include_raw_content=True,
            )
        except Exception as e:
            return repr(e)

    async def _arun(
        self,
        query: str,
        run_manager: Optional[AsyncCallbackManagerForToolRun] = None,
    ) -> Union[List[Dict], str]:
        """Use the tool asynchronously."""
        try:
            return await self.api_wrapper.results_async(
                query,
                self.max_results,
                include_answer=True,
                include_raw_content=True,
            )
        except Exception as e:
            return repr(e)


FINAL_ANSWER_ACTION = "Final Answer:"
class LooseReActJsonSingleInputOutputParser(ReActJsonSingleInputOutputParser):
    def parse(self, text: str) -> AgentAction | AgentFinish:
        try:
            return super().parse(text)
        except OutputParserException as e:
            output = text
            if FINAL_ANSWER_ACTION in text:
                output = text.split(FINAL_ANSWER_ACTION)[-1].strip()
            return AgentFinish({"output": output}, text)



web_search_system_prompt = """You are a helpful, intelligent and respectful assistant with access to the Internet via the `tavily_search_results_json` search engine tool. \
You provide answers and responses as accurately as possible to the user queries and questions, using the tools available to you. \
You may use your own knowledge to reply to the user. However, if you are not confident about your knowledge, or you do not have the up-to-date knowledge and abilitiy to answer the questions, please use the search tool to query appropriately.

You understand that you have to craft an informative and search-engine-friendly query given the user's question for the engine to retrieve the most relevant information. \
You also understand that if the question is complex, you may need to reason your thoughts step by step, and may call the search engine multiple times if needed. However, you must use the least API call as possible!
If you have used the search engine, you should include in your final response citations of the website links you have retrieved.

To use the search engine, you must first speak out your thought, then follow by an action as a json blob and understand the observation, and produce the final answer. ALWAYS use the following format:

Question: the input user question you must answer
Thought: you should always think about what to do in the first step
Action:
```
{{
    "action": "tavily_search_results_json",
    "action_input": {{
        "query": "search query 1"
    }}
}}
```
Observation: the result of the search query 1 you just performed
Thought: you continue to think about what to query next, if necessary
Action:
```
{{
    "action": "tavily_search_results_json",
    "action_input": {{
        "query": "search query 1"
    }}
}}
```
Observation: the result of the search query 2 you just performed
... (this Thought/Action/Observation can repeat N times)
Thought: I now know the final answer
Final answer: the final answer to the original user's input question
Citation: ...

You are provided the following concrete examples, please study them and understand your task.

### Example 1

Question: Who is the wife of the current US president?
Thought: This question is twofold and a single search query may not suffice. First I need to find out who is the current US president, then I need to find out who his wife is.
Action:
```
{{
    "action": "tavily_search_results_json",
    "action_input": {{
        "query": "Current US president"
    }}
}}
```
Observation: [{{'url': 'https://en.wikipedia.org/wiki/Joe_Biden', 'content': 'Joe Biden is the current US president. He is the 46th US president.'}}]
Thought: Now I need to find out who is the wife of Joe Biden
Action:
```
{{
    "action": "tavily_search_results_json",
    "action_input": {{
        "query": "Who is the wife of Joe Biden?"
    }}
}}
```
Observation: [{{'url': 'https://en.wikipedia.org/wiki/Jill_Biden', 'content': 'The wife of Joe Biden is Jill Biden, who is an American educator.'}}]
Thought: I now know the final answer
Final answer: The wife of the current US president is Jill Biden.
Citation: 
* https://en.wikipedia.org/wiki/Joe_Biden
* https://en.wikipedia.org/wiki/Jill_Biden

### Example 2

Question: What is langchain?
Thought: I think I should query the internet to understand what is langchain
Action:
```
{{
    "action": "tavily_search_results_json",
    "action_input": {{
        "query": "what is langchain?"
    }}
}}
```
Observation: [{{'url': 'https://python.langchain.com/docs/get_started/introduction/', 'content': 'LangChain is a framework for developing applications powered by large language models (LLMs).'}}]
Thought: I now know the final answer
Final answer: From my search query, Langchain is a framework for building applications using Large Language Models or LLMs.
Citation: 
* https://python.langchain.com/docs/get_started/introduction/


Let's begin! Below is the question from the user.
"""
# FINAL REMARKS: The user may not speak English and may ask you questions in any language. Thus, while your Thought, Action and Observation is in English, your `Final answer` should be in the same language as the user's query.


"""


"""



def create_web_search_engine(model_engine=None):
    # from langchain_community.tools.tavily_search import TavilySearchResults
    if model_engine is None:
        raise ValueError(f'model_engine empty')

    from langchain_core.utils.function_calling import (
        convert_to_openai_function,
        convert_to_openai_tool,
    )
    from langchain_core.exceptions import OutputParserException
    from langchain_core.agents import AgentAction, AgentFinish
    web_search_llm = AnyEnginePipeline.from_engine(model_engine)
    web_search_chat_model = ChatAnyEnginePipeline(llm=web_search_llm)
    if "TAVILY_API_KEY" not in os.environ:
        raise ValueError(f'TAVILY_API_KEY is not found to use websearch, please `export TAVILY_API_KEY=YOUR_TAVILY_API_KEY`')

    tools = [NewTavilySearchResults(max_results=1)]
    formatted_tools = [convert_to_openai_tool(tool) for tool in tools]
    prompt_template = ChatPromptTemplate.from_messages(
        [
            # (
            #     "system",
            #     web_search_system_prompt,
            # ),
            (
                "human",
                web_search_system_prompt + "\n{input}\n{agent_scratchpad}"
                # "{input}\n\n{agent_scratchpad}"
            )
        ]
    )
    prompt = prompt_template.partial(
        tools=formatted_tools,
        tool_names=", ".join([t.name for t in tools]),
    )
    chat_model_with_stop = web_search_chat_model.bind(stop=["\nObservation"])
    agent = (
        {
            "input": lambda x: x["input"],
            "agent_scratchpad": lambda x: format_log_to_str(x["intermediate_steps"]),
        }
        | prompt
        | chat_model_with_stop
        | LooseReActJsonSingleInputOutputParser()
    )
        # | ReActJsonSingleInputOutputParser()

    # instantiate AgentExecutor
    agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
    # agent_executor.invoke({"input": "What is langchain?"})
    return web_search_llm, web_search_chat_model, agent_executor