import math
import threading
from collections import defaultdict, deque
from dataclasses import dataclass
from datetime import datetime
from typing import Dict, List, Optional, Tuple

import cv2
import numpy as np
import rclpy
from cv_bridge import CvBridge
from rclpy.node import Node
from rclpy.qos import HistoryPolicy, QoSProfile, ReliabilityPolicy
from sensor_msgs.msg import Image

from .recorder_service import RecorderService
from .trigger_policy import TriggerContext, TriggerPolicy
from .upload_client import ServerUploadClient, UploadConfig

try:
    from hbm_runtime import HB_HBMRuntime
except Exception:
    HB_HBMRuntime = None

CLASS_HEAD = 0
CLASS_EYE = 1
CLASS_ID = 2

ID_CLASS_MAP = {
    0: "6006",
    1: "6002",
    2: "6005",
    3: "6001",
    4: "6003",
}


@dataclass
class Det:
    x1: int
    y1: int
    x2: int
    y2: int
    cls_id: int
    score: float


class SimpleTracker:
    def __init__(self, max_lost: int = 30, match_dist: float = 150.0):
        self.next_id = 1
        self.tracks: Dict[int, Dict[str, object]] = {}
        self.lost: Dict[int, int] = {}
        self.max_lost = max_lost
        self.match_dist = match_dist

    def update(self, detections: List[Tuple[int, int, int, int]]) -> Dict[int, Dict[str, object]]:
        updated_tracks: Dict[int, Dict[str, object]] = {}
        used_track_ids = set()
        updated_lost: Dict[int, int] = {}

        for x1, y1, x2, y2 in detections:
            cx = (x1 + x2) / 2
            cy = (y1 + y2) / 2
            best_id = None
            best_dist = 1e9

            for tid, data in self.tracks.items():
                if tid in used_track_ids:
                    continue
                px, py = data["center"]
                dist = math.hypot(cx - px, cy - py)
                if dist < best_dist and dist < self.match_dist:
                    best_dist = dist
                    best_id = tid

            if best_id is None:
                best_id = self.next_id
                self.next_id += 1

            used_track_ids.add(best_id)
            updated_tracks[best_id] = {"bbox": [x1, y1, x2, y2], "center": [cx, cy]}
            updated_lost[best_id] = 0

        for tid, data in self.tracks.items():
            if tid not in updated_tracks:
                lost_count = self.lost.get(tid, 0) + 1
                if lost_count < self.max_lost:
                    updated_tracks[tid] = data
                    updated_lost[tid] = lost_count

        self.tracks = updated_tracks
        self.lost = updated_lost
        return self.tracks


class HbmYoloDetector:
    def __init__(self, model_path: str, target_size: int, conf_thres: float, iou_thres: float):
        if HB_HBMRuntime is None:
            raise RuntimeError("hbm_runtime not available")
        self.model = HB_HBMRuntime(model_path)
        self.model_name = self.model.model_names[0]
        self.input_names = self.model.input_names[self.model_name]
        self.conf_thres = conf_thres
        self.iou_thres = iou_thres
        self.target_size = target_size

    def infer(self, image: np.ndarray) -> List[Det]:
        h0, w0 = image.shape[:2]
        r = min(self.target_size / h0, self.target_size / w0)
        new_w, new_h = int(w0 * r), int(h0 * r)
        resized = cv2.resize(image, (new_w, new_h))

        pad_w = self.target_size - new_w
        pad_h = self.target_size - new_h
        top = pad_h // 2
        bottom = pad_h - top
        left = pad_w // 2
        right = pad_w - left

        img_letterbox = cv2.copyMakeBorder(resized, top, bottom, left, right, cv2.BORDER_CONSTANT, value=(114, 114, 114))
        x = img_letterbox.astype(np.float32) / 255.0
        x = np.transpose(x, (2, 0, 1))[None, ...]

        inputs = {name: x.astype(np.float32) for name in self.input_names}
        outputs = self.model.run(inputs)

        out = None
        for _, data in outputs[self.model_name].items():
            out = data
            break
        if out is None:
            return []

        preds = out[0].transpose(1, 0)
        boxes = preds[:, 0:4]
        cls_scores = preds[:, 4:]

        cls_probs = 1.0 / (1.0 + np.exp(-cls_scores))
        cls_ids = np.argmax(cls_probs, axis=1)
        scores = np.max(cls_probs, axis=1)

        keep = scores > self.conf_thres
        boxes, cls_ids, scores = boxes[keep], cls_ids[keep], scores[keep]
        if len(boxes) == 0:
            return []

        nms_boxes = []
        for cx, cy, w, h in boxes:
            nms_boxes.append([float(cx - w / 2), float(cy - h / 2), float(w), float(h)])
        idx = cv2.dnn.NMSBoxes(nms_boxes, scores.tolist(), self.conf_thres, self.iou_thres)
        if len(idx) == 0:
            return []

        dets: List[Det] = []
        for i in idx.flatten():
            cx, cy, w, h = boxes[i]
            x1 = int((cx - w / 2 - left) / r)
            y1 = int((cy - h / 2 - top) / r)
            x2 = int((cx + w / 2 - left) / r)
            y2 = int((cy + h / 2 - top) / r)

            x1 = max(0, min(w0 - 1, x1))
            y1 = max(0, min(h0 - 1, y1))
            x2 = max(0, min(w0 - 1, x2))
            y2 = max(0, min(h0 - 1, y2))
            if x2 <= x1 or y2 <= y1:
                continue

            dets.append(Det(x1, y1, x2, y2, int(cls_ids[i]), float(scores[i])))
        return dets


