import io
import time
from typing import List, Dict, Callable, Any
from langchain.memory import ConversationBufferMemory
from langchain.schema import (
    HumanMessage,
    SystemMessage,
)
from langchain.agents import initialize_agent
from langchain.agents import AgentType
from langchain.chat_models import ChatOpenAI
from langchain_community.tools.tavily_search import TavilySearchResults
import os
from langchain_core.tools import BaseTool
from langchain.tools import DuckDuckGoSearchResults
from langchain_community.tools.arxiv.tool import ArxivQueryRun
from serpapi import GoogleSearch
from langchain.agents import initialize_agent, Tool
import replicate
import base64
from openai import OpenAI
from PIL import Image
import gradio as gr
from css import *

os.environ["REPLICATE_API_TOKEN"] = "r8_IYJpjwjrxegcUfBeBbyUxErJXXsnHDM4AlSQQ"
os.environ["OPENAI_API_KEY"] = "sb-6a683cb3bd63a9b72040aa2dd08feff8b68f08a0e1d959f5"
os.environ['OPENAI_BASE_URL'] = "https://api.openai-sb.com/v1/"
os.environ["SERPAPI_API_KEY"] = "dcc98b22d5f7d413979a175ff7d75b721c5992a3ee1e2363020b2bbdf4f82404"  # you
# os.environ['TAVILY_API_KEY'] = "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" #you
os.environ['TAVILY_API_KEY'] = "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"  # zeng
# image_path = "C:/Users/Lenovo/Desktop/demo-repository-master/image/img_2.png"
model = ChatOpenAI(model_name="gpt-4", temperature=0.6)


def search_baidu(query) -> str:
    params = {
        "engine": "baidu_news",
        "q": query,
        "ct": "1",
        "api_key": "dcc98b22d5f7d413979a175ff7d75b721c5992a3ee1e2363020b2bbdf4f82404"
    }
    search = GoogleSearch(params)
    results = search.get_dict()
    organic_results = results.get("organic_results", [])
    if not organic_results:
        return ""
    final_output = "\n\n".join([
        f"Title: {news.get('title', 'No Title')}\nSnippet: {news.get('snippet', 'No Snippet')}\nDate: {news.get('date', 'No Date')}\nSource:{news.get('source', 'No Source')}"
        for news in organic_results
    ])
    return final_output


def search_bing(query) -> str:
    params = {
        "engine": "bing_news",
        "q": query,
        "ct": "1",
        "api_key": "dcc98b22d5f7d413979a175ff7d75b721c5992a3ee1e2363020b2bbdf4f82404"
    }
    search = GoogleSearch(params)
    results = search.get_dict()
    organic_results = results.get("organic_results", [])
    if not organic_results:
        return ""
    final_output = "\n\n".join([
        f"Title: {news.get('title', 'No Title')}\nSnippet: {news.get('snippet', 'No Snippet')}\nDate: {news.get('date', 'No Date')}\nSource:{news.get('source', 'No Source')}"
        for news in organic_results
    ])
    return final_output


def img_size(image):
    # img = Image.open(image)
    img = Image.fromarray(image.astype("uint8"))
    width, height = img.size
    while width >= 500 or height >= 400:
        width = width * 0.8
        height = height * 0.8
    width = int(width)
    height = int(height)
    resized_img = img.resize((width, height))
    # resized_img.save(image_path)
    return resized_img


def search_image(query) -> str:
    params = {
        "engine": "google_images",
        "q": query,
        "gl": "cn",
    }
    search = GoogleSearch(params)
    results = search.get_dict()
    thumbnails = [search['thumbnail'] for search in results['suggested_searches']][:3]
    information = ""
    client = OpenAI()
    for idx, thumbnail in enumerate(thumbnails):
        response = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": "What’s in this image?Give a brief answer"},
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": thumbnail,
                            },
                        },
                    ],
                }
            ],
            max_tokens=100,
        )
        response_content = response.choices[0].message.content
        information = information + str(idx + 1) + ": " + response_content + "\n"
    return (information)


