# Copyright (c) Alibaba Cloud.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree

import copy
import re
import gc
import torch
from argparse import ArgumentParser
from threading import Thread
from qwen_vl_utils import process_vision_info
from transformers import AutoProcessor, Qwen2VLForConditionalGeneration, TextIteratorStreamer
import os

from fastapi import FastAPI
from pydantic import BaseModel
from typing import Optional

app = FastAPI()

DEFAULT_CKPT_PATH = '/home/practice/model/Qwen2-VL-72B-Instruct'
REVISION = 'v1.0.4'
BOX_TAG_PATTERN = r"<box>([\s\S]*?)</box>"
PUNCTUATION = "！？。＂＃＄％＆＇（）＊＋，－／：；＜＝＞＠［＼］＾＿｀｛｜｝～｟｠｢｣､、〃》「」『』【】〔〕〖〗〘〙〚〛〜〝〞〟〰〾〿–—‘’‛“”„‟…‧﹏."

def get_args():
    parser = ArgumentParser()

    parser.add_argument('-c', '--checkpoint_path', type=str, default=DEFAULT_CKPT_PATH,
                        help='Checkpoint name or path, default to %(default)r')
    parser.add_argument('--cpu_only', action='store_true', help='Run demo with CPU only')
    parser.add_argument('--flash_attn2', action='store_true', default=False,
                        help='Enable flash_attention_2 when loading the model.')
    parser.add_argument('--share', action='store_true', default=False,
                        help='Create a publicly shareable link for the interface.')
    parser.add_argument('--inbrowser', action='store_true', default=False,
                        help='Automatically launch the interface in a new tab on the default browser.')
    parser.add_argument('--server_port', type=int, default=6020, help='Demo server port.')
    parser.add_argument('--server_name', type=str, default='127.0.0.1', help='Demo server name.')
    parser.add_argument('--dcu_id', type=str, default=None, help='Specify the GPU ID to load the model onto.')

    args = parser.parse_args()
    return args


def load_model_processor(args):
    if args.cpu_only:
        device_map = 'cpu'
    else:
        os.environ['CUDA_VISIBLE_DEVICES'] = args.dcu_id
        device_map = 'auto'
        print(f"Visible CUDA devices: {os.environ.get('CUDA_VISIBLE_DEVICES')}")

    if args.flash_attn2:
        model = Qwen2VLForConditionalGeneration.from_pretrained(
            args.checkpoint_path,
            torch_dtype=torch.float16,
            attn_implementation='flash_attention_2',
            device_map=device_map
        )
    else:
        model = Qwen2VLForConditionalGeneration.from_pretrained(
            args.checkpoint_path,
            torch_dtype=torch.float16,
            device_map=device_map
        )

    processor = AutoProcessor.from_pretrained(args.checkpoint_path)
    return model, processor

def parse_text(text):
    lines = text.split("\n")
    lines = [line for line in lines if line.strip() != ""]  # 去除空行
    count = 0
    parsed_lines = []

    for i, line in enumerate(lines):
        if "```" in line:
            count += 1
            items = line.split("`")
            if count % 2 == 1:
                # 开始代码块
                parsed_lines.append(f'<pre><code class="language-{items[-1]}">')
            else:
                # 结束代码块
                parsed_lines.append(f"</code></pre>")
        else:
            if i > 0 and count % 2 == 1:
                # 转义代码块内的特殊字符
                line = line.replace("`", r"\`")
                line = line.replace("<", "&lt;")
                line = line.replace(">", "&gt;")
                line = line.replace(" ", "&nbsp;")
                line = line.replace("*", "&ast;")
                line = line.replace("_", "&lowbar;")
                line = line.replace("-", "&#45;")
                line = line.replace(".", "&#46;")
                line = line.replace("!", "&#33;")
                line = line.replace("(", "&#40;")
                line = line.replace(")", "&#41;")
                line = line.replace("$", "&#36;")
            # 使用空格连接行
            if parsed_lines:
                parsed_lines[-1] += " " + line
            else:
                parsed_lines.append(line)

    text = "".join(parsed_lines)
    return text