class CowTempCheckNode(Node):
    def __init__(self):
        super().__init__("cow_temp_check_node")
        self.rgb_topic = self.declare_parameter("rgb_topic", "/camera/front/image_raw").value
        self.ir_temp_topic = self.declare_parameter("ir_temp_topic", "/ir/temperature_celsius_corrected").value
        self.output_topic = self.declare_parameter("output_topic", "/cow/temp_check/image").value

        self.det_model_path = self.declare_parameter(
            "det_model_path", "/home/sunrise/opt/dev/project/aiagent/models/cowCheck/det/det_640x640.hbm"
        ).value
        self.id_model_path = self.declare_parameter(
            "id_model_path", "/home/sunrise/opt/dev/project/aiagent/models/cowCheck/id/id_128x128.hbm"
        ).value

        self.det_conf = float(self.declare_parameter("det_conf", 0.6).value)
        self.det_iou = float(self.declare_parameter("det_iou", 0.45).value)
        self.id_conf = float(self.declare_parameter("id_conf", 0.6).value)
        self.id_iou = float(self.declare_parameter("id_iou", 0.45).value)

        self.track_max_lost = int(self.declare_parameter("track_max_lost", 30).value)
        self.rgb_reliability = self.declare_parameter("rgb_reliability", "best_effort").value
        self.ir_reliability = self.declare_parameter("ir_reliability", "best_effort").value

        self.record_output_dir = self.declare_parameter(
            "record_output_dir", "/home/sunrise/opt/dev/project/aiagent/capabilities/cowTempCheck/records"
        ).value
        self.record_duration_sec = int(self.declare_parameter("record_duration_sec", 3).value)
        self.record_fps = int(self.declare_parameter("record_fps", 15).value)
        self.trigger_center_ratio_x = float(self.declare_parameter("trigger_center_ratio_x", 0.3).value)
        self.trigger_cooldown_sec = int(self.declare_parameter("trigger_cooldown_sec", 30).value)

        self.api_username = self.declare_parameter("api_username", "admin").value
        self.api_password = self.declare_parameter("api_password", "gbd2025").value
        self.api_device_id = self.declare_parameter("api_device_id", "CAM-001").value
        self.enable_ir_video_url = bool(self.declare_parameter("enable_ir_video_url", True).value)
        self.enable_upload = bool(self.declare_parameter("enable_upload", True).value)

        self.bridge = CvBridge()
        self.last_temp: Optional[np.ndarray] = None

        self.tracker = SimpleTracker(max_lost=self.track_max_lost)
        self.track_cowid: Dict[int, str] = {}
        self.track_id_history: Dict[int, deque] = defaultdict(lambda: deque(maxlen=20))

        if not self.det_model_path or not self.id_model_path:
            raise RuntimeError("det_model_path or id_model_path not set")

        self.detector = HbmYoloDetector(self.det_model_path, target_size=640, conf_thres=self.det_conf, iou_thres=self.det_iou)
        self.id_detector = HbmYoloDetector(self.id_model_path, target_size=128, conf_thres=self.id_conf, iou_thres=self.id_iou)

        self.trigger_policy = TriggerPolicy(center_ratio_x=self.trigger_center_ratio_x, cooldown_sec=self.trigger_cooldown_sec)
        self.recorder = RecorderService(self.record_output_dir, duration_sec=self.record_duration_sec, fps=self.record_fps)

        self.upload_client = ServerUploadClient(
            UploadConfig(
                username=self.api_username,
                password=self.api_password,
                device_id=self.api_device_id,
                enable_ir_video_url=self.enable_ir_video_url,
            )
        )

        rgb_qos = self._make_qos(self.rgb_reliability)
        ir_qos = self._make_qos(self.ir_reliability)

        self.create_subscription(Image, self.rgb_topic, self.on_rgb, rgb_qos)
        self.create_subscription(Image, self.ir_temp_topic, self.on_temp, ir_qos)
        self.pub_img = self.create_publisher(Image, self.output_topic, 10)

    def _make_qos(self, reliability: str) -> QoSProfile:
        qos = QoSProfile(depth=5)
        qos.history = HistoryPolicy.KEEP_LAST
        qos.reliability = ReliabilityPolicy.RELIABLE if reliability == "reliable" else ReliabilityPolicy.BEST_EFFORT
        return qos

    def _map_rgb_bbox_to_temp(self, bbox, rgb_shape, temp_shape):
        x1, y1, x2, y2 = bbox
        rh, rw = rgb_shape[:2]
        th, tw = temp_shape[:2]
        sx = tw / max(1, rw)
        sy = th / max(1, rh)
        tx1, ty1 = int(x1 * sx), int(y1 * sy)
        tx2, ty2 = int(x2 * sx), int(y2 * sy)
        tx1, ty1 = max(0, min(tw - 1, tx1)), max(0, min(th - 1, ty1))
        tx2, ty2 = max(0, min(tw - 1, tx2)), max(0, min(th - 1, ty2))
        if tx2 <= tx1:
            tx2 = min(tw - 1, tx1 + 1)
        if ty2 <= ty1:
            ty2 = min(th - 1, ty1 + 1)
        return tx1, ty1, tx2, ty2

    def _eye_max_temp(self, eye_bbox, rgb_shape):
        if self.last_temp is None:
            return None
        tx1, ty1, tx2, ty2 = self._map_rgb_bbox_to_temp(eye_bbox, rgb_shape, self.last_temp.shape)
        roi = self.last_temp[ty1:ty2, tx1:tx2]
        if roi.size == 0:
            return None
        finite = roi[np.isfinite(roi)]
        if finite.size == 0:
            return None
        return float(np.max(finite))

    def on_temp(self, msg: Image):
        self.last_temp = self.bridge.imgmsg_to_cv2(msg, desired_encoding="32FC1").astype(np.float32)

    def _assign_id(self, frame, tracks, id_boxes):
        used_tracks = set()
        for d in id_boxes:
            ix1, iy1, ix2, iy2 = d.x1, d.y1, d.x2, d.y2
            cx, cy = (ix1 + ix2) / 2, (iy1 + iy2) / 2
            best_tid, best_dist = None, 1e9

            for tid, data in tracks.items():
                hx1, hy1, hx2, hy2 = data["bbox"]
                hcx, hcy = (hx1 + hx2) / 2, (hy1 + hy2) / 2
                dist = math.hypot(cx - hcx, cy - hcy)
                if dist < best_dist and dist < 200:
                    best_dist, best_tid = dist, tid

            if best_tid is None or best_tid in used_tracks:
                continue
            used_tracks.add(best_tid)

            crop = frame[iy1:iy2, ix1:ix2]
            if crop is None or crop.size == 0:
                continue
            crop = cv2.resize(crop, (128, 128), interpolation=cv2.INTER_LINEAR)
            dd = self.id_detector.infer(crop)
            if not dd:
                continue

            best = max(dd, key=lambda x: x.score)
            digit = ID_CLASS_MAP.get(int(best.cls_id), str(best.cls_id))
            self.track_id_history[best_tid].append(digit)
            hist = self.track_id_history[best_tid]
            cow_id = max(set(hist), key=hist.count)
            self.track_cowid[best_tid] = cow_id

            cv2.putText(frame, f"ID:{digit}", (ix1, max(0, iy1 - 10)), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 0, 0), 2)

    def _spawn_upload(self, rec):
        def _run():
            try:
                self.upload_client.upload_record(rec)
                self.get_logger().info(f"upload success: {rec.species_id} {rec.rgb_video_path}")
            except Exception as e:
                self.get_logger().error(f"upload failed: {e}")

        threading.Thread(target=_run, daemon=True).start()

    def on_rgb(self, msg: Image):
        frame = self.bridge.imgmsg_to_cv2(msg, desired_encoding="bgr8")
        dets = self.detector.infer(frame)

        heads, eyes, ids = [], [], []
        for d in dets:
            if d.cls_id == CLASS_HEAD:
                heads.append((d.x1, d.y1, d.x2, d.y2))
            elif d.cls_id == CLASS_EYE:
                eyes.append(d)
            elif d.cls_id == CLASS_ID:
                ids.append(d)

        tracks = self.tracker.update(heads)
        self._assign_id(frame, tracks, ids)

        best_species_id = None
        best_head_center_x = None
        best_eye_temp = None
        has_eye_temp = False

        for d in eyes:
            cv2.rectangle(frame, (d.x1, d.y1), (d.x2, d.y2), (255, 0, 0), 2)
            cv2.putText(frame, "EYE", (d.x1, max(0, d.y1 - 5)), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 1)

            tmax = self._eye_max_temp((d.x1, d.y1, d.x2, d.y2), frame.shape)
            if tmax is not None:
                has_eye_temp = True
                if best_eye_temp is None or tmax > best_eye_temp:
                    best_eye_temp = tmax
                cv2.putText(frame, f"{tmax:.2f}C", (d.x2 + 5, max(0, d.y1 - 5)), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 165, 255), 2)

        for d in ids:
            cv2.rectangle(frame, (d.x1, d.y1), (d.x2, d.y2), (0, 0, 255), 2)
            cv2.putText(frame, "ID_BOX", (d.x1, max(0, d.y1 - 5)), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1)

        for tid, data in tracks.items():
            x1, y1, x2, y2 = map(int, data["bbox"])
            cow_id = self.track_cowid.get(tid, "Unknown")
            cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
            cv2.putText(frame, f"Cow:{cow_id} T:{tid}", (x1, max(0, y1 - 10)), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)

            if cow_id != "Unknown" and best_species_id is None:
                best_species_id = cow_id
                best_head_center_x = (x1 + x2) / 2.0

        now = datetime.now()
        if best_species_id is not None and best_head_center_x is not None:
            ctx = TriggerContext(
                has_eye=len(eyes) > 0,
                has_temp=has_eye_temp,
                has_id=True,
                head_center_x=best_head_center_x,
                image_width=frame.shape[1],
                species_id=best_species_id,
            )
            if self.trigger_policy.should_trigger(ctx, now):
                if self.last_temp is not None and best_eye_temp is not None:
                    started = self.recorder.start(
                        species_id=best_species_id,
                        temp_value=best_eye_temp,
                        temp_time=now,
                        rgb_shape=(frame.shape[0], frame.shape[1]),
                        ir_shape=(self.last_temp.shape[0], self.last_temp.shape[1]),
                    )
                    if started:
                        self.get_logger().info(f"record start: species={best_species_id}, temp={best_eye_temp:.2f}")

        if self.recorder.recording and self.last_temp is not None:
            rec = self.recorder.write(frame, self.last_temp)
            if rec is not None:
                self.get_logger().info(f"record done: {rec.rgb_video_path}, {rec.ir_video_path}")
                if self.enable_upload:
                    self._spawn_upload(rec)

        out = self.bridge.cv2_to_imgmsg(frame, encoding="bgr8")
        out.header = msg.header
        self.pub_img.publish(out)


def main():
    rclpy.init()
    node = CowTempCheckNode()
    try:
        rclpy.spin(node)
    finally:
        node.destroy_node()
        rclpy.shutdown()


if __name__ == "__main__":
    main()