search_baidu = Tool(
    name="Baidu News Search",  # 工具名称
    func=search_baidu,  # 引用search函数
    description="搜索引擎，当你需要回答当前问题的时候调用，输入是检索query"  # 工具描述
)

search_bing = Tool(
    name="Baidu News Search",  # 工具名称
    func=search_bing,  # 引用search函数
    description="搜索引擎，当你需要回答当前问题的时候调用，输入是检索query"  # 工具描述
)

search_image = Tool(
    name="Image Search",  # 工具名称
    func=search_image,  # 引用search函数
    description="搜索图片引擎，当你需要检索相关图片的时候调用，输入是检索query，输出是与query有关的图片的信息"  # 工具描述
)


class ReplicateModel:
    def __init__(self, model_name: str):
        self.model_name = model_name

    def predict(self, input_data: dict) -> Any:
        # 调用 Replicate API
        output = replicate.run(
            f"{self.model_name}",
            input=input_data
        )
        return output


class DialogueAgent:
    def __init__(
            self,
            name: str,
            system_message: SystemMessage,
            model: model,
    ) -> None:
        self.name = name
        self.system_message = system_message
        self.model = model
        self.prefix = f"{self.name}: "
        self.reset()

    def reset(self):
        self.message_history = ["Here is the conversation so far."]

    def send(self) -> str:
        """
        Applies the chatmodel to the message history
        and returns the message string
        """
        message = self.model(
            [
                self.system_message,
                HumanMessage(content="\n".join(self.message_history + [self.prefix])),
            ]
        )
        return message.content

    def receive(self, name: str, message: str) -> None:
        """
        Concatenates {message} spoken by {name} into message history
        """
        self.message_history.append(f"{name}: {message}")


class Replicate_DialogueAgent:
    def __init__(
            self,
            name: str,
            system_message: SystemMessage,
            model_id: str,  # 修改为接收模型ID
    ) -> None:
        self.name = name
        self.system_message = system_message
        self.model_id = model_id  # 存储模型ID
        self.prefix = f"{self.name}: "
        self.reset()

    def reset(self):
        self.message_history = ["Here is the conversation so far."]

    def send(self) -> str:
        """
        Applies the replicate model to the message history
        and returns the message string
        """
        # 创建输入数据
        input_data = {
            "prompt": "\n".join(self.message_history + [self.prefix]),
        }

        # 调用 replicate 模型
        output = replicate.run(
            self.model_id,
            input=input_data
        )

        return output  # 返回模型的输出

    def receive(self, name: str, message: str) -> None:
        """
        Concatenates {message} spoken by {name} into message history
        """
        self.message_history.append(f"{name}: {message}")


class DialogueSimulator:
    def __init__(
            self,
            agents: List[DialogueAgent],
            selection_function: Callable[[int, List[DialogueAgent]], int],
    ) -> None:
        self.agents = agents
        self._step = 0
        self.select_next_speaker = selection_function

    def reset(self):
        for agent in self.agents:
            agent.reset()

    def inject(self, name: str, message: str):
        """
        Initiates the conversation with a {message} from {name}
        """
        for agent in self.agents:
            agent.receive(name, message)

        # increment time
        self._step += 1

    def step(self) -> tuple[str, str]:
        # 1. choose the next speaker
        speaker_idx = self.select_next_speaker(self._step, self.agents)
        speaker = self.agents[speaker_idx]

        # 2. next speaker sends message
        message = speaker.send()

        # 3. everyone receives message
        for receiver in self.agents:
            receiver.receive(speaker.name, message)

        # 4. increment time
        self._step += 1

        return speaker.name, message


