import io
import os
import shutil
import base64
import gradio as gr
from PIL import Image, ImageDraw

from MobileAgent.text_localization import ocr
from MobileAgent.icon_localization import det
from MobileAgent.local_server import mobile_agent_infer

from modelscope import snapshot_download
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks


chatbot_css = """
<style>
.chat-container {
    display: flex;
    flex-direction: column;
    overflow-y: auto;
    max-height: 630px;
    margin: 10px;
}
.user-message, .bot-message {
    margin: 5px;
    padding: 10px;
    border-radius: 10px;
}
.user-message {
    text-align: right;
    background-color: #7B68EE;
    color: white;
    align-self: flex-end;
}
.bot-message {
    text-align: left;
    background-color: #ADD8E6;
    color: black;
    align-self: flex-start;
}
.user-image {
    text-align: right;
    align-self: flex-end;
    max-width: 150px;
    max-height: 300px;
}
.bot-image {
    text-align: left;
    align-self: flex-start;
    max-width: 200px;
    max-height: 400px;
}
</style>
"""


temp_file = "temp"
screenshot = "screenshot"
cache = "cache"
if not os.path.exists(temp_file):
    os.mkdir(temp_file)
if not os.path.exists(screenshot):
    os.mkdir(screenshot)
if not os.path.exists(cache):
    os.mkdir(cache)


groundingdino_dir = snapshot_download('AI-ModelScope/GroundingDINO', revision='v1.0.0')
groundingdino_model = pipeline('grounding-dino-task', model=groundingdino_dir)
ocr_detection = pipeline(Tasks.ocr_detection, model='damo/cv_resnet18_ocr-detection-line-level_damo')
ocr_recognition = pipeline(Tasks.ocr_recognition, model='damo/cv_convnextTiny_ocr-recognition-document_damo')


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


def get_all_files_in_folder(folder_path):
    file_list = []
    for file_name in os.listdir(folder_path):
        file_list.append(file_name)
    return file_list


def crop(image, box, i):
    image = Image.open(image)
    x1, y1, x2, y2 = int(box[0]), int(box[1]), int(box[2]), int(box[3])
    if x1 >= x2-10 or y1 >= y2-10:
        return
    cropped_image = image.crop((x1, y1, x2, y2))
    cropped_image.save(f"./temp/{i}.png", format="PNG")


def merge_text_blocks(text_list, coordinates_list):
    merged_text_blocks = []
    merged_coordinates = []

    sorted_indices = sorted(range(len(coordinates_list)), key=lambda k: (coordinates_list[k][1], coordinates_list[k][0]))
    sorted_text_list = [text_list[i] for i in sorted_indices]
    sorted_coordinates_list = [coordinates_list[i] for i in sorted_indices]

    num_blocks = len(sorted_text_list)
    merge = [False] * num_blocks

    for i in range(num_blocks):
        if merge[i]:
            continue
        
        anchor = i
        
        group_text = [sorted_text_list[anchor]]
        group_coordinates = [sorted_coordinates_list[anchor]]

        for j in range(i+1, num_blocks):
            if merge[j]:
                continue

            if abs(sorted_coordinates_list[anchor][0] - sorted_coordinates_list[j][0]) < 10 and \
            sorted_coordinates_list[j][1] - sorted_coordinates_list[anchor][3] >= -10 and sorted_coordinates_list[j][1] - sorted_coordinates_list[anchor][3] < 30 and \
            abs(sorted_coordinates_list[anchor][3] - sorted_coordinates_list[anchor][1] - (sorted_coordinates_list[j][3] - sorted_coordinates_list[j][1])) < 10:
                group_text.append(sorted_text_list[j])
                group_coordinates.append(sorted_coordinates_list[j])
                merge[anchor] = True
                anchor = j
                merge[anchor] = True

        merged_text = "\n".join(group_text)
        min_x1 = min(group_coordinates, key=lambda x: x[0])[0]
        min_y1 = min(group_coordinates, key=lambda x: x[1])[1]
        max_x2 = max(group_coordinates, key=lambda x: x[2])[2]
        max_y2 = max(group_coordinates, key=lambda x: x[3])[3]

        merged_text_blocks.append(merged_text)
        merged_coordinates.append([min_x1, min_y1, max_x2, max_y2])

    return merged_text_blocks, merged_coordinates


