import json
import os
import cv2
import threading
from datetime import datetime
from queue import Queue
from typing import Dict

from app.core.model import Model
from app.core.preprocess import Preprocess
from app.services.model import TaskStatus, TaskStore


class Worker:
    def __init__(self):
        self.detection = Model().getModel()
        self.preprocess = Preprocess().getPreprocess()
        self.task_queue = Queue()
        self.task_store: Dict[str, TaskStore] = {}

        threading.Thread(target=self.worker, daemon=True).start()

    def worker(self):
        from app.main import UPLOAD_DIR
        while True:
            task_id = self.task_queue.get()
            if task_id is None:
                break

            task = self.task_store.get(task_id)
            if not task:
                continue

            try:
                task.status = TaskStatus.PROCESSING.name
                task.progress = 0
                print(f"开始处理任务 {task_id}...")

                # 创建输出目录
                output_dir = os.path.join(UPLOAD_DIR, task_id, "outputs")
                os.makedirs(output_dir, exist_ok=True)

                # 获取图像的标签列表
                image_labels = self.preprocess.preprocess(task.images)  # 返回一个0和1的列表，0代表跳过，1代表进行检测

                for idx, (input_img_path, label) in enumerate(zip(task.images, image_labels)):
                    print(f"处理任务 {task_id}, 处理图片 {input_img_path}...")

                    if label == 0:
                        # 如果标签是0，跳过模型检测，输出路径和坐标为空
                        task.result.append(
                            {"input_img_path": input_img_path, "output_img_path": "", "coords": "[]"}
                        )
                    else:
                        # 进行模型检测
                        img_res, coords_res = self.detection.detect(input_img_path)

                        coords_res = [{"name": name, "coords": coords} for name, coords in coords_res]
                        coords_json = json.dumps(coords_res, ensure_ascii=False)

                        out_img_path = os.path.join(output_dir, f"{idx}.jpg")
                        cv2.imwrite(out_img_path, img_res)

                        task.result.append(
                            {"input_img_path": input_img_path, "output_img_path": out_img_path, "coords": coords_json}
                        )

                    task.progress = int((idx + 1) / len(task.images) * 100)

                task.status = TaskStatus.COMPLETED.name
                task.completedAt = datetime.now()
                task.message = "处理完成"

                print(f"任务 {task_id} 处理完成")

            except Exception as e:
                task.status = TaskStatus.FAILED.name
                task.message = str(e)
            finally:
                self.task_queue.task_done()