class Replicate_DialogueSimulator:
    def __init__(
            self,
            agents: List[Replicate_DialogueAgent],
            selection_function: Callable[[int, List[Replicate_DialogueAgent]], int],
    ) -> None:
        self.agents = agents
        self._step = 0
        self.select_next_speaker = selection_function

    def reset(self):
        for agent in self.agents:
            agent.reset()

    def inject(self, name: str, message: str):
        """
        Initiates the conversation with a {message} from {name}
        """
        for agent in self.agents:
            agent.receive(name, message)

        # increment time
        self._step += 1

    def step(self) -> tuple[str, str]:
        # 1. choose the next speaker
        speaker_idx = self.select_next_speaker(self._step, self.agents)
        speaker = self.agents[speaker_idx]

        # 2. next speaker sends message
        message = speaker.send()

        # 3. everyone receives message
        for receiver in self.agents:
            receiver.receive(speaker.name, message)

        # 4. increment time
        self._step += 1

        return speaker.name, message


class DialogueAgentWithTools(DialogueAgent):
    def __init__(self, name: str, system_message: SystemMessage, model: ChatOpenAI, tools: List[BaseTool]):
        super().__init__(name, system_message, model)
        self.tools = tools  # 手动传递工具

    def send(self) -> dict[str, Any]:
        agent_chain = initialize_agent(
            self.tools,
            self.model,
            agent=AgentType.CHAT_CONVERSATIONAL_REACT_DESCRIPTION,
            verbose=True,
            memory=ConversationBufferMemory(memory_key="chat_history", return_messages=True),
        )
        message_input = "\n".join([self.system_message.content] + self.message_history + [self.prefix])
        response = agent_chain.invoke({"input": message_input}, handle_parsing_errors=True)

        # response = agent_chain.invoke(
        #     {"input": "\n".join([self.system_message.content] + self.message_history + [self.prefix],handle_parsing_errors=True)}
        # )
        return response


class Repliacte_DialogueAgentWithTools(Replicate_DialogueAgent):
    def __init__(self, name: str, system_message: SystemMessage, replicate_model: ReplicateModel,
                 tools: List[BaseTool]):
        super().__init__(name, system_message, None)  # 不再使用 ChatOpenAI 模型
        self.replicate_model = replicate_model  # 存储 Replicate 模型
        self.tools = tools  # 手动传递工具

    def send(self) -> dict[str, Any]:
        # 准备输入数据
        input_data = {
            "prompt": "\n".join([self.system_message.content] + self.message_history + [self.prefix])
        }
        # 调用 Replicate 模型
        response = self.replicate_model.predict(input_data)
        return response


ddg_search = DuckDuckGoSearchResults()
arxiv_query = ArxivQueryRun()
tavily_tool = TavilySearchResults(max_result=2)

tools = [tavily_tool, search_baidu]


def select_next_speaker(step: int, agents: List[DialogueAgent]) -> int:
    idx = step % len(agents)
    return idx


names = {
    "Evidence-Driven Analyst": ["arxiv", "ddg-search", "wikipedia"],
    "Skeptic Critic": ["arxiv", "ddg-search", "wikipedia"],
}

agent_descriptor_system_message = SystemMessage(
    content="""Evidence Driven Analysts tend to seek positive evidence to prove the authenticity of news, focusing on data, facts, and verifiable sources of information.
Skeptic Critic focuses more on potential issues, false information, and potentially misleading content in news, taking a skeptical attitude towards every detail of the news and emphasizing the search for rebuttal evidence."""
)


def generate_agent_description(name, conversation_description, word_limit):
    agent_specifier_prompt = [
        agent_descriptor_system_message,
        HumanMessage(
            content=f"""{conversation_description}
                        请为{name}提供一个关于新闻真实性判断的描述，限制在{word_limit}个单词内。 
                        直接对{name}说话，并给出他们的观点。 
                        不要添加其他内容。"""
        ),
    ]
    agent_description = ChatOpenAI(temperature=1.0)(agent_specifier_prompt).content
    return agent_description


