import cv2
import numpy as np
from pathlib import Path

from ultralytics.engine.results import Results

from hardware.camera.camera import Camera
from models.detection.detector import Detector
from models.detection.board_manager import BoardManager
from models.detection.pieces_manager import PiecesManager


class DetectionService:

    edges_detector : Detector
    pieces_detector : Detector

    board_manager : BoardManager
    pieces_manager : PiecesManager

    scale_size : tuple[int, int]

    camera :  Camera

    def __init__(self):
        current_file = Path(__file__).resolve()
        project_root = current_file.parent.parent

        self.edges_detector = Detector(project_root / "assets" / "models" / "edges.pt")
        self.pieces_detector = Detector(project_root / "assets" / "models" / "unified-nano-refined.pt")

        self.pieces_manager = PiecesManager()
        self.board_manager = BoardManager()
        self.scale_size = (800, 800)
        self.camera = Camera()

    def start(self):
        self.camera.open()

    def stop(self):
        self.camera.close()

    def analyze_single_frame(self) -> tuple[bytes, str | None]:
        frame = self.camera.take_photo()
        encoded_frame = cv2.imencode('.jpg', frame, [int(cv2.IMWRITE_JPEG_QUALITY), 80])[1].tobytes()

        result = self.__run_complete_detection(frame)

        edges_prediction = result["edges"]
        pieces_prediction = result["pieces"]

        processed_frame = self.board_manager.process_frame(edges_prediction[0], frame, self.scale_size)
        if processed_frame is None:
            return encoded_frame, None

        warped_corners, matrix = processed_frame

        detections = self.pieces_manager.extract_pieces(pieces_prediction)

        board = self.pieces_manager.pieces_to_board(detections, warped_corners, matrix, self.scale_size)

        return encoded_frame, self.pieces_manager.board_to_fen(board)

    def __run_complete_detection(self, frame : np.ndarray, display=False) -> dict[str, list[Results]] :
        pieces_prediction = self.__run_pieces_detection(frame)
        edges_prediction = self.__run_edges_detection(frame)

        if display:
            edges_annotated_frame = edges_prediction[0].plot()
            pieces_annotated_frame = pieces_prediction[0].plot(img=edges_annotated_frame)
            self.__display_frame(pieces_annotated_frame)

        return { "edges" : edges_prediction, "pieces" : pieces_prediction}


    def __run_pieces_detection(self, frame : np.ndarray, display=False) -> list[Results]:
        prediction = self.pieces_detector.make_prediction(frame)
        if display:
            self.__display_frame(prediction[0].plot())
        return prediction


    def __run_edges_detection(self, frame : np.ndarray, display=False) -> list[Results]:
        prediction = self.edges_detector.make_prediction(frame)
        if display:
            self.__display_frame(prediction[0].plot())
        return prediction

    def __display_frame(self, frame : np.ndarray):
        cv2.namedWindow("Frame", cv2.WINDOW_NORMAL)
        cv2.resizeWindow("Frame", self.scale_size[0], self.scale_size[1])
        cv2.imshow("Frame", frame)
        cv2.waitKey(0)
        cv2.destroyAllWindows()
        return