"""
slam_toolbox_persistence.py — SLAM Toolbox map serialization and persistence.

Responsibilities:
  1. Auto-save SLAM Toolbox map to ~/saltybot-data/maps/ with timestamp (5 min intervals)
  2. Auto-load most recent map on startup if available
  3. Serialize/deserialize via slam_toolbox built-in ROS services
  4. Export maps to Nav2-compatible YAML+PGM format on demand
  5. Manage map files and metadata

Map lifecycle with slam_toolbox:
  - slam_toolbox publishes /map (OccupancyGrid) during mapping
  - This node calls /slam_toolbox/serialize to save the map state
  - On startup, auto-loads the most recent serialized map via /slam_toolbox/deserialize
  - Maps are stored as .posegraph (binary serialization format)

ROS2 services exposed:
  /saltybot/save_map         (saltybot_mapping/SaveMap) — Manual save with name
  /saltybot/load_map         (saltybot_mapping/LoadMap) — Manual load by name
  /saltybot/list_maps        (saltybot_mapping/ListMaps) — List available maps
  /saltybot/export_occupancy (saltybot_mapping/ExportOccupancy) — Export to PGM+YAML

ROS2 topics published:
  /mapping/slam_status       (std_msgs/String, JSON) — Map metadata & status
"""

import json
import time
import os
import shutil
from pathlib import Path
from datetime import datetime
from typing import Optional, Tuple

import rclpy
from rclpy.node import Node
from rclpy.qos import QoSProfile, ReliabilityPolicy, HistoryPolicy, DurabilityPolicy

from std_msgs.msg import String
from nav_msgs.msg import OccupancyGrid
from slam_toolbox.srv import SerializePose, DeserializePose

from saltybot_mapping.srv import SaveMap, LoadMap, ListMaps, ExportOccupancy
from .map_exporter import write_occupancy_export


class SlamToolboxPersistenceNode(Node):
    """Manage SLAM Toolbox map serialization and persistence."""

    def __init__(self):
        super().__init__('slam_toolbox_persistence')

        # ── Parameters ──────────────────────────────────────────────────────
        self.declare_parameter('maps_dir', str(Path.home() / 'saltybot-data' / 'maps'))
        self.declare_parameter('exports_dir', str(Path.home() / 'saltybot-data' / 'maps' / 'exports'))
        self.declare_parameter('autosave_interval', 300.0)  # 5 minutes
        self.declare_parameter('autoload_on_startup', True)
        self.declare_parameter('keep_autosaves_n', 5)

        self._maps_dir = Path(self.get_parameter('maps_dir').value)
        self._exports_dir = Path(self.get_parameter('exports_dir').value)
        self._autosave_interval = self.get_parameter('autosave_interval').value
        self._autoload_on_startup = self.get_parameter('autoload_on_startup').value
        self._keep_autosaves_n = self.get_parameter('keep_autosaves_n').value

        # Create directories if they don't exist
        self._maps_dir.mkdir(parents=True, exist_ok=True)
        self._exports_dir.mkdir(parents=True, exist_ok=True)

        # State tracking
        self._last_occupancy_msg: Optional[OccupancyGrid] = None
        self._last_autosave_time = time.time()
        self._slam_ready = False

        # ── Wait for slam_toolbox services ──────────────────────────────────
        self.get_logger().info('Waiting for slam_toolbox services...')
        self._serialize_client = self.create_client(SerializePose, '/slam_toolbox/serialize_map')
        self._deserialize_client = self.create_client(DeserializePose, '/slam_toolbox/deserialize_map')

        # Wait with timeout
        if not self._serialize_client.wait_for_service(timeout_sec=30.0):
            self.get_logger().error('slam_toolbox/serialize service not available!')
        if not self._deserialize_client.wait_for_service(timeout_sec=30.0):
            self.get_logger().error('slam_toolbox/deserialize service not available!')

        self._slam_ready = True
        self.get_logger().info(f'slam_toolbox services ready. maps_dir={self._maps_dir}')

        # ── Subscriptions ────────────────────────────────────────────────────
        transient_local = QoSProfile(
            reliability=ReliabilityPolicy.RELIABLE,
            durability=DurabilityPolicy.TRANSIENT_LOCAL,
            history=HistoryPolicy.KEEP_LAST,
            depth=1,
        )

        self.create_subscription(
            OccupancyGrid, '/map',
            self._on_occupancy_grid, transient_local)

        # ── Publishers ───────────────────────────────────────────────────────
        self._pub_status = self.create_publisher(String, '/mapping/slam_status', 10)

        # ── Services ─────────────────────────────────────────────────────────
        self.create_service(SaveMap, '/saltybot/save_map', self._handle_save_map)
        self.create_service(LoadMap, '/saltybot/load_map', self._handle_load_map)
        self.create_service(ListMaps, '/saltybot/list_maps', self._handle_list_maps)
        self.create_service(ExportOccupancy, '/saltybot/export_occupancy',
                          self._handle_export_occupancy)

        # ── Timers ───────────────────────────────────────────────────────────
        self.create_timer(1.0, self._status_tick)
        self.create_timer(self._autosave_interval, self._autosave_tick)

        # ── Auto-load most recent map on startup ────────────────────────────
        if self._autoload_on_startup:
            recent_map = self._find_most_recent_map()
            if recent_map:
                self.get_logger().info(f'Auto-loading most recent map: {recent_map}')
                self._call_deserialize(recent_map)
            else:
                self.get_logger().info('No saved maps found for auto-load.')

    # ── Subscription callbacks ─────────────────────────────────────────────

    def _on_occupancy_grid(self, msg: OccupancyGrid) -> None:
        """Receive occupancy grid updates from slam_toolbox."""
        self._last_occupancy_msg = msg

    # ── Status tick ───────────────────────────────────────────────────────

    def _status_tick(self) -> None:
        """Publish status every second."""
        maps_list = self._list_maps()
        status = {
            'maps_count': len(maps_list),
            'most_recent': maps_list[0]['name'] if maps_list else None,
            'slam_ready': self._slam_ready,
            'has_occupancy': self._last_occupancy_msg is not None,
            'timestamp': datetime.now().isoformat(),
        }
        msg = String()
        msg.data = json.dumps(status)
        self._pub_status.publish(msg)

    # ── Autosave tick ─────────────────────────────────────────────────────

    def _autosave_tick(self) -> None:
        """Auto-save map every N seconds."""
        if not self._slam_ready or not self._last_occupancy_msg:
            return

        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        name = f'autosave_{timestamp}'

        if self._call_serialize(name):
            self._prune_autosaves()
            self.get_logger().info(f'Auto-saved: {name}')
        else:
            self.get_logger().warn(f'Auto-save failed: {name}')

    # ── Service handlers ───────────────────────────────────────────────────

    def _handle_save_map(self, request: SaveMap.Request,
                        response: SaveMap.Response) -> SaveMap.Response:
        """Handle manual map save request."""
        map_name = request.map_name or self._auto_name()

        if self._call_serialize(map_name):
            path = self._maps_dir / f'{map_name}.posegraph'
            response.success = True
            response.message = f'Map saved: {map_name}'
            response.saved_path = str(path)
            self.get_logger().info(response.message)
        else:
            response.success = False
            response.message = f'Failed to save map: {map_name}'
            response.saved_path = ''
            self.get_logger().warn(response.message)

        return response

    def _handle_load_map(self, request: LoadMap.Request,
                        response: LoadMap.Response) -> LoadMap.Response:
        """Handle map load request."""
        if self._call_deserialize(request.map_name):
            path = self._maps_dir / f'{request.map_name}.posegraph'
            response.success = True
            response.message = f'Map loaded: {request.map_name}'
            response.loaded_path = str(path)
            self.get_logger().info(response.message)
        else:
            response.success = False
            response.message = f'Failed to load map: {request.map_name}'
            response.loaded_path = ''
            self.get_logger().warn(response.message)

        return response

    def _handle_list_maps(self, request: ListMaps.Request,
                         response: ListMaps.Response) -> ListMaps.Response:
        """List all saved maps."""
        maps = self._list_maps()
        response.map_names = [m['name'] for m in maps]
        response.map_paths = [m['path'] for m in maps]
        response.sizes_bytes = [m['size'] for m in maps]
        response.modified_times = [m['mtime'] for m in maps]
        response.count = len(maps)
        return response

    def _handle_export_occupancy(self, request: ExportOccupancy.Request,
                                response: ExportOccupancy.Response) -> ExportOccupancy.Response:
        """Export current occupancy grid to PGM+YAML."""
        if not self._last_occupancy_msg:
            response.success = False
            response.message = 'No occupancy grid available'
            return response

        try:
            out_dir = request.output_dir or str(self._exports_dir)
            pgm_path, yaml_path = write_occupancy_export(
                self._last_occupancy_msg, out_dir, request.map_name or '')
            response.success = True
            response.pgm_path = pgm_path
            response.yaml_path = yaml_path
            response.message = f'Exported to {pgm_path}'
            self.get_logger().info(response.message)
        except Exception as e:
            response.success = False
            response.message = str(e)
            self.get_logger().error(f'Export failed: {e}')

        return response

    # ── slam_toolbox service wrappers ──────────────────────────────────────

    def _call_serialize(self, map_name: str) -> bool:
        """Call slam_toolbox serialize service to save map."""
        if not self._slam_ready:
            return False

        try:
            request = SerializePose.Request()
            request.filename = str(self._maps_dir / map_name)
            future = self._serialize_client.call_async(request)
            # Wait for result with timeout
            rclpy.spin_until_future_complete(self, future, timeout_sec=10.0)
            if future.result() is not None:
                return future.result().success
            return False
        except Exception as e:
            self.get_logger().error(f'Serialize call failed: {e}')
            return False

    def _call_deserialize(self, map_name: str) -> bool:
        """Call slam_toolbox deserialize service to load map."""
        if not self._slam_ready:
            return False

        try:
            request = DeserializePose.Request()
            request.filename = str(self._maps_dir / map_name)
            future = self._deserialize_client.call_async(request)
            # Wait for result with timeout
            rclpy.spin_until_future_complete(self, future, timeout_sec=10.0)
            if future.result() is not None:
                return future.result().success
            return False
        except Exception as e:
            self.get_logger().error(f'Deserialize call failed: {e}')
            return False

    # ── Utility methods ────────────────────────────────────────────────────

    def _list_maps(self) -> list:
        """List all saved maps, sorted by modification time (newest first)."""
        maps = []
        for file in sorted(self._maps_dir.glob('*.posegraph'),
                          key=lambda p: p.stat().st_mtime, reverse=True):
            try:
                maps.append({
                    'name': file.stem,
                    'path': str(file),
                    'size': file.stat().st_size,
                    'mtime': int(file.stat().st_mtime),
                })
            except OSError:
                pass
        return maps

    def _find_most_recent_map(self) -> Optional[str]:
        """Find the most recent map file (by modification time)."""
        maps = self._list_maps()
        return maps[0]['name'] if maps else None

    def _prune_autosaves(self) -> None:
        """Keep only N most recent autosaves."""
        autosaves = sorted(
            self._maps_dir.glob('autosave_*.posegraph'),
            key=lambda p: p.stat().st_mtime,
            reverse=True
        )
        for old_file in autosaves[self._keep_autosaves_n:]:
            try:
                old_file.unlink()
                self.get_logger().debug(f'Pruned old autosave: {old_file.name}')
            except OSError as e:
                self.get_logger().warn(f'Failed to prune {old_file.name}: {e}')

    def _auto_name(self) -> str:
        """Generate auto-name with timestamp."""
        return datetime.now().strftime('map_%Y%m%d_%H%M%S')


# ── Entry point ───────────────────────────────────────────────────────────────

def main(args=None):
    rclpy.init(args=args)
    node = SlamToolboxPersistenceNode()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.shutdown()


if __name__ == '__main__':
    main()
