from picamera2 import Picamera2
from threading import Lock, RLock
import threading
import time
import numpy as np
import cv2
from pathlib import Path


class CameraManager:
    def __init__(self, state):
        self.state = state
        self.initialized = False

        self.camera_lock = RLock()

        self.cameras_runtime = {}
        self._reconfigure_needed = False
        self._sensor_modes_cache = None

    # =========================================================
    # Estado interno
    # =========================================================

    def mark_reconfigure_needed(self):
        with self.camera_lock:
            self._reconfigure_needed = True

    def _init_camera_runtime(self, cam_spec):
        return {
            "camera_id": cam_spec.id,
            "camera_index": cam_spec.index,
            "role": cam_spec.role,
            "interface": getattr(cam_spec, "interface", "CSI"),

            "backend": None,           # picamera2 | opencv
            "picam2": None,
            "cap": None,

            "last_frame": None,
            "frame_id": 0,
            "frame_ts": None,
            "buffer": None,

            "stop_event": threading.Event(),
            "thread": None,

            "last_signature": None,
            "last_read_ts": None,
            "last_new_frame_ts": None,

            "frame_lock": Lock(),
            "consecutive_failures": 0,
        }

    def _make_frame_signature(self, frame: np.ndarray):
        if frame is None:
            return None

        h, w = frame.shape[:2]

        # pega uma grade simples e barata
        step_y = max(1, h // 16)
        step_x = max(1, w // 16)

        sample = frame[::step_y, ::step_x]

        # reduz para assinatura curtinha
        return (
            sample.shape,
            int(sample.mean()),
            int(sample.std()),
            int(sample[0, 0, 0]) if sample.ndim == 3 else int(sample[0, 0]),
            int(sample[-1, -1, 1]) if sample.ndim == 3 and sample.shape[2] > 1 else 0,
        )

    # =========================================================
    # Inicialização
    # =========================================================

    def begin(self):
        with self.camera_lock:
            self.stop()

            required_ids = set(self.state.get_bootstrap_camera_ids_for_frame_type())
            if not required_ids:
                print("[WARN] Nenhuma câmera requerida para o frame_type/capture_mode atual")

            cams_to_open = [cam for cam in self.state.cameras if cam.id in required_ids]

            cams_to_open.sort(
                key=lambda cam: 0 if getattr(cam, "interface", "CSI").upper() == "USB" else 1
            )

            for cam in cams_to_open:
                if cam.id not in required_ids:
                    self.state.set_camera_connected(cam.index, False)
                    continue

                try:
                    runtime = self._open_camera(cam)
                    self.cameras_runtime[cam.id] = runtime

                    self.state.set_camera_connected(
                        cam.index,
                        False,
                        width=cam.width,
                        height=cam.height,
                        bayer_pattern=cam.bayer_pattern,
                        bit_depth=cam.bit_depth,
                    )

                except Exception as e:
                    print(f"[WARN] Falha ao abrir {cam.id} (index={cam.index}, role={cam.role}): {e}")
                    self.state.set_camera_connected(cam.index, False)

            # Aplica controles logo após abrir
            if self.cameras_runtime:
                try:
                    self.apply_controls()
                except Exception as e:
                    print(f"[WARN] Falha ao aplicar controles iniciais: {e}")

            for cam_id in list(self.cameras_runtime.keys()):
                self._start_thread(cam_id)

            deadline = time.perf_counter() + 1.0
            while time.perf_counter() < deadline:
                if self.state.camera_count_active > 0:
                    break
                time.sleep(0.02)

            self.initialized = len(self.cameras_runtime) > 0
            self._reconfigure_needed = False

        return self.initialized

    def _open_camera(self, cam):
        interface = getattr(cam, "interface", "CSI").upper()

        if interface == "CSI":
            return self._open_csi_camera(cam)

        if interface == "USB":
            return self._open_usb_camera(cam)

        raise RuntimeError(f"Interface de câmera não suportada: {interface}")

    def _open_csi_camera(self, cam):
        runtime = self._init_camera_runtime(cam)

        picam2 = Picamera2(camera_num=cam.index)

        config = picam2.create_video_configuration(
            #main={"size": (640, 480), "format": "RGB888"},
            raw={"size": (cam.width, cam.height)},
            buffer_count=6
        )

        picam2.configure(config)
        picam2.start()

        runtime["backend"] = "picamera2"
        runtime["picam2"] = picam2
        return runtime

    def _open_usb_camera(self, cam):
        runtime = self._init_camera_runtime(cam)

        backend_name = str(getattr(cam, "usb_backend", "V4L2")).upper()
        api_preference = cv2.CAP_V4L2 if backend_name == "V4L2" else cv2.CAP_ANY

        source = self._resolve_usb_video_path(cam)
        if source:
            cap = cv2.VideoCapture(source, api_preference)
        else:
            cap = cv2.VideoCapture(cam.index, api_preference)

        if not cap.isOpened():
            # fallback no índice se o path falhar
            if source:
                cap.release()
                cap = cv2.VideoCapture(cam.index, api_preference)

        if not cap.isOpened():
            raise RuntimeError(
                f"Falha ao abrir câmera USB. device_path={getattr(cam, 'device_path', None)} index={cam.index}"
            )

        # Configuração desejada
        cap.set(cv2.CAP_PROP_FRAME_WIDTH, cam.width)
        cap.set(cv2.CAP_PROP_FRAME_HEIGHT, cam.height)
        cap.set(cv2.CAP_PROP_FPS, float(self.state.fps))

        # Tenta reduzir buffer
        try:
            cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)
        except Exception:
            pass

        # Tenta MJPG, se a câmera suportar
        try:
            fourcc = cv2.VideoWriter_fourcc(*"MJPG")
            cap.set(cv2.CAP_PROP_FOURCC, fourcc)
        except Exception:
            pass

        # Pequeno warmup
        time.sleep(0.25)

        frame = None
        ok = False
        for i in range(15):
            ok, frame = cap.read()
            if ok and frame is not None:
                break
            time.sleep(0.05)

        if not ok or frame is None:
            cap.release()
            raise RuntimeError(
                f"Falha ao capturar frame inicial da câmera USB "
                f"(device_path={getattr(cam, 'device_path', None)}, index={cam.index})"
            )

        actual_h, actual_w = frame.shape[:2]
        channels = frame.shape[2] if frame.ndim == 3 else 1

        cam.width = int(actual_w)
        cam.height = int(actual_h)

        initial_ts = time.perf_counter()
        runtime["backend"] = "opencv"
        runtime["cap"] = cap
        runtime["buffer"] = frame.copy()
        runtime["last_frame"] = runtime["buffer"]
        runtime["frame_id"] = 0
        runtime["last_signature"] = self._make_frame_signature(frame)
        runtime["last_read_ts"] = initial_ts
        runtime["last_new_frame_ts"] = initial_ts
        runtime["frame_ts"] = initial_ts

        print(
            f"[INFO] USB {cam.id} aberta: src={getattr(cam, 'device_path', None) or cam.index} "
            f"{actual_w}x{actual_h}, channels={channels}, backend={backend_name}"
        )

        return runtime

    def _resolve_usb_video_path(self, cam):
        # 1. Se já veio um caminho explícito e ele existe, usa
        device_path = getattr(cam, "device_path", None)
        if device_path and Path(device_path).exists():
            return str(Path(device_path).resolve()) if Path(device_path).is_symlink() else device_path

        # 2. Tenta by-id
        by_id_dir = Path("/dev/v4l/by-id")
        if by_id_dir.exists():
            candidates = sorted(by_id_dir.glob("*video-index0"))
            if candidates:
                # Para 1 webcam USB, o primeiro já costuma resolver
                return str(candidates[0])

        # 3. Tenta by-path
        by_path_dir = Path("/dev/v4l/by-path")
        if by_path_dir.exists():
            candidates = sorted(by_path_dir.glob("*video-index0"))
            if candidates:
                return str(candidates[0])

        # 4. Fallback bruto
        return None

    # =========================================================
    # Threads de captura
    # =========================================================

    def _start_thread(self, cam_id):
        runtime = self.cameras_runtime[cam_id]
        runtime["stop_event"].clear()
        runtime["consecutive_failures"] = 0

        t = threading.Thread(
            target=self._update_loop,
            args=(cam_id,),
            daemon=True
        )
        runtime["thread"] = t
        t.start()

    def _update_loop(self, cam_id):
        runtime = self.cameras_runtime[cam_id]
        backend = runtime["backend"]

        while not runtime["stop_event"].is_set():
            try:
                if backend == "picamera2":
                    self._update_loop_picamera2(runtime)
                elif backend == "opencv":
                    self._update_loop_opencv(runtime)
                else:
                    raise RuntimeError(f"Backend desconhecido: {backend}")

                runtime["consecutive_failures"] = 0

            except Exception as e:
                runtime["consecutive_failures"] += 1
                print(f"[ERRO LOOP {cam_id}/{backend}] {e}")

                if runtime["consecutive_failures"] >= 5:
                    self.state.set_camera_connected(runtime["camera_index"], False)
                    runtime["stop_event"].set()
                    print(f"[WARN] Desativando {cam_id} após falhas consecutivas no loop")
                    break

                time.sleep(0.05)

    def _update_loop_picamera2(self, runtime):
        request = None
        try:
            picam2 = runtime["picam2"]
            request = picam2.capture_request()

            raw = request.make_array("raw")

            with runtime["frame_lock"]:
                first_valid_frame = runtime["frame_id"] == 0

                if (
                    runtime["buffer"] is None or
                    runtime["buffer"].shape != raw.shape or
                    runtime["buffer"].dtype != raw.dtype
                ):
                    runtime["buffer"] = raw.copy()
                else:
                    np.copyto(runtime["buffer"], raw)

                read_ts = time.perf_counter()
                runtime["last_frame"] = runtime["buffer"]
                runtime["frame_id"] += 1
                runtime["frame_ts"] = read_ts
                runtime["last_read_ts"] = read_ts
                runtime["last_new_frame_ts"] = read_ts

            if first_valid_frame:
                self.state.set_camera_connected(
                    runtime["camera_index"],
                    True,
                    width=raw.shape[1],
                    height=raw.shape[0],
                )

        finally:
            if request is not None:
                try:
                    request.release()
                except Exception:
                    pass

    def _update_loop_opencv(self, runtime):
        cap = runtime["cap"]

        target_fps = max(1.0, float(self.state.fps or 10))
        min_period = 1.0 / target_fps
        t_start = time.perf_counter()

        ok, frame = cap.read()
        read_ts = time.perf_counter()

        if not ok or frame is None:
            raise RuntimeError("Falha ao ler frame da câmera USB")

        signature = self._make_frame_signature(frame)

        with runtime["frame_lock"]:
            first_valid_frame = runtime["frame_id"] == 0
            runtime["last_read_ts"] = read_ts

            if signature != runtime.get("last_signature"):
                if (
                    runtime["buffer"] is None or
                    runtime["buffer"].shape != frame.shape or
                    runtime["buffer"].dtype != frame.dtype
                ):
                    runtime["buffer"] = frame.copy()
                else:
                    np.copyto(runtime["buffer"], frame)

                runtime["last_frame"] = runtime["buffer"]
                runtime["frame_id"] += 1
                runtime["frame_ts"] = read_ts
                runtime["last_new_frame_ts"] = read_ts
                runtime["last_signature"] = signature

        if first_valid_frame:
            self.state.set_camera_connected(
                runtime["camera_index"],
                True,
                width=frame.shape[1],
                height=frame.shape[0],
                bit_depth=8,
            )

        dt = time.perf_counter() - t_start
        sleep_s = min_period - dt
        if sleep_s > 0:
            time.sleep(sleep_s)

    # =========================================================
    # Leitura consolidada
    # =========================================================

    def capture_raw_frames(self):
        result = {}

        for cam_id, runtime in self.cameras_runtime.items():
            with runtime["frame_lock"]:
                if runtime["last_frame"] is None:
                    continue

                frame = runtime["last_frame"].copy()
                h, w = frame.shape[:2]
                channels = frame.shape[2] if frame.ndim == 3 else 1

                result[cam_id] = (
                    frame,
                    w,
                    h,
                    channels,
                    runtime["frame_id"],
                    runtime["frame_ts"],
                    runtime["last_read_ts"]
                )

        return result

    # =========================================================
    # Controles
    # =========================================================

    def apply_controls(self):
        with self.camera_lock:
            for cam_id, runtime in self.cameras_runtime.items():
                backend = runtime.get("backend")

                if backend == "picamera2":
                    self._apply_controls_picamera2(runtime)

                elif backend == "opencv":
                    self._apply_controls_opencv(runtime)

        return True

    def _apply_controls_picamera2(self, runtime):
        picam2 = runtime.get("picam2")
        if picam2 is None:
            return

        controls = {}
        frame_us = int(1_000_000 / max(1, self.state.fps or 10))

        controls["FrameDurationLimits"] = (frame_us, frame_us)
        controls["AeEnable"] = bool(self.state.ae_enable)
        controls["AwbEnable"] = bool(self.state.awb_enable)

        if not self.state.ae_enable:
            if self.state.exposure_time_us is not None:
                controls["ExposureTime"] = int(self.state.exposure_time_us)
            if self.state.analogue_gain is not None:
                controls["AnalogueGain"] = float(self.state.analogue_gain)

        # Só faz sentido real para câmera colorida no backend Picamera2.
        # Como RE/NIR são mono, normalmente AWB/ColourGains serão ignorados.
        if not self.state.awb_enable and self.state.colour_gains is not None:
            r_gain, b_gain = self.state.colour_gains
            controls["ColourGains"] = (float(r_gain), float(b_gain))

        try:
            picam2.set_controls(controls)
        except Exception as e:
            print(f"[ERRO CONTROLS PICAM2] {e} | controls={controls}")

    def _apply_controls_opencv(self, runtime):
        cap = runtime.get("cap")
        if cap is None:
            return

        # FPS
        try:
            cap.set(cv2.CAP_PROP_FPS, float(self.state.fps))
        except Exception:
            pass

        # Exposição
        if self.state.exposure_time_us is not None:
            try:
                # Nem toda webcam respeita isso; tentativa best-effort
                cap.set(cv2.CAP_PROP_AUTO_EXPOSURE, 0.25)  # manual em muitos backends V4L2
            except Exception:
                pass

            try:
                # OpenCV costuma usar escala dependente do driver.
                # Mantemos como tentativa simples, depois calibramos na prática.
                cap.set(cv2.CAP_PROP_EXPOSURE, float(self.state.exposure_time_us))
            except Exception:
                pass

        # Ganho
        if self.state.analogue_gain is not None:
            try:
                cap.set(cv2.CAP_PROP_GAIN, float(self.state.analogue_gain))
            except Exception:
                pass

        # AWB / WB
        try:
            if self.state.awb_enable:
                cap.set(cv2.CAP_PROP_AUTO_WB, 1)
            else:
                cap.set(cv2.CAP_PROP_AUTO_WB, 0)
        except Exception:
            pass

    # =========================================================
    # Sensor modes
    # =========================================================

    def get_sensor_modes(self):
        if self._sensor_modes_cache is not None:
            return self._sensor_modes_cache

        result = []

        for cam in self.state.cameras:
            if getattr(cam, "interface", "CSI").upper() != "CSI":
                continue

            temp = None
            try:
                temp = Picamera2(camera_num=cam.index)
                modes = temp.sensor_modes

                cam_modes = []
                for i, m in enumerate(modes):
                    cam_modes.append({
                        "camera_id": cam.id,
                        "camera_index": cam.index,
                        "role": cam.role,
                        "interface": cam.interface,
                        "mode_index": i,
                        "format": str(m.get("format")) if m.get("format") is not None else None,
                        "size": list(m.get("size")) if m.get("size") is not None else None,
                        "bit_depth": m.get("bit_depth"),
                        "fps": m.get("fps"),
                        "crop_limits": list(m.get("crop_limits")) if m.get("crop_limits") is not None else None,
                        "exposure_limits": list(m.get("exposure_limits")) if m.get("exposure_limits") is not None else None,
                    })

                result.extend(cam_modes)

            except Exception as e:
                result.append({
                    "camera_id": cam.id,
                    "camera_index": cam.index,
                    "role": cam.role,
                    "interface": cam.interface,
                    "error": str(e),
                })

            finally:
                if temp is not None:
                    try:
                        temp.close()
                    except Exception:
                        pass

        self._sensor_modes_cache = result
        return result

    # =========================================================
    # Encerramento
    # =========================================================

    def stop(self):
        with self.camera_lock:
            for runtime in self.cameras_runtime.values():
                runtime["stop_event"].set()

            for runtime in self.cameras_runtime.values():
                t = runtime.get("thread")
                if t:
                    t.join(timeout=1.5)

            for runtime in self.cameras_runtime.values():
                cam = runtime.get("picam2")
                if cam:
                    try:
                        cam.stop()
                    except Exception:
                        pass
                    try:
                        cam.close()
                    except Exception:
                        pass

                cap = runtime.get("cap")
                if cap:
                    try:
                        cap.release()
                    except Exception:
                        pass

            self.cameras_runtime.clear()
            self.initialized = False
            self._reconfigure_needed = False
            self._sensor_modes_cache = None

            for cam in self.state.cameras:
                self.state.set_camera_connected(cam.index, False)