import json
import os
import re

import requests
import torch
import logging
from PIL import Image
from logging.handlers import RotatingFileHandler

class detect_tasks():
    def __init__(self):
        #线程名称
        self._thread_name = ''
        self.device = None
        # 初始化检测模型
        self.qwen_tokenizer = None
        self.qwen_model = None

        # 初始化Milvus集合
        self.collection = None

    def remove_duplicate_lines(self,text):
        seen = set()
        result = []
        for line in text.split('。'):  # 按句号分割
            if line.strip() and line not in seen:
                seen.add(line)
                result.append(line)
        return '。'.join(result)

    def init_logging(self,logname):
        # 创建实例专属logger
        self.logger = logging.getLogger(f"{self.__class__}_{id(self)}")
        self.logger.setLevel(logging.INFO)
        # 避免重复添加handler
        if not self.logger.handlers:
            handler = RotatingFileHandler(
                filename=os.path.join("logs", logname+'_log.log'),
                maxBytes=10 * 1024 * 1024,
                backupCount=3,
                encoding='utf-8'
            )
            formatter = logging.Formatter(
                '%(asctime)s - %(filename)s:%(lineno)d - %(funcName)s() - %(levelname)s: %(message)s'
            )
            handler.setFormatter(formatter)
            self.logger.addHandler(handler)

    def tark_do(self,image_id,filesavepath,ragurl,max_tokens):
        try :
            # 1. 从集合A获取向量和元数据
            res_a = self.collection.query(
                expr=f"id == {image_id}",
                output_fields=["id", "zh_desc_class", "text_vector", "bounding_box", "video_point_name", "task_id",
                               "task_name","event_level_id","event_level_name",
                               "video_point_id", "detect_num", "is_waning", "waning_value", "rule_id", "detect_id",
                               "detect_time", "image_path", "image_desc_path","video_path"],
                consistency_level="Strong"
            )
            # 图片和视频地址
            image = Image.open(f"{res_a[0]['image_desc_path']}").convert("RGB")  # 替换为您的图片路径
            conversation = [
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "image",
                        },
                        {"type": "text", "text": "请按以下要求描述图片：\n1. 列出主要物体\n2.不进行推理和think\n返回小于2000字的整段描述,描述中的物体信息不加数字序号"},
                    ],
                }
            ]

            # Preprocess the inputs
            self.logger.info("aaaa")
            text_prompt = self.qwen_tokenizer.apply_chat_template(conversation, add_generation_prompt=True)
            # Excepted output: '<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>Describe this image.<|im_end|>\n<|im_start|>assistant\n'
            self.logger.info("bbbb")
            inputs = self.qwen_tokenizer(
                text=[text_prompt], images=[image], padding=True, return_tensors="pt"
            )
            inputs = inputs.to("cuda")
            self.logger.info("cccc")
            torch.cuda.empty_cache()
            with torch.no_grad():
                output_ids = self.qwen_model.generate(**inputs, max_new_tokens=50)
            print(output_ids.device)
            self.logger.info("dddd")
            generated_ids = [
                output_ids[len(input_ids):]
                for input_ids, output_ids in zip(inputs.input_ids, output_ids)
            ]
            image_text = self.qwen_tokenizer.batch_decode(
                generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True
            )
            self.logger.info("ffff")
            image_des = (image_text[0]).replace('\n', '')

            if len(image_des)>4 and image_des.startswith("这张图片"):
                image_des = image_des[4:]
            self.logger.info(image_des)
            image_des = self.remove_duplicate_lines(image_des)

            is_waning = 0

            data = {
                "id": image_id,
                "event_level_id": res_a[0]['event_level_id'],  # event_level_id
                "event_level_name": res_a[0]['event_level_name'],  # event_level_id
                "rule_id": res_a[0]["rule_id"],
                "video_point_id": res_a[0]['video_point_id'],  # video_point_id
                "video_point_name": res_a[0]['video_point_name'],
                "is_waning": is_waning,
                "zh_desc_class": image_des,  # text_vector
                "bounding_box": res_a[0]['bounding_box'],  # bounding_box
                "task_id": res_a[0]['task_id'],  # task_id
                "task_name": res_a[0]['task_name'],  # task_id
                "detect_id": res_a[0]['detect_id'],  # detect_id
                "detect_time": res_a[0]['detect_time'],  # detect_time
                "detect_num": res_a[0]['detect_num'],
                "waning_value": res_a[0]['waning_value'],
                "image_path": res_a[0]['image_path'],  # image_path
                "image_desc_path": res_a[0]['image_desc_path'],  # image_path
                "video_path": res_a[0]['video_path'],
                "text_vector": res_a[0]['text_vector']
            }
            # 保存到milvus
            image_id = self.collection.upsert(data).primary_keys
            data = {
                "id": str(image_id[0]),
                "video_point_id": res_a[0]['video_point_id'],
                "video_path": res_a[0]["video_point_name"],
                "zh_desc_class": image_des,
                "detect_time": res_a[0]['detect_time'],
                "image_path": f"{res_a[0]['image_path']}",
                "task_name": res_a[0]["task_name"],
                "event_level_name": res_a[0]["event_level_name"],
                "rtsp_address": f"{res_a[0]['video_path']}"
            }

            data = {'collection_name': "smartrag","data": data,"description":""}
            requests.post(ragurl+"/insert_json_data",
                json=data
            )
            return image_id
        except Exception as e:
            self.logger.info(f"{self._thread_name}线程：执行模型解析时出错:任务：{image_id} :{e}")
            return 0