def generate_system_message_ch(name, description, tools, topic, img_information):
    return f"""Here is the topic of discussion: {topic}
Your name is {name}.

Your description is as follows: {description}

Your goal is to use the provided tools :{tools} and information about news related images:{img_information} to evaluate the truthfulness of the claims in the news article.

Use the following format:
Question: the input information (e.g., claim, post, news) you must assess.

Thought: to evaluate the first detail in the input information, always think about what to do.

Action: the action to take, should be one of [{tools}].

Action Input: the input to the action.

Observation: the result of the action.

Thought: to evaluate the second detail in the input information, always think about what to do.

Action: the action to take, should be one of [{tools}].

Action Input: the input to the action.

Observation: the result of the action.

... (repeat Thought/Action/Action Input/Observation sequence as necessary)

Thought: I now know the final answer.

Final Answer: The definitive response, including reasons on all details and a conclusion, for assessing the original input information.

Additional Notes:
1. If the input information contains multiple statements or details, assess them one by one.
2. The results of tools may also be incorrect; pay attention to conflicts and distinguish them based on your knowledge.
3. The conclusion in the Final Answer should include the detected labels,and you need to give a clear label without any possible label results.
   - For information lacking clear evidence (e.g., rumors), candidate labels are: [Highly Likely Correct, Highly Likely Incorrect, Slightly Incorrect, Slightly Correct, Neutral].
   - For information with clear evidence (e.g., verified by specific proof), candidate labels are: [Completely Incorrect, Mostly Incorrect, Mixed Authenticity, Mostly Correct, Completely Correct].
4. You should respond in Chinese.
5. If you feel the need to search for image information. You can use the tools in {tools} to obtain information

DO look up information with your tools to support your arguments.
DO cite your sources clearly.

DO NOT fabricate fake citations.
DO NOT make assumptions without evidence.

Stop speaking the moment you finish your evaluation.
"""


def generate_system_message_en(name, description, tools, topic, img_information):
    return f"""Here is the topic of discussion: {topic}
Your name is {name}.

Your description is as follows: {description}

Your goal is to use the provided tools :{tools} and information about news related images:{img_information} to evaluate the truthfulness of the claims in the news article.

Use the following format:
Question: the input information (e.g., claim, post, news) you must assess.

Thought: to evaluate the first detail in the input information, always think about what to do.

Action: the action to take, should be one of [{tools}].

Action Input: the input to the action.

Observation: the result of the action.

Thought: to evaluate the second detail in the input information, always think about what to do.

Action: the action to take, should be one of [{tools}].

Action Input: the input to the action.

Observation: the result of the action.

... (repeat Thought/Action/Action Input/Observation sequence as necessary)

Thought: I now know the final answer.

Final Answer: The definitive response, including reasons on all details and a conclusion, for assessing the original input information.

Additional Notes:
1. If the input information contains multiple statements or details, assess them one by one.
2. The results of tools may also be incorrect; pay attention to conflicts and distinguish them based on your knowledge.
3. The conclusion in the Final Answer should include the detected labels,and you need to give a clear label without any possible label results.
   - For information lacking clear evidence (e.g., rumors), candidate labels are: [Highly Likely Correct, Highly Likely Incorrect, Slightly Incorrect, Slightly Correct, Neutral].
   - For information with clear evidence (e.g., verified by specific proof), candidate labels are: [Completely Incorrect, Mostly Incorrect, Mixed Authenticity, Mostly Correct, Completely Correct].
4. You should respond in English.
5. If you feel the need to search for image information. You can use the tools in {tools} to obtain information

DO look up information with your tools to support your arguments.
DO cite your sources clearly.

DO NOT fabricate fake citations.
DO NOT make assumptions without evidence.

Stop speaking the moment you finish your evaluation.
"""


def encode_image(image):
    image = Image.fromarray(image.astype("uint8"))
    buffered = io.BytesIO()
    image.save(buffered, format="PNG")
    return base64.b64encode(buffered.getvalue()).decode("utf-8")
    # with open(image_path, "rb") as image_file:
    #     return base64.b64encode(image_file.read()).decode("utf-8")


def image_summarize(img_base64, prompt):
    chat = ChatOpenAI(model="gpt-4o", max_tokens=256)
    msg = chat.invoke(
        [
            HumanMessage(
                content=[
                    {"type": "text", "text": prompt},
                    {
                        "type": "image_url",
                        "image_url": {"url": f"data:image/jpeg;base64,{img_base64}"},
                    },
                ]
            )
        ]
    )
    return msg.content


