import os

from lib.json_fun import f_detections_to_objects
from lib.qwen_fun import chat, hazard_inspection, load_json_data, search_knowledge_base, upload_files_and_get_urls_concurrently
from encodings.punycode import T
"""
测试 在给定标注框与类别，原始视频经过转换之后，AI能否准确识别物体特征
"""

from tkinter import N
from datetime import datetime
import json
from lib.qwen_fun_vid import generate_video_to_objects
from lib.sam3 import SAM3
from ultralytics.models.sam import SAM3VideoSemanticPredictor
import json
import numpy as np
from pathlib import Path

def run():
    vid_name = "santai5" # 视频名称（不含后缀）
    vid_end = ".mp4" # 视频后缀
    use_url_cache = True # 是否使用 URL 缓存，避免重复上传视频
    enable_thinking = False # 是否启用思考模式

    run_vid_process = False # 是否运行视频处理流程（提取物体视频）

    input_video_path = f"input/{vid_name}{vid_end}"
    output_dir = f"output/{vid_name}"
    frame_detections_path = f"{output_dir}/frame_detections.json"
    objects_json_path = f"{output_dir}/objects.json"
    vid_dir = f"{output_dir}/obj_vids"
    obj = json.load(open(objects_json_path, "r", encoding="utf-8"))
    class_list = obj["class_list"]
    obj_dict = obj["track_id_list"]
    rule_dict: dict = load_json_data('知识库/rule.json')
    video_url_file = f"{output_dir}/video_url.json"

    # 提取物体信息
    f_detections_to_objects(
        frame_detections_path,
        objects_json_path
    )

    if run_vid_process:
        # 生成物体视频
        generate_video_to_objects(
            obj_dict,
            input_video_path,
            output_dir=vid_dir,
        )

    # 上传视频并获取 URL
    vid_dict = get_vid_dict(vid_dir, obj_dict, video_url_file, use_url_cache)

    hazard_inspection(
        output_dir,
        obj_dict,
        rule_dict,
        class_list,
        vid_dict,
        fps=30,
        enable_thinking=enable_thinking
    )

def get_vid_dict(vid_dir: str, obj_dict: dict, video_url_file: str, use_url_cache: bool) -> dict:
    """获取视频字典
    参数:
        vid_dir: 物体视频存放目录
        obj_dict: 物品字典，键为物品TrackID，值为物品信息
        video_url_file: 视频URL文件路径
        use_url_cache: 是否使用缓存的URL
    返回:
        视频字典，键为物品TrackID，值为视频本地地址和视频URL
    """

    vid_dict = {}

    # 检查URL是否已存在
    if use_url_cache:
        if os.path.exists(video_url_file):
            try:
                with open(video_url_file, "r", encoding="utf-8") as f:
                    vid_dict = json.load(f)
            except Exception as e:
                print(f"读取URL文件失败: {e}")
    for track_id in obj_dict.keys():
        if track_id in vid_dict:
            continue
        vid_path = f"{vid_dir}/obj_{int(track_id):03d}.mp4"
        vid_dict[track_id] = {"vid_path": vid_path, "vid_url": None}

    # 找到未上传的视频
    unuploaded_vids = []
    for track_id, vid_info in vid_dict.items():
        if vid_info["vid_url"] is None:
            unuploaded_vids.append(vid_info["vid_path"])

    # 上传未上传的视频并获取 URL
    uploaded_urls = upload_files_and_get_urls_concurrently(
        file_path_list=unuploaded_vids,
        max_workers=8
    )

    # 更新视频URL
    for track_id, vid_info in vid_dict.items():
        if vid_info["vid_url"] is None:
            idx = unuploaded_vids.index(vid_info["vid_path"])
            vid_info["vid_url"] = uploaded_urls[idx]

    # 保存字典为 JSON 文件
    with open(video_url_file, "w", encoding="utf-8") as f:
        json.dump(vid_dict, f, ensure_ascii=False, indent=4)

    return vid_dict


# 启动应用
if __name__ == "__main__":
    # objects_json_path = "output/santai5/objects.json"
    # input_video_path = "input/santai5.mp4"
    # obj_dict = json.load(open(objects_json_path, "r", encoding="utf-8"))["track_id_list"]

    # # 记录开始时间
    # start_time = datetime.now()
    # generate_video_to_objects(
    #     obj_dict,
    #     input_video_path,
    #     output_dir="output/santai5/obj_vids",
    # )

    run()