from __future__ import annotations

import json
import os
import shutil
import threading
import time
from dataclasses import asdict, dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional, Union

try:
    import cv2
except ImportError:  # pragma: no cover
    cv2 = None

try:
    from ultralytics import YOLO
except ImportError:  # pragma: no cover
    YOLO = None


@dataclass
class DetectionConfig:
    enabled: bool = False
    stream_url: str = ""
    model_path: str = ""
    target_labels: List[str] = field(default_factory=list)
    output_dir: str = "~/savedetect"
    detect_interval: int = 5
    confidence_threshold: float = 0.5
    max_saved_images: int = 1000
    auto_start: bool = False
    stop_on_flight: bool = True
    status: str = "stopped"
    last_alert_time: Optional[str] = None
    alert_count: int = 0
    frame_count: int = 0
    error: Optional[str] = None


class DetectionManager:
    """YOLO object detection manager for RTSP streams."""

    MODELS_DIR: Path = Path(__file__).parent / "models"
    ALERT_LOG_FILE: Path = Path(__file__).parent / "detection_alerts.jsonl"

    def __init__(self, config_file: Optional[Union[str, Path]] = None) -> None:
        """Initialize the detection manager.

        Args:
            config_file: Path to the JSON configuration file. Defaults to
                ``detection_config.json`` in the same directory as this module.
        """
        if config_file is None:
            config_file = Path(__file__).parent / "detection_config.json"
        self._config_path: Path = Path(config_file)
        self._config: DetectionConfig = self._load_config()
        self._thread: Optional[threading.Thread] = None
        self._stop_event: threading.Event = threading.Event()
        self._alert_queue: List[Dict[str, Any]] = []
        self._alert_lock: threading.Lock = threading.Lock()
        self._last_flush_time: float = 0.0
        self._flight_monitor_thread: Optional[threading.Thread] = None
        self._flight_monitor_stop: threading.Event = threading.Event()
        self._model: Optional[Any] = None
        self._cap: Optional[Any] = None
        self._lock: threading.Lock = threading.Lock()

    # ---------- config persistence ----------

    def _load_config(self) -> DetectionConfig:
        """Load configuration from disk or return defaults."""
        if self._config_path.exists():
            try:
                with open(self._config_path, "r", encoding="utf-8") as f:
                    data = json.load(f)
                return DetectionConfig(**data)
            except Exception as exc:
                print(f"[WARN] Failed to load detection config: {exc}")
        return DetectionConfig()

    def _save_config(self) -> None:
        """Persist current configuration to disk."""
        try:
            self._config_path.parent.mkdir(parents=True, exist_ok=True)
            with open(self._config_path, "w", encoding="utf-8") as f:
                json.dump(asdict(self._config), f, ensure_ascii=False, indent=2)
                os.fsync(f.fileno())
            os.sync()
        except Exception as exc:
            print(f"[WARN] Failed to save detection config: {exc}")

    # ---------- thread priority ----------

    @staticmethod
    def _set_thread_priority() -> None:
        """Lower thread priority to avoid interfering with flight control."""
        try:
            os.nice(19)
        except Exception:
            pass
        try:
            os.setpriority(os.PRIO_PROCESS, 0, 19)
        except Exception:
            pass

    # ---------- flight safety ----------

    @staticmethod
    def _is_flying() -> bool:
        """Return True if the drone is armed or in flight."""
        return Path("/tmp/arm.lock").exists() or Path("/tmp/flight.lock").exists()

    # ---------- image cleanup ----------

    def _cleanup_old_images(self) -> None:
        """Remove oldest images when the saved count exceeds the limit."""
        output_dir = Path(self._config.output_dir).expanduser()
        if not output_dir.exists():
            return
        try:
            images = sorted(
                output_dir.glob("*.jpg"),
                key=lambda p: p.stat().st_mtime,
            )
            excess = len(images) - self._config.max_saved_images
            if excess > 0:
                for img in images[:excess]:
                    img.unlink(missing_ok=True)
        except Exception:
            pass

    # ---------- alert batching ----------

    def _enqueue_alert(self, alert: Dict[str, Any]) -> None:
        """Queue an alert and attempt a debounced flush."""
        with self._alert_lock:
            self._alert_queue.append(alert)
            self._config.alert_count += 1
            self._config.last_alert_time = datetime.utcnow().isoformat()
        self._flush_alerts()

    def _flush_alerts(self, force: bool = False) -> None:
        """Write queued alerts to the JSONL log if the debounce window has passed."""
        now = time.time()
        if not force and now - self._last_flush_time < 5.0:
            return
        with self._alert_lock:
            if not self._alert_queue:
                return
            batch = self._alert_queue[:]
            self._alert_queue.clear()
        try:
            self.ALERT_LOG_FILE.parent.mkdir(parents=True, exist_ok=True)
            with open(self.ALERT_LOG_FILE, "a", encoding="utf-8") as f:
                for alert in batch:
                    f.write(json.dumps(alert, ensure_ascii=False, default=str) + "\n")
                os.fsync(f.fileno())
            os.sync()
            self._last_flush_time = now
        except Exception as exc:
            print(f"[WARN] Failed to flush alerts: {exc}")

    # ---------- detection loop ----------

    def _run_detection(self) -> None:
        """Main detection worker loop."""
        self._set_thread_priority()

        with self._lock:
            self._config.status = "running"
            self._config.error = None
            self._save_config()

        if cv2 is None:
            with self._lock:
                self._config.status = "error"
                self._config.error = "opencv-python (cv2) not available"
                self._save_config()
            return

        if YOLO is None:
            with self._lock:
                self._config.status = "error"
                self._config.error = "ultralytics not available"
                self._save_config()
            return

        output_dir = Path(self._config.output_dir).expanduser()
        output_dir.mkdir(parents=True, exist_ok=True)

        self._cap = cv2.VideoCapture(self._config.stream_url)
        if not self._cap.isOpened():
            with self._lock:
                self._config.status = "error"
                self._config.error = f"Failed to open stream: {self._config.stream_url}"
                self._save_config()
            return

        model_path = self._config.model_path
        if not model_path:
            models = self.list_models()
            if models:
                model_path = str(self.MODELS_DIR / models[0])
        if not model_path or not Path(model_path).exists():
            with self._lock:
                self._config.status = "error"
                self._config.error = f"Model not found: {model_path}"
                self._save_config()
            self._cap.release()
            self._cap = None
            return

        try:
            self._model = YOLO(model_path)
        except Exception as exc:
            with self._lock:
                self._config.status = "error"
                self._config.error = f"Failed to load model: {exc}"
                self._save_config()
            self._cap.release()
            self._cap = None
            return

        frame_counter = 0
        last_flush_check = time.time()

        try:
            while not self._stop_event.is_set():
                if self._config.stop_on_flight and self._is_flying():
                    time.sleep(1.0)
                    continue

                ret, frame = self._cap.read()
                if not ret:
                    time.sleep(0.5)
                    continue

                frame_counter += 1
                with self._lock:
                    self._config.frame_count += 1

                if frame_counter % self._config.detect_interval != 0:
                    continue

                try:
                    results = self._model(frame, verbose=False)
                except Exception:
                    continue

                for result in results:
                    if result.boxes is None:
                        continue
                    for box in result.boxes:
                        cls_id = int(box.cls[0]) if box.cls is not None else -1
                        conf = float(box.conf[0]) if box.conf is not None else 0.0
                        names = getattr(result, "names", {})
                        label = (
                            names.get(cls_id, str(cls_id))
                            if isinstance(names, dict)
                            else str(cls_id)
                        )

                        if conf < self._config.confidence_threshold:
                            continue

                        if (
                            self._config.target_labels
                            and label not in self._config.target_labels
                        ):
                            continue

                        timestamp = datetime.utcnow().isoformat()
                        safe_ts = timestamp.replace(":", "-")
                        filename = f"{safe_ts}_{label}_{conf:.2f}.jpg"
                        save_path = output_dir / filename
                        try:
                            cv2.imwrite(str(save_path), frame)
                        except Exception:
                            pass

                        self._cleanup_old_images()

                        alert = {
                            "timestamp": timestamp,
                            "label": label,
                            "confidence": round(conf, 4),
                            "image_path": str(save_path),
                            "frame_count": self._config.frame_count,
                        }
                        self._enqueue_alert(alert)

                now = time.time()
                if now - last_flush_check >= 1.0:
                    self._flush_alerts()
                    last_flush_check = now

        except Exception as exc:
            with self._lock:
                self._config.status = "error"
                self._config.error = str(exc)
                self._save_config()
        finally:
            if self._cap is not None:
                self._cap.release()
                self._cap = None
            self._model = None
            with self._lock:
                if self._config.status != "error":
                    self._config.status = "stopped"
                self._save_config()
            self._flush_alerts(force=True)

    # ---------- public control API ----------

    def start(self) -> None:
        """Start the background detection thread."""
        with self._lock:
            if self._thread is not None and self._thread.is_alive():
                return
            self._stop_event.clear()
            self._config.enabled = True
            self._save_config()
            thread = threading.Thread(target=self._run_detection, daemon=True)
            self._thread = thread
        thread.start()

    def stop(self) -> None:
        """Signal the detection thread to stop and wait for it to finish."""
        self._stop_event.set()
        with self._lock:
            if self._thread is not None:
                self._thread.join(timeout=10.0)
                self._thread = None
            self._config.enabled = False
            self._config.status = "stopped"
            self._save_config()
        self._flush_alerts(force=True)

    def start_auto(self) -> None:
        """Enable auto-start and begin detection if enabled."""
        with self._lock:
            self._config.auto_start = True
            self._save_config()
        if self._config.enabled:
            self.start()

    def start_flight_monitor(self) -> None:
        """Start a background thread that pauses detection during flight."""
        with self._lock:
            if (
                self._flight_monitor_thread is not None
                and self._flight_monitor_thread.is_alive()
            ):
                return
            self._flight_monitor_stop.clear()

        def _monitor() -> None:
            was_flying = False
            while not self._flight_monitor_stop.is_set():
                is_flying = self._is_flying()
                if is_flying and not was_flying:
                    if self._config.stop_on_flight:
                        self.stop()
                elif not is_flying and was_flying:
                    if self._config.auto_start:
                        self.start()
                was_flying = is_flying
                time.sleep(2.0)

        thread = threading.Thread(target=_monitor, daemon=True)
        thread.start()
        with self._lock:
            self._flight_monitor_thread = thread

    def stop_flight_monitor(self) -> None:
        """Signal the flight monitor thread to stop."""
        self._flight_monitor_stop.set()
        with self._lock:
            if self._flight_monitor_thread is not None:
                self._flight_monitor_thread.join(timeout=5.0)
                self._flight_monitor_thread = None

    # ---------- settings API ----------

    def get_settings(self) -> Dict[str, Any]:
        """Return the current configuration as a dictionary."""
        with self._lock:
            return asdict(self._config)

    def update_settings(self, settings: Dict[str, Any]) -> None:
        """Update configuration fields from the provided dictionary."""
        with self._lock:
            for key, value in settings.items():
                if hasattr(self._config, key):
                    setattr(self._config, key, value)
            self._save_config()

    # ---------- model management ----------

    def list_models(self) -> List[str]:
        """Return a sorted list of ``.pt`` model filenames in the models directory."""
        if not self.MODELS_DIR.exists():
            return []
        return sorted([p.name for p in self.MODELS_DIR.glob("*.pt")])

    def upload_model(
        self, source_path: Union[str, Path], filename: Optional[str] = None
    ) -> str:
        """Copy a model file into the managed models directory.

        Args:
            source_path: Path to the source ``.pt`` file.
            filename: Optional destination filename. Defaults to the source basename.

        Returns:
            Absolute path to the copied model file.
        """
        src = Path(source_path)
        if not src.exists():
            raise FileNotFoundError(f"Source model not found: {source_path}")
        self.MODELS_DIR.mkdir(parents=True, exist_ok=True)
        dest_name = filename if filename is not None else src.name
        dest = self.MODELS_DIR / dest_name
        shutil.copy2(str(src), str(dest))
        return str(dest)

    def delete_model(self, filename: str) -> bool:
        """Remove a model file from the models directory.

        Args:
            filename: Name of the model file to delete.

        Returns:
            True if the file existed and was removed, False otherwise.
        """
        target = self.MODELS_DIR / filename
        if target.exists():
            target.unlink()
            return True
        return False