def generate_img_summaries(image):
    # img_size(path)
    image_summaries = []
    prompt = """You are an assistant responsible for compiling images for retrieval\
        These abstracts will be embedded and used to retrieve the original images\
        Provide a detailed summary of the optimized images for retrieval."""
    base64_image = encode_image(image)
    image_summaries.append(image_summarize(base64_image, prompt))
    return image_summaries


selected_language = "ch"
selected_model = "gpt"


def SelectLanguage(option):
    global selected_language
    if option == "英文":
        selected_language = "en"
    else:
        selected_language = "ch"


language_option = ["中文", "英文"]


def SelectModel(option):
    global selected_model
    if option == "自研":
        selected_model = "gpt"
    elif option == "gpt-4o":
        selected_model = "gpt"
    elif option == "llama3-70b":
        selected_model = "llama3-70b"
    elif option == "llama3-8b":
        selected_model = "llama3-8b"
    else:
        selected_model = "mistral"


model_option = ["自研", "gpt-4o","llama3-70b", "llama3-8b", "mistral"]


def start(topic, image_summaries):
    max_iters = 3
    n = 0
    word_limit = 50
    result = ""
    global selected_model
    global selected_language
    # image_summaries = generate_img_summaries(image)
    # print(image_summaries)
    conversation_description = f"""Here is the news of conversation: {topic}
    The participants are: {', '.join(names.keys())}"""

    agent_descriptions = {name: generate_agent_description(name, conversation_description, word_limit) for name in
                          names}
    if selected_language == "ch":
        agent_system_messages = {
            name: generate_system_message_ch(name, description, tools, topic, image_summaries)
            for (name, tools), description in zip(names.items(), agent_descriptions.values())
        }
    else:
        agent_system_messages = {
            name: generate_system_message_en(name, description, tools, topic, image_summaries)
            for (name, tools), description in zip(names.items(), agent_descriptions.values())
        }

    topic_specifier_prompt = [
        SystemMessage(content="你可以使主题更具体。"),
        HumanMessage(
            content=f"""{topic}
            你是主持人。
            请具体化这个新闻真实性的主题。
            请以{word_limit}个单词或更少的字数回复指定的问题。 
            这是与新闻有关的图片的信息：{image_summaries}。
            直接对参与者说话：{*names,}。
            不要添加其他内容。"""
        ),
    ]
    specified_topic = ChatOpenAI(temperature=1.0)(topic_specifier_prompt).content
    if selected_model == "gpt":
        agents = [
            DialogueAgentWithTools(
                name=name,
                system_message=SystemMessage(content=system_message),
                model=ChatOpenAI(model="gpt-4o", temperature=0.6),
                tools=[tavily_tool, search_baidu, search_image]
            )
            for (name, tools), system_message in zip(
                names.items(), agent_system_messages.values()
            )
        ]
        simulator = DialogueSimulator(agents=agents, selection_function=select_next_speaker)
        simulator.reset()
        simulator.inject("Moderator", specified_topic)
        while n < max_iters:
            name, message = simulator.step()
            result = result + "(" + name + "): " + message['output'] + "\n"
            # print(f"({name}): {message['output']}")
            # print("\n")
            n += 1

    else:
        if selected_model == "llama3-70b":
            replicate_model = ReplicateModel(model_name="meta/meta-llama-3-70b-instruct")
        elif selected_model == "llama3-8b":
            replicate_model = ReplicateModel(model_name="meta/meta-llama-3-8b-instruct")
        elif selected_model == "mistral":
            replicate_model = ReplicateModel(model_name="mistralai/mixtral-8x7b-instruct-v0.1")
        agents = [
            Repliacte_DialogueAgentWithTools(
                name=name,
                system_message=SystemMessage(content=system_message),
                # model=ChatOpenAI(model="gpt-4o", temperature=0.6),
                replicate_model=replicate_model,
                tools=[tavily_tool, search_baidu]
            )
            for (name, tools), system_message in zip(
                names.items(), agent_system_messages.values()
            )
        ]
        simulator = Replicate_DialogueSimulator(agents=agents, selection_function=select_next_speaker)
        simulator.reset()
        simulator.inject("Moderator", specified_topic)
        while n < max_iters:
            name, message = simulator.step()
            formatted_output = ''.join(message)
            result = result + "(" + name + "): " + formatted_output + "\n"
            # print(f"({name}): {formatted_output}")
            # print("\n")
            n += 1
    return result


