import os
from dataclasses import dataclass
from datetime import datetime
from typing import Optional, Tuple

import cv2
import numpy as np


@dataclass
class RecordResult:
    species_id: str
    temp_value: float
    temp_time: datetime
    video_start_time: datetime
    video_end_time: datetime
    video_duration: int
    rgb_video_path: str
    ir_video_path: str


class RecorderService:
    def __init__(self, output_dir: str, duration_sec: int = 3, fps: int = 15):
        self.output_dir = output_dir
        self.duration_sec = duration_sec
        self.fps = fps

        os.makedirs(self.output_dir, exist_ok=True)

        self.recording = False
        self.frames_target = self.duration_sec * self.fps
        self.frames_written = 0

        self.rgb_writer: Optional[cv2.VideoWriter] = None
        self.ir_writer: Optional[cv2.VideoWriter] = None

        self.current_species_id: str = "Unknown"
        self.current_temp_value: float = 0.0
        self.current_temp_time: Optional[datetime] = None
        self.start_time: Optional[datetime] = None

        self.rgb_path = ""
        self.ir_path = ""

    @staticmethod
    def _fourcc():
        return cv2.VideoWriter_fourcc(*"mp4v")

    @staticmethod
    def _to_ir_colormap(temp32f: np.ndarray) -> np.ndarray:
        finite = temp32f[np.isfinite(temp32f)]
        if finite.size == 0:
            h, w = temp32f.shape[:2]
            return np.zeros((h, w, 3), dtype=np.uint8)

        tmin = np.percentile(finite, 5)
        tmax = np.percentile(finite, 95)
        if tmax <= tmin:
            tmax = tmin + 1.0

        norm = np.clip((temp32f - tmin) / (tmax - tmin), 0.0, 1.0)
        gray = (norm * 255.0).astype(np.uint8)
        return cv2.applyColorMap(gray, cv2.COLORMAP_JET)

    def start(self, species_id: str, temp_value: float, temp_time: datetime, rgb_shape: Tuple[int, int], ir_shape: Tuple[int, int]) -> bool:
        if self.recording:
            return False

        h_rgb, w_rgb = rgb_shape
        h_ir, w_ir = ir_shape

        ts = temp_time.strftime("%Y%m%d_%H%M%S")
        self.rgb_path = os.path.join(self.output_dir, f"{species_id}_{ts}_rgb.mp4")
        self.ir_path = os.path.join(self.output_dir, f"{species_id}_{ts}_ir.mp4")

        self.rgb_writer = cv2.VideoWriter(self.rgb_path, self._fourcc(), self.fps, (w_rgb, h_rgb))
        self.ir_writer = cv2.VideoWriter(self.ir_path, self._fourcc(), self.fps, (w_ir, h_ir))

        if not self.rgb_writer.isOpened() or not self.ir_writer.isOpened():
            if self.rgb_writer is not None:
                self.rgb_writer.release()
            if self.ir_writer is not None:
                self.ir_writer.release()
            self.rgb_writer = None
            self.ir_writer = None
            return False

        self.recording = True
        self.frames_written = 0
        self.current_species_id = species_id
        self.current_temp_value = temp_value
        self.current_temp_time = temp_time
        self.start_time = datetime.now()
        return True

    def write(self, rgb_frame: np.ndarray, ir_temp_frame: np.ndarray) -> Optional[RecordResult]:
        if not self.recording or self.rgb_writer is None or self.ir_writer is None:
            return None

        ir_color = self._to_ir_colormap(ir_temp_frame)

        self.rgb_writer.write(rgb_frame)
        self.ir_writer.write(ir_color)
        self.frames_written += 1

        if self.frames_written < self.frames_target:
            return None

        self.rgb_writer.release()
        self.ir_writer.release()

        self.rgb_writer = None
        self.ir_writer = None
        self.recording = False

        end_time = datetime.now()
        result = RecordResult(
            species_id=self.current_species_id,
            temp_value=self.current_temp_value,
            temp_time=self.current_temp_time or end_time,
            video_start_time=self.start_time or end_time,
            video_end_time=end_time,
            video_duration=self.duration_sec,
            rgb_video_path=self.rgb_path,
            ir_video_path=self.ir_path,
        )
        return result

