from langchain.agents import create_tool_calling_agent
from langchain.agents import AgentExecutor
import os
from langchain_openai import ChatOpenAI
from langchain.agents import  Tool
from serp import GoogleSerperAPIWrapper, get_youtube_url
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.messages import HumanMessage, AIMessage
import base64
from PIL import Image
import io

def encode_image(image_path):
    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode('utf-8')

os.environ["SERPER_API_KEY"] = '2a'
os.environ['OPENAI_API_KEY'] = "sk-"

llm = ChatOpenAI(temperature=0, model_name='gpt-4o', openai_api_key=os.environ['OPENAI_API_KEY'])
search_web = GoogleSerperAPIWrapper()
search_images = GoogleSerperAPIWrapper(type="images")
tools = [
    Tool(
        name="web_search",
        func=search_web.run,
        description="useful when you need to extract from the internet a list of websites snippet and a **valid URL**"
    ),
    Tool(
        name="image_search",
        func=search_images.run,
        description="useful when you need to extract from the internet a list of images with the their titles and a **valid URL**"
    ),
    Tool(
        name="video_search",
        func=get_youtube_url,
        description="useful when you need to extract from the internet a list of videos. The output is a list with a **valid URL**"
    ),
]

# prompt = ChatPromptTemplate.from_messages([
#     self.system_prompt,
#     self.source_prompt,
#     self.generate_eval_message(url)])

agent_prompt = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            "You are a helpful assistant that can provide informative urls from the web for any request.Review the attached image and collect from the internet resources which are related and helpful to continue writing the document. This includes website, images and videos Group the resources you collect by type and subject. You should collect at least 2 images and 2 websites and 1 video."
        ),
        ("human", "{input}"),
        ("placeholder", "{agent_scratchpad}"),
    ]
)

agent = create_tool_calling_agent(llm, tools, agent_prompt)

agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)

import gradio as gr
import os
from openai import OpenAI

with gr.Blocks() as demo:
    with gr.Row():
        image = gr.Image(label="image", height=600)
        chatbot = gr.Chatbot()

    #prompt = gr.Textbox(label="prompt")
    button = gr.Button()
    serper_api = gr.Textbox(label="Serper API key")
    openai_key = gr.Textbox(label="OpenAI API key")


    def respond(chat_history, image):
        out = agent_executor.invoke({'input': ''})

        chat_history.append(('', out['output']))
        return chat_history


    def update_serper_api(serper_api):
        print(os.environ['OPENAI_API_KEY'])
        print(serper_api)
        os.environ["SERPER_API_KEY"] = serper_api
        search_web = GoogleSerperAPIWrapper()
        search_images = GoogleSerperAPIWrapper(type="images")
        global tools
        tools = [
            Tool(
                name="web_search",
                func=search_web.run,
                description="useful when you need to extract from the internet a list of websites snippet and a **valid URL**"
            ),
            Tool(
                name="image_search",
                func=search_images.run,
                description="useful when you need to extract from the internet a list of images with the their titles and a **valid URL**"
            ),
            Tool(
                name="video_search",
                func=get_youtube_url,
                description="useful when you need to extract from the internet a list of videos. The output is a list with a **valid URL**"
            ),
        ]
        agent = create_tool_calling_agent(llm, tools, agent_prompt)
        global agent_executor
        agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)


    def update_agent(openai_key):
        os.environ['OPENAI_API_KEY'] = openai_key
        print(os.environ['OPENAI_API_KEY'])
        global llm
        llm = ChatOpenAI(temperature=0, model_name='gpt-4o', openai_api_key=os.environ['OPENAI_API_KEY'])
        agent = create_tool_calling_agent(llm, tools, agent_prompt)
        global agent_executor
        agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)

    def change_image(image):
        image_pil = Image.fromarray(image)

        # Save the image to a bytes buffer
        buffer = io.BytesIO()
        image_pil.save(buffer, format="PNG")  # You can also use "JPEG" if needed

        # Get the byte data from the buffer and encode it to base64
        image_bytes = buffer.getvalue()
        image_base64 = base64.b64encode(image_bytes).decode('utf-8')
        message_content = [{"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,"
                                                                                      f"{image_base64}"}}]
        image_message = HumanMessage(content=message_content)
        global agent_prompt
        agent_prompt = ChatPromptTemplate.from_messages(
            [
                (
                    "system",
                    "You are a helpful assistant that can provide informative urls from the web for any request.Review the attached image and collect from the internet resources which are related and helpful to continue writing the document. This includes website, images Group the resources you collect by type and subject. You should collect at least 3 images and 3 websites."
                ),
                image_message,
                ("human", "{input}"),
                ("placeholder", "{agent_scratchpad}"),
            ]
        )



        agent = create_tool_calling_agent(llm, tools, agent_prompt)
        global agent_executor
        agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)

    button.click(respond, [chatbot, image], [chatbot])
    openai_key.submit(update_agent, [openai_key], [])
    serper_api.submit(update_serper_api, [serper_api], [])
    image.change(change_image,[image],[])
demo.queue().launch(share=True)