# start("9月18日，Trump 在纽约举行第二次暗杀未遂事件后首场竞选集会，现场共有1.8万名支持者参加。", "gpt")  #topic是新闻，model_select是要选择的模型（gpt，llama3-70b,llama3-8b,mistral）

title = "# 虚假信息检测"

with gr.Blocks(css=css1) as demo:
    gr.Markdown(title, elem_id="title")
    with gr.Row():
        with gr.Column(scale=4):
            chatbot = gr.Chatbot(elem_classes="gradio-output")
            # with gr.Row():
            #     with gr.Column(scale=1):
            #         language_select = gr.Dropdown(choices=language_option, elem_classes="gradio-input",
            #                               label="请选择要使用的语言")
            #     with gr.Column(scale=1):
            #         model_select = gr.Dropdown(choices=model_option, elem_classes="gradio-input", label="请选择要使用的大模型")
            # input_box = gr.Textbox(label="输入", elem_classes="gradio-input", placeholder="请输入要判断的新闻", lines=3)
            img_info = gr.Textbox(label="提取到的信息", lines=5, elem_classes="gradio-output")

        with gr.Column(scale=2):
            img_input = gr.Image(label="上传图像", type="numpy")
            img_output = gr.Image(label="处理后的图像", type="numpy", visible=False)
            input_box = gr.Textbox(label="输入", elem_classes="gradio-input", placeholder="请输入要判断的新闻", lines=3)
            # img_info = gr.Textbox(label="提取到的信息", lines=5, elem_classes="gradio-output")
            ans_box = gr.Textbox(label="gpt-4o", lines=5, elem_classes="gradio-output", visible=False)
            dialogue_box = gr.Textbox(label="React", lines=5, elem_classes="gradio-output", visible=False)
            with gr.Row():
                # with gr.Column(scale=1):
                language_select = gr.Dropdown(choices=language_option, elem_classes="gradio-input",
                                          label="请选择要使用的语言",scale=1)
                # with gr.Column(scale=1):
                model_select = gr.Dropdown(choices=model_option, elem_classes="gradio-input", label="请选择要使用的大模型",scale=1)
            with gr.Row():
                # with gr.Column(scale=1):
                clear = gr.Button("清空页面", elem_classes="gradio-button", scale=1)
                # with gr.Column(scale=1):
                submit_btn = gr.Button("提交", elem_classes="gradio-button", scale=1)
        language_select.change(SelectLanguage, language_select)
        model_select.change(SelectModel, model_select)


    def user(user_input, history):
        if history is None:
            history = []
        return user_input, history + [[user_input, None]]


    def bot(history, rag_box):
        if history is None or len(history) == 0:
            return
        bot_rag = str(rag_box)
        history[-1][1] = ""
        for character in bot_rag:
            history[-1][1] += character
            time.sleep(0.01)
            yield history


    submit_btn.click(img_size, img_input, img_output
                     ).then(
        generate_img_summaries, img_output, img_info
    ).then(
        start, [input_box, img_info], dialogue_box
    ).then(
        user, [input_box, chatbot], [input_box, chatbot]
    ).then(
        bot, [chatbot, dialogue_box], chatbot
    )

    clear.click(lambda: (None, None, None, None, None, None), inputs=None,
                outputs=[chatbot, ans_box, dialogue_box, img_input, img_info, img_output])

if __name__ == "__main__":
    demo.launch()