def get_perception_infos(screenshot_file):
    width, height = Image.open(screenshot_file).size
    
    text, coordinates = ocr(screenshot_file, ocr_detection, ocr_recognition)
    text, coordinates = merge_text_blocks(text, coordinates)
    
    perception_infos = []
    for i in range(len(coordinates)):
        perception_info = {"text": "text: " + text[i], "coordinates": coordinates[i]}
        perception_infos.append(perception_info)
        
    coordinates = det(screenshot_file, "icon", groundingdino_model)
    
    for i in range(len(coordinates)):
        perception_info = {"text": "icon", "coordinates": coordinates[i]}
        perception_infos.append(perception_info)
        
    image_box = []
    image_id = []
    for i in range(len(perception_infos)):
        if perception_infos[i]['text'] == 'icon':
            image_box.append(perception_infos[i]['coordinates'])
            image_id.append(i)

    for i in range(len(image_box)):
        crop(screenshot_file, image_box[i], image_id[i])

    images = get_all_files_in_folder(temp_file)
    if len(images) > 0:
        images = sorted(images, key=lambda x: int(x.split('/')[-1].split('.')[0]))
        image_id = [int(image.split('/')[-1].split('.')[0]) for image in images]
        icon_map = {}
        prompt = 'This image is an icon from a phone screen. Please briefly describe the shape and color of this icon in one sentence.'

        string_image = []
        for i in range(len(images)):
            image_path = os.path.join(temp_file, images[i])
            string_image.append({"image_name": images[i], "image_file": encode_image(image_path)})
        query_data = {"task": "caption", "images": string_image, "query": prompt}
        response_query = mobile_agent_infer(query_data)
        icon_map = response_query["icon_map"]

        for i, j in zip(image_id, range(1, len(image_id)+1)):
            if icon_map.get(str(j)):
                perception_infos[i]['text'] = "icon: " + icon_map[str(j)]

    for i in range(len(perception_infos)):
        perception_infos[i]['coordinates'] = [int((perception_infos[i]['coordinates'][0]+perception_infos[i]['coordinates'][2])/2), int((perception_infos[i]['coordinates'][1]+perception_infos[i]['coordinates'][3])/2)]
        
    return perception_infos, width, height


def image_to_base64(image):
    buffered = io.BytesIO()
    image.save(buffered, format="PNG")
    img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
    img_html = f'<img src="data:image/png;base64,{img_str}" />'
    return img_html


def chatbot(image, instruction, add_info, history, chat_log):
    if history == {}:
        thought_history = []
        summary_history = []
        action_history = []
        summary = ""
        action = ""
        completed_requirements = ""
        memory = ""
        insight = ""
        error_flag = False
        user_msg = "<div class='user-message'>{}</div>".format(instruction)
    else:
        thought_history = history["thought_history"]
        summary_history = history["summary_history"]
        action_history = history["action_history"]
        summary = history["summary"]
        action = history["action"]
        completed_requirements = history["completed_requirements"]
        memory = history["memory"][0]
        insight = history["insight"]
        error_flag = history["error_flag"]
        user_msg = "<div class='user-message'>{}</div>".format("I have uploaded the screenshot. Please continue operating.")
        
    images = get_all_files_in_folder(cache)
    if len(images) > 0 and len(images) <= 100:
        images = sorted(images, key=lambda x: int(x.split('/')[-1].split('.')[0]))
        image_id = [int(image.split('/')[-1].split('.')[0]) for image in images]
        cur_image_id = image_id[-1] + 1
    elif len(images) > 100:
        images = sorted(images, key=lambda x: int(x.split('/')[-1].split('.')[0]))
        image_id = [int(image.split('/')[-1].split('.')[0]) for image in images]
        cur_image_id = image_id[-1] + 1
        os.remove(os.path.join(cache, str(image_id[0])+".png"))
    else:
        cur_image_id = 1
    
    image.save(os.path.join(cache, str(cur_image_id) + ".png"), format="PNG")
    screenshot_file = os.path.join(cache, str(cur_image_id) + ".png")
    perception_infos, width, height = get_perception_infos(screenshot_file)
    shutil.rmtree(temp_file)
    os.mkdir(temp_file)
    
    local_screenshot_file = encode_image(screenshot_file)
    query_data = {
        "task": "decision",
        "screenshot_file": local_screenshot_file,
        "instruction": instruction,
        "perception_infos": perception_infos,
        "width": width,
        "height": height,
        "summary_history": summary_history,
        "action_history": action_history,
        "summary": summary,
        "action": action,
        "add_info": add_info,
        "error_flag": error_flag,
        "completed_requirements": completed_requirements,
        "memory": memory,
        "memory_switch": True,
        "insight": insight
    }

    response_query = mobile_agent_infer(query_data)
    output_action = response_query["decision"]
    output_memory = response_query["memory"]
    if output_action == "No token":
        bot_response = ["<div class='bot-message'>{}</div>".format("Sorry, the resources can be exhausted today.")]
        chat_html = "<div class='chat-container'>{}</div>".format("".join(bot_response))
        return chatbot_css + chat_html, history, chat_log
    
    thought = output_action.split("### Thought ###")[-1].split("### Action ###")[0].replace("\n", " ").replace(":", "").replace("  ", " ").strip()
    summary = output_action.split("### Operation ###")[-1].replace("\n", " ").replace("  ", " ").strip()
    action = output_action.split("### Action ###")[-1].split("### Operation ###")[0].replace("\n", " ").replace("  ", " ").strip()

    output_memory = output_memory.split("### Important content ###")[-1].split("\n\n")[0].strip() + "\n"
    if "None" not in output_memory and output_memory not in memory:
        memory += output_memory
    
    if "Open app" in action:
        bot_response = "Please click the red circle and upload the current screenshot again."
        app_name = action.split("(")[-1].split(")")[0]
        text, coordinate = ocr(screenshot_file, ocr_detection, ocr_recognition)
        for ti in range(len(text)):
            if app_name == text[ti]:
                name_coordinate = [int((coordinate[ti][0] + coordinate[ti][2])/2), int((coordinate[ti][1] + coordinate[ti][3])/2)]
                x, y = name_coordinate[0], name_coordinate[1]
                radius = 75
                draw = ImageDraw.Draw(image)
                draw.ellipse([x - radius, y - radius, x + radius, y + radius], outline='red', width=10)
                break
        
    elif "Tap" in action:
        bot_response = "Please click the red circle and upload the current screenshot again."
        coordinate = action.split("(")[-1].split(")")[0].split(", ")
        x, y = int(coordinate[0]), int(coordinate[1])
        radius = 75
        draw = ImageDraw.Draw(image)
        draw.ellipse([x - radius, y - radius, x + radius, y + radius], outline='red', width=10)

    elif "Swipe" in action:
        bot_response = "Please slide from red circle to blue circle and upload the current screenshot again."
        coordinate1 = action.split("Swipe (")[-1].split("), (")[0].split(", ")
        coordinate2 = action.split("), (")[-1].split(")")[0].split(", ")
        x1, y1 = int(coordinate1[0]), int(coordinate1[1])
        x2, y2 = int(coordinate2[0]), int(coordinate2[1])
        radius = 75
        draw = ImageDraw.Draw(image)
        draw.ellipse([x1 - radius, y1 - radius, x1 + radius, y1 + radius], outline='red', width=10)
        draw.ellipse([x2 - radius, y2 - radius, x2 + radius, y2 + radius], outline='blue', width=10)
        
    elif "Type" in action:
        if "(text)" not in action:
            text = action.split("(")[-1].split(")")[0]
        else:
            text = action.split(" \"")[-1].split("\"")[0]
        bot_response = f"Please type the \"{text}\" and upload the current screenshot again."
        
    elif "Back" in action:
        bot_response = f"Please back to previous page and upload the current screenshot again."
        
    elif "Home" in action:
        bot_response = f"Please back to home page and upload the current screenshot again."
        
    elif "Stop" in action:
        bot_response = f"Task completed."
    
    bot_text1 = "<div class='bot-message'>{}</div>".format("### Decision ###")
    bot_thought = "<div class='bot-message'>{}</div>".format("Thought: " + thought)
    bot_action = "<div class='bot-message'>{}</div>".format("Action: " + action)
    bot_operation = "<div class='bot-message'>{}</div>".format("Operation: " + summary)
    bot_text2 = "<div class='bot-message'>{}</div>".format("### Memory ###")
    bot_memory = "<div class='bot-message'>{}</div>".format(output_memory)
    bot_response = "<div class='bot-message'>{}</div>".format(bot_response)
    if image is not None:
        bot_img_html = image_to_base64(image)
        bot_response = "<div class='bot-image'>{}</div>".format(bot_img_html) + bot_response

    chat_log.append(user_msg)
    
    thought_history.append(thought)
    summary_history.append(summary)
    action_history.append(action)
    
    history["thought_history"] = thought_history
    history["summary_history"] = summary_history
    history["action_history"] = action_history
    history["summary"] = summary
    history["action"] = action
    history["memory"] = memory,
    history["memory_switch"] = True,
    history["insight"] = insight
    history["error_flag"] = error_flag
    
    query_data = {
        "task": "planning",
        "instruction": instruction,
        "thought_history": thought_history,
        "summary_history": summary_history,
        "action_history": action_history,
        "completed_requirements": "",
        "add_info": add_info
    }

    response_query = mobile_agent_infer(query_data)
    output_planning = response_query["planning"]
    if output_planning == "No token":
        bot_response = ["<div class='bot-message'>{}</div>".format("Sorry, the resources can be exhausted today.")]
        chat_html = "<div class='chat-container'>{}</div>".format("".join(bot_response))
        return chatbot_css + chat_html, history, chat_log
    
    output_planning = output_planning.split("### Completed contents ###")[-1].replace("\n", " ").strip()
    history["completed_requirements"] = output_planning
    
    bot_text3 = "<div class='bot-message'>{}</div>".format("### Planning ###")
    output_planning = "<div class='bot-message'>{}</div>".format(output_planning)
    
    chat_log.append(bot_text3)
    chat_log.append(output_planning)
    chat_log.append(bot_text1)
    chat_log.append(bot_thought)
    chat_log.append(bot_action)
    chat_log.append(bot_operation)
    chat_log.append(bot_text2)
    chat_log.append(bot_memory)
    chat_log.append(bot_response)

    chat_html = "<div class='chat-container'>{}</div>".format("".join(chat_log))

    return chatbot_css + chat_html, history, chat_log


