# 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 configparser
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
from vllm import LLM, SamplingParams
import os
from loguru import logger

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

logger.add("parse.log", rotation="10 MB", level="INFO",
           format="{time} {level} {message}", encoding='utf-8', enqueue=True)


app = FastAPI()

DEFAULT_CKPT_PATH = '/home/practice/model/Qwen2-VL-7B-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('--gpu_nums', type=int, default=1, help='Number of GPUs to use for tensor parallelism.')
    parser.add_argument('--dcu_id', type=str, default=None, help='Specify the GPU ID to load the model onto.')
    parser.add_argument(
        '--config_path',
        default='/home/practice/magic_pdf-main/magic_pdf/config.ini',
        )
    args = parser.parse_args()
    return args

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

    llm = LLM(
        model=args.checkpoint_path,
        limit_mm_per_prompt={"image": 10, "video": 10},
        trust_remote_code=True,
        tensor_parallel_size=args.gpu_nums,  # 用args.gpu_nums根据实际情况调整
        dtype='float16',  # 或者 'bfloat16'
    )

    processor = AutoProcessor.from_pretrained(args.checkpoint_path)
    return llm, 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(llm, 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)
    mm_data = {}
    if image_inputs is not None:
        mm_data["image"] = image_inputs
    if video_inputs is not None:
        mm_data["video"] = video_inputs

    llm_inputs = {
        "prompt": text,
        "multi_modal_data": mm_data,
    }

    sampling_params = SamplingParams(
        temperature=0.1,
        top_p=0.001,
        repetition_penalty=1.05,
        max_tokens=256,
        stop_token_ids=[],
    )

    outputs = llm.generate([llm_inputs], sampling_params=sampling_params)
    generated_text = outputs[0].outputs[0].text

    yield parse_text(generated_text)

def create_predict_fn(llm, 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(llm, 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()
llm, processor = load_model_processor(args)

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

@app.get("/health")
async def health_check():
    return {"status": "healthy"}


@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(llm, processor, messages):
        generated_text = unparse_text(response)

    _gc()

    return {"Generated Text": generated_text}

if __name__ == "__main__":
    import uvicorn
    args = get_args()
    config = configparser.ConfigParser()
    config.read(args.config_path)
    # host = config.get('server', 'ocr_host')
    # port = int(config.get('server', 'ocr_port'))
    host, port = config.get('server', 'ocr_server').split('://')[1].split(':')[0], int(
        config.get('server', 'ocr_server').split('://')[1].split(':')[1])
    uvicorn.run(app, host=args.host, port=args.port)
