from __future__ import annotations

import hashlib
import os
import time

import numpy as np
import torch

from .core import normalize_bbox, normalize_keypoints


class Models:
    def __init__(self) -> None:
        self.det_config = os.environ["DET_CONFIG"]
        self.det_checkpoint = os.environ["DET_CHECKPOINT"]
        self.pose_config = os.environ["POSE_CONFIG"]
        self.pose_checkpoint = os.environ["POSE_CHECKPOINT"]
        self.detector = self.pose = None

    @staticmethod
    def sha256(path: str) -> str:
        digest = hashlib.sha256()
        with open(path, "rb") as fh:
            for chunk in iter(lambda: fh.read(8 * 1024 * 1024), b""):
                digest.update(chunk)
        return digest.hexdigest()

    def load(self) -> None:
        from mmdet.apis import init_detector
        from mmpose.apis import init_model
        self.detector = init_detector(self.det_config, self.det_checkpoint, device="cuda:0")
        self.pose = init_model(self.pose_config, self.pose_checkpoint, device="cuda:0")

    @torch.inference_mode()
    def infer(self, frame: np.ndarray, threshold: float, max_people: int) -> tuple[list[dict], float, float]:
        from mmdet.apis import inference_detector
        from mmengine.registry import init_default_scope
        from mmpose.apis import inference_topdown
        height, width = frame.shape[:2]
        t0 = time.perf_counter_ns()
        init_default_scope("mmdet")
        result = inference_detector(self.detector, frame)
        pred = result.pred_instances.cpu().numpy()
        keep = (pred.labels == 0) & (pred.scores >= threshold)
        boxes = pred.bboxes[keep]
        scores = pred.scores[keep]
        if len(boxes):
            order = np.argsort(scores)[::-1][:max_people]
            boxes, scores = boxes[order], scores[order]
        det_ms = (time.perf_counter_ns() - t0) / 1e6
        t1 = time.perf_counter_ns()
        init_default_scope("mmpose")
        pose_results = inference_topdown(self.pose, frame, bboxes=boxes) if len(boxes) else []
        pose_ms = (time.perf_counter_ns() - t1) / 1e6
        people = []
        for idx, (box, det_score, sample) in enumerate(zip(boxes, scores, pose_results)):
            points = sample.pred_instances.keypoints[0]
            kp_scores = sample.pred_instances.keypoint_scores[0]
            combined = np.column_stack((points, kp_scores))
            wholebody = normalize_keypoints(combined, width, height, 133)
            people.append({
                "detection_id": idx,
                "bbox_xyxy_norm": normalize_bbox(box, width, height),
                "det_score": float(det_score),
                "keypoint_schema": "coco_wholebody_133",
                "keypoints_norm": wholebody,
                "body17": wholebody[:17],
                "pose_score": float(np.mean(np.nan_to_num(kp_scores[:17]))),
            })
        return people, det_ms, pose_ms