#!/usr/bin/env python3
"""
OTA Firmware Update Node - Downloads, builds, deploys, and rolls back firmware.

Features:
- Downloads release archives from Gitea (seb/saltylab-firmware)
- Runs colcon build in staging directory
- Swaps symlink to activate new version
- Restarts ROS2 services (systemd)
- Rolls back on build failure
- Tracks versions in ~/saltybot-data/versions.json
- Safety: blocks update if robot is moving (velocity > threshold)
- Triggers via MQTT /saltybot/ota_command or dashboard button
"""

import json
import os
import subprocess
import shutil
import requests
import threading
from datetime import datetime
from pathlib import Path
from typing import Dict, Any, Optional, Tuple
import rclpy
from rclpy.node import Node
from std_msgs.msg import String


class OTAUpdater(Node):
    """OTA firmware update manager for SaltyBot."""

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

        # Configuration
        self.data_dir = Path(os.path.expanduser('~/.saltybot-data'))
        self.data_dir.mkdir(parents=True, exist_ok=True)
        self.versions_file = self.data_dir / 'versions.json'

        self.staging_dir = Path(os.path.expanduser('~/saltybot-ota-staging'))
        self.install_dir = Path(os.path.expanduser('~/saltybot-ros2-install'))

        self.gitea_api = 'https://gitea.vayrette.com/api/v1/repos/seb/saltylab-firmware'
        self.max_velocity_threshold = 0.05  # m/s - block update if moving faster

        # Runtime state
        self.updating = False
        self.current_velocity = 0.0

        self.get_logger().info(f'OTA Updater initialized')

        # Subscriptions
        self.create_subscription(String, '/saltybot/ota_command', self._on_ota_command, 10)
        self.create_subscription(String, '/odom', self._on_odometry, 10)

        # Publisher for status
        self.status_pub = self.create_publisher(String, '/saltybot/ota_status', 10)

    def _on_ota_command(self, msg: String):
        """Handle OTA update request from MQTT or dashboard."""
        try:
            cmd = msg.data.strip()
            if cmd == 'check':
                self._check_for_updates()
            elif cmd.startswith('update:'):
                version = cmd.split(':', 1)[1].strip()
                self._start_update_thread(version)
            elif cmd == 'rollback':
                self._rollback_update()
        except Exception as e:
            self.get_logger().error(f'Error handling OTA command: {e}')

    def _on_odometry(self, msg: String):
        """Track current velocity from odometry for safety checks."""
        try:
            # Parse velocity from odom message (simplified)
            data = json.loads(msg.data)
            vx = data.get('vx', 0.0)
            vy = data.get('vy', 0.0)
            self.current_velocity = (vx**2 + vy**2) ** 0.5
        except:
            pass

    def _check_for_updates(self):
        """Check Gitea for new releases."""
        try:
            response = requests.get(f'{self.gitea_api}/releases', timeout=5)
            response.raise_for_status()
            releases = response.json()

            if releases:
                latest = releases[0]
                version = latest.get('tag_name', 'unknown')
                current = self._get_current_version()

                status = {
                    'timestamp': datetime.now().isoformat(),
                    'status': 'update_available' if version != current else 'up_to_date',
                    'current': current,
                    'latest': version,
                }
            else:
                status = {'status': 'no_releases'}

            self._publish_status(status)
        except Exception as e:
            self.get_logger().error(f'Error checking for updates: {e}')
            self._publish_status({'status': 'check_failed', 'error': str(e)})

    def _start_update_thread(self, version: str):
        """Start update in background thread."""
        if self.updating:
            self._publish_status({'status': 'already_updating'})
            return

        thread = threading.Thread(target=self._update_firmware, args=(version,), daemon=True)
        thread.start()

    def _update_firmware(self, version: str):
        """Execute firmware update: download, build, deploy, with rollback."""
        self.updating = True
        try:
            # Safety check
            if self.current_velocity > self.max_velocity_threshold:
                self._publish_status({'status': 'blocked_robot_moving', 'velocity': self.current_velocity})
                self.updating = False
                return

            self._publish_status({'status': 'starting_update', 'version': version})

            # Step 1: Backup current installation
            backup_dir = self._backup_current_install()

            # Step 2: Download release
            archive_path = self._download_release(version)

            # Step 3: Extract to staging
            self._extract_to_staging(archive_path, version)

            # Step 4: Build with colcon
            if not self._colcon_build():
                self.get_logger().error('Build failed, rolling back')
                self._restore_from_backup(backup_dir)
                self._publish_status({'status': 'build_failed_rolled_back', 'version': version})
                self.updating = False
                return

            # Step 5: Swap symlink
            self._swap_install_symlink(version)

            # Step 6: Restart services
            self._restart_ros_services()

            # Step 7: Update version tracking
            self._update_version_file(version)

            # Cleanup
            shutil.rmtree(backup_dir, ignore_errors=True)

            self._publish_status({'status': 'update_complete', 'version': version, 'timestamp': datetime.now().isoformat()})

        except Exception as e:
            self.get_logger().error(f'Update failed: {e}')
            self._publish_status({'status': 'update_failed', 'error': str(e)})
        finally:
            self.updating = False

    def _backup_current_install(self) -> Path:
        """Backup current installation for rollback."""
        try:
            timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
            backup_dir = self.data_dir / f'backup_{timestamp}'

            if self.install_dir.exists():
                shutil.copytree(self.install_dir, backup_dir, dirs_exist_ok=True)

            self.get_logger().info(f'Backup created: {backup_dir}')
            return backup_dir
        except Exception as e:
            self.get_logger().warning(f'Backup failed: {e}')
            return None

    def _download_release(self, version: str) -> Path:
        """Download release archive from Gitea."""
        try:
            self._publish_status({'status': 'downloading', 'version': version})

            # Get release info
            response = requests.get(f'{self.gitea_api}/releases/tags/{version}', timeout=10)
            response.raise_for_status()
            release = response.json()

            # Download source code archive
            archive_url = release.get('tarball_url')
            if not archive_url:
                raise Exception(f'No tarball URL for release {version}')

            archive_path = self.data_dir / f'saltylab-firmware-{version}.tar.gz'

            response = requests.get(archive_url, timeout=60, stream=True)
            response.raise_for_status()

            with open(archive_path, 'wb') as f:
                for chunk in response.iter_content(chunk_size=8192):
                    if chunk:
                        f.write(chunk)

            self.get_logger().info(f'Downloaded {archive_path}')
            return archive_path

        except Exception as e:
            self.get_logger().error(f'Download failed: {e}')
            raise

    def _extract_to_staging(self, archive_path: Path, version: str):
        """Extract release archive to staging directory."""
        try:
            self._publish_status({'status': 'extracting', 'version': version})

            # Clean staging
            if self.staging_dir.exists():
                shutil.rmtree(self.staging_dir)
            self.staging_dir.mkdir(parents=True)

            # Extract
            subprocess.run(['tar', 'xzf', str(archive_path), '-C', str(self.staging_dir)],
                         check=True, capture_output=True)

            # Move extracted content to correct location
            extracted = list(self.staging_dir.glob('*'))
            if len(extracted) == 1 and extracted[0].is_dir():
                # Rename extracted directory
                src_dir = extracted[0]
                final_dir = self.staging_dir / 'firmware'
                src_dir.rename(final_dir)

            self.get_logger().info(f'Extracted to {self.staging_dir}')

        except Exception as e:
            self.get_logger().error(f'Extract failed: {e}')
            raise

    def _colcon_build(self) -> bool:
        """Build firmware with colcon."""
        try:
            self._publish_status({'status': 'building'})

            build_dir = self.staging_dir / 'firmware' / 'jetson' / 'ros2_ws'

            result = subprocess.run(
                ['colcon', 'build', '--symlink-install'],
                cwd=str(build_dir),
                capture_output=True,
                timeout=3600,  # 1 hour timeout
                text=True
            )

            if result.returncode != 0:
                self.get_logger().error(f'Build failed: {result.stderr}')
                return False

            self.get_logger().info('Build succeeded')
            return True

        except subprocess.TimeoutExpired:
            self.get_logger().error('Build timed out')
            return False
        except Exception as e:
            self.get_logger().error(f'Build error: {e}')
            return False

    def _swap_install_symlink(self, version: str):
        """Swap symlink to activate new installation."""
        try:
            self._publish_status({'status': 'deploying', 'version': version})

            new_install = self.staging_dir / 'firmware' / 'jetson' / 'ros2_ws' / 'install'
            symlink = self.install_dir

            # Remove old symlink
            if symlink.is_symlink():
                symlink.unlink()
            elif symlink.exists():
                shutil.rmtree(symlink)

            # Create new symlink
            symlink.parent.mkdir(parents=True, exist_ok=True)
            symlink.symlink_to(new_install)

            self.get_logger().info(f'Symlink swapped to {new_install}')

        except Exception as e:
            self.get_logger().error(f'Deploy failed: {e}')
            raise

    def _restart_ros_services(self):
        """Restart ROS2 systemd services."""
        try:
            self._publish_status({'status': 'restarting_services'})

            # Restart main ROS2 service
            subprocess.run(['sudo', 'systemctl', 'restart', 'saltybot-ros2'],
                         check=False, capture_output=True)

            self.get_logger().info('ROS2 services restarted')

        except Exception as e:
            self.get_logger().warning(f'Service restart failed: {e}')

    def _update_version_file(self, version: str):
        """Update version tracking file."""
        try:
            versions = {}
            if self.versions_file.exists():
                with open(self.versions_file, 'r') as f:
                    versions = json.load(f)

            versions['current'] = version
            versions['updated'] = datetime.now().isoformat()
            versions['history'] = versions.get('history', [])
            versions['history'].append({
                'version': version,
                'timestamp': datetime.now().isoformat(),
                'status': 'success'
            })

            with open(self.versions_file, 'w') as f:
                json.dump(versions, f, indent=2)

            self.get_logger().info(f'Version file updated: {version}')

        except Exception as e:
            self.get_logger().warning(f'Version file update failed: {e}')

    def _restore_from_backup(self, backup_dir: Path):
        """Restore installation from backup."""
        try:
            if backup_dir and backup_dir.exists():
                if self.install_dir.exists():
                    shutil.rmtree(self.install_dir)

                shutil.copytree(backup_dir, self.install_dir)
                self.get_logger().info(f'Restored from backup')
        except Exception as e:
            self.get_logger().error(f'Restore failed: {e}')

    def _rollback_update(self):
        """Rollback to previous version."""
        try:
            self._publish_status({'status': 'rolling_back'})

            versions = {}
            if self.versions_file.exists():
                with open(self.versions_file, 'r') as f:
                    versions = json.load(f)

            history = versions.get('history', [])
            if len(history) > 1:
                previous = history[-2]['version']
                self._start_update_thread(previous)
            else:
                self._publish_status({'status': 'rollback_unavailable'})

        except Exception as e:
            self.get_logger().error(f'Rollback failed: {e}')
            self._publish_status({'status': 'rollback_failed', 'error': str(e)})

    def _get_current_version(self) -> str:
        """Get currently installed version."""
        try:
            if self.versions_file.exists():
                with open(self.versions_file, 'r') as f:
                    data = json.load(f)
                    return data.get('current', 'unknown')
            return 'unknown'
        except:
            return 'unknown'

    def _publish_status(self, status: Dict[str, Any]):
        """Publish OTA status update."""
        try:
            msg = String()
            msg.data = json.dumps(status)
            self.status_pub.publish(msg)
        except Exception as e:
            self.get_logger().error(f'Status publish failed: {e}')


def main(args=None):
    rclpy.init(args=args)
    node = OTAUpdater()

    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.shutdown()


if __name__ == '__main__':
    main()