def unparse_text(parsed_text):
    in_code_block = False
    lines = parsed_text.split("\n")
    unparsed_lines = []
    
    for line in lines:
        if "<pre><code" in line:
            in_code_block = True
            # 移除开始标签
            line = line.split(">", 1)[1]
        elif "</code></pre>" in line:
            in_code_block = False
            # 移除结束标签
            line = line.rsplit("<", 1)[0]
        
        # 反转 HTML 实体
        line = line.replace("&lt;", "<")
        line = line.replace("&gt;", ">")
        line = line.replace("&nbsp;", " ")
        line = line.replace("&ast;", "*")
        line = line.replace("&lowbar;", "_")
        line = line.replace("&#45;", "-")
        line = line.replace("&#46;", ".")
        line = line.replace("&#33;", "!")
        line = line.replace("&#40;", "(")
        line = line.replace("&#41;", ")")
        line = line.replace("&#36;", "$")
        
        # 如果在代码块内，还原反斜杠转义
        if in_code_block:
            line = line.replace(r"\`", "`")
        
        unparsed_lines.append(line)
    
    # 合并所有行
    unparsed_text = "\n".join(unparsed_lines)
    return unparsed_text

def remove_image_special(text):
    text = text.replace('<ref>', '').replace('</ref>', '')
    return re.sub(r'<box>.*?(</box>|$)', '', text)

def is_video_file(filename):
    video_extensions = ['.mp4', '.avi', '.mkv', '.mov', '.wmv', '.flv', '.webm', '.mpeg']
    return any(filename.lower().endswith(ext) for ext in video_extensions)

def transform_messages(original_messages):
    transformed_messages = []
    for message in original_messages:
        new_content = []
        for item in message['content']:
            if 'image' in item:
                new_item = {'type': 'image', 'image': item['image']}
            elif 'text' in item:
                new_item = {'type': 'text', 'text': item['text']}
            elif 'video' in item:
                new_item = {'type': 'video', 'video': item['video']}
            else:
                continue
            new_content.append(new_item)

        new_message = {'role': message['role'], 'content': new_content}
        transformed_messages.append(new_message)

    return transformed_messages

def _gc():
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

def call_local_model(model, processor, messages):
    messages = transform_messages(messages)

    text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    image_inputs, video_inputs = process_vision_info(messages)
    inputs = processor(text=[text], images=image_inputs, videos=video_inputs, padding=True, return_tensors='pt')
    inputs = inputs.to(model.device)

    tokenizer = processor.tokenizer
    streamer = TextIteratorStreamer(tokenizer, timeout=2000.0, skip_prompt=True, skip_special_tokens=True)

    gen_kwargs = {'max_new_tokens': 512, 'streamer': streamer, **inputs}

    thread = Thread(target=model.generate, kwargs=gen_kwargs)
    thread.start()

    generated_text = ''
    for new_text in streamer:
        generated_text += new_text
        yield parse_text(generated_text)

def create_predict_fn(model, processor):
    def predict(_chatbot, task_history):
        chat_query = _chatbot[-1][0]
        query = task_history[-1][0]
        if len(chat_query) == 0:
            _chatbot.pop()
            task_history.pop()
            return _chatbot
        print('User: ' + parse_text(query))
        history_cp = copy.deepcopy(task_history)
        full_response = ''
        messages = []
        content = []
        for q, a in history_cp:
            if isinstance(q, (tuple, list)):
                if is_video_file(q[0]):
                    content.append({'video': f'file://{q[0]}'})
                else:
                    content.append({'image': f'file://{q[0]}'})
            else:
                content.append({'text': q})
                messages.append({'role': 'user', 'content': content})
                messages.append({'role': 'assistant', 'content': [{'text': a}]})
                content = []
        messages.pop()

        for response in call_local_model(model, processor, messages):
            _chatbot[-1] = (parse_text(chat_query), remove_image_special(parse_text(response)))

            yield _chatbot
            full_response = parse_text(response)

        task_history[-1] = (query, full_response)
        print('Qwen-VL-Chat: ' + unparse_text(full_response))
        yield _chatbot

    return predict

# 启用加载模型
args = get_args()
model, processor = load_model_processor(args)

class Item(BaseModel):
    image_path: str
    text: str

@app.post("/predict")
async def predict(item: Item):
    messages = [
        {
            'role': 'user',
            'content': [
                {'image': item.image_path},
                {'text': item.text}
            ]
        }
    ]

    generated_text = ''
    for response in call_local_model(model, processor, messages):
        generated_text = unparse_text(response)

    _gc()

    return {"Generated Text": generated_text}

if __name__ == "__main__":
    import uvicorn
    args = get_args()
    uvicorn.run(app, host=args.server_name, port=args.server_port)