def lock_input(instruction):
    return gr.update(value=instruction, interactive=False), gr.update(value=None)


def reset_demo():
    return gr.update(value="", interactive=True), gr.update(value="If you want to tap an icon of an app, use the action \"Open app\"", interactive=True), "<div class='chat-container'></div>", {}, []


tos_markdown = ("""<div style="display:flex; gap: 0.25rem;" align="center">
    <a href='https://github.com/X-PLUG/MobileAgent'><img src='https://img.shields.io/badge/Github-Code-blue'></a>
    <a href="https://arxiv.org/abs/2406.01014"><img src="https://img.shields.io/badge/Arxiv-2406.01014-red"></a>
    <a href='https://github.com/X-PLUG/MobileAgent/stargazers'><img src='https://img.shields.io/github/stars/X-PLUG/MobileAgent.svg?style=social'></a>
</div>
If you like our project, please give us a star ✨ on Github for latest update.

**Terms of use**
1. Input your instruction in \"Instruction\", for example \"Turn on the dark mode\".
2. You can input helpful operation knowledge in \"Knowledge\".
3. Click \"Submit\" to get the operation. You need to operate your mobile device according to the operation and then upload the screenshot after your operation.
4. The 5 cases in \"Examples\" are a complete flow. Click and submit from top to bottom to experience.
5. Due to limited resources, each operation may take a long time, please be patient and wait.

**使用说明**
1. 在“Instruction”中输入你的指令，例如“打开深色模式”。
2. 你可以在“Knowledge”中输入帮助性的操作知识。
3. 点击“Submit”来获得操作。你需要根据输出来操作手机，并且上传操作后的截图。
4. “Example”中的5个例子是一个任务。从上到下点击它们并且点击“Submit”来体验。
5. 由于资源有限，每次操作的时间会比较长，请耐心等待。""")

title_markdowm = ("""# Mobile-Agent-v2: Mobile Device Operation Assistant with Effective Navigation via Multi-Agent Collaboration""")

instruction_input = gr.Textbox(label="Instruction", placeholder="Input your instruction")
knowledge_input = gr.Textbox(label="Knowledge", placeholder="Input your knowledge", value="If you want to tap an icon of an app, use the action \"Open app\"")
with gr.Blocks() as demo:
    history_state = gr.State(value={})
    history_output = gr.State(value=[])
    with gr.Row():
        gr.Markdown(title_markdowm)
    with gr.Row():
        with gr.Column(scale=5):
            gr.Markdown(tos_markdown)
            with gr.Row():
                image_input = gr.Image(label="Screenshot", type="pil", height=550, width=230)
                gr.Examples(examples=[
                    ["./example/example_1.jpg", "Turn on the dark mode"],
                    ["./example/example_2.jpg", "Turn on the dark mode"],
                    ["./example/example_3.jpg", "Turn on the dark mode"],
                    ["./example/example_4.jpg", "Turn on the dark mode"],
                    ["./example/example_5.jpg", "Turn on the dark mode"],
                ], inputs=[image_input, instruction_input, knowledge_input])
            
        with gr.Column(scale=6):
            instruction_input.render()
            knowledge_input.render()
            with gr.Row():
                start_button = gr.Button("Submit")
                clear_button = gr.Button("Clear")
            output_component = gr.HTML(label="Chat history", value="<div class='chat-container'></div>")
            
    start_button.click(
        fn=lambda image, instruction, add_info, history, output: chatbot(image, instruction, add_info, history, output),
        inputs=[image_input, instruction_input, knowledge_input, history_state, history_output],
        outputs=[output_component, history_state, history_output]
    )

    clear_button.click(
        fn=reset_demo,
        inputs=[],
        outputs=[instruction_input, knowledge_input, output_component, history_state, history_output]
    )

demo.queue().launch(share=True)