"""battery_node.py — Battery management for saltybot (Issue #125).

Subscribes to /saltybot/telemetry/battery (JSON from esp32_cmd_node) and:
  - Publishes sensor_msgs/BatteryState on /saltybot/battery
  - Publishes JSON alerts on /saltybot/battery/alert at threshold crossings
  - Reduces speed limit at low SoC via /saltybot/speed_limit (std_msgs/Float32)
  - Commands safe stop at critical SoC via /cmd_vel zero + /saltybot/arm False
  - Logs history to SQLite (default: /var/log/saltybot/battery.db)
  - Detects charging (current_ma < charge_detect_threshold_ma)

Alert levels (SoC thresholds):
  20%  WARNING   — reduce max speed to 60%
  10%  CRITICAL  — reduce max speed to 30%, log prominently
   5%  EMERGENCY — publish zero /cmd_vel, disarm, log + alert

SoC source priority:
<<<<<<< HEAD
  1. soc_pct field from ESP32 BATTERY telemetry (fuel gauge or lookup on ESP32 BALANCE)
=======
  1. soc_pct field from ESP32-S3 BATTERY telemetry (fuel gauge or lookup on ESP32-S3)
>>>>>>> 291dd68 (feat: remove all STM32/Mamba/BlackPill references — ESP32-S3 only)
  2. Voltage-based lookup table (3S LiPo curve) if soc_pct == 0 and voltage known

Parameters (config/battery_params.yaml):
  db_path              /var/log/saltybot/battery.db
  cell_count           3        (number of LiPo cells)
  warn_pct             20.0
  critical_pct         10.0
  emergency_pct         5.0
  warn_speed_factor     0.6
  critical_speed_factor 0.3
  charge_detect_ma     -100     (current_ma below this → charging)
  history_retention_days 7
  publish_rate_hz       1.0     (re-publish rate for /saltybot/battery)
"""

from __future__ import annotations

import json
import math
import os
import sqlite3
import time
from typing import Optional

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

from geometry_msgs.msg import Twist
from sensor_msgs.msg import BatteryState
from std_msgs.msg import Float32, String
from std_srvs.srv import SetBool

# ── SoC voltage lookup table (3S LiPo, per-cell voltage → SoC %) ─────────────
# Based on typical LiPo discharge curve at 1C
_SOC_CURVE_PER_CELL = [
    # (cell_v, soc_pct) — sorted ascending by voltage
    (3.00,  0),
    (3.27,  5),
    (3.37, 10),
    (3.53, 20),
    (3.62, 30),
    (3.68, 40),
    (3.73, 50),
    (3.78, 60),
    (3.83, 70),
    (3.90, 80),
    (4.00, 90),
    (4.20, 100),
]

# Alert level constants
_ALERT_NONE      = "ok"
_ALERT_WARN      = "warning"
_ALERT_CRITICAL  = "critical"
_ALERT_EMERGENCY = "emergency"


def _soc_from_voltage(voltage_v: float, cell_count: int) -> int:
    """Estimate SoC from pack voltage using per-cell LiPo lookup table."""
    if cell_count <= 0:
        return 0
    cell_v = voltage_v / cell_count
    curve = _SOC_CURVE_PER_CELL

    if cell_v <= curve[0][0]:
        return 0
    if cell_v >= curve[-1][0]:
        return 100

    # Linear interpolation between adjacent points
    for i in range(len(curve) - 1):
        v_lo, soc_lo = curve[i]
        v_hi, soc_hi = curve[i + 1]
        if v_lo <= cell_v <= v_hi:
            t = (cell_v - v_lo) / (v_hi - v_lo)
            return int(soc_lo + t * (soc_hi - soc_lo))
    return 0


class BatteryNode(Node):
    """Battery management and monitoring node."""

    def __init__(self) -> None:
        super().__init__("battery_node")

        # ── Parameters ────────────────────────────────────────────────────────
        self.declare_parameter("db_path",               "/var/log/saltybot/battery.db")
        self.declare_parameter("cell_count",            3)
        self.declare_parameter("warn_pct",              20.0)
        self.declare_parameter("critical_pct",          10.0)
        self.declare_parameter("emergency_pct",          5.0)
        self.declare_parameter("warn_speed_factor",      0.6)
        self.declare_parameter("critical_speed_factor",  0.3)
        self.declare_parameter("charge_detect_ma",      -100)
        self.declare_parameter("history_retention_days", 7)
        self.declare_parameter("publish_rate_hz",        1.0)

        self._db_path             = self.get_parameter("db_path").value
        self._cell_count          = self.get_parameter("cell_count").value
        self._warn_pct            = self.get_parameter("warn_pct").value
        self._critical_pct        = self.get_parameter("critical_pct").value
        self._emergency_pct       = self.get_parameter("emergency_pct").value
        self._warn_speed_factor   = self.get_parameter("warn_speed_factor").value
        self._critical_speed_factor = self.get_parameter("critical_speed_factor").value
        self._charge_detect_ma    = self.get_parameter("charge_detect_ma").value
        self._retention_days      = self.get_parameter("history_retention_days").value
        publish_rate              = self.get_parameter("publish_rate_hz").value

        # ── QoS ───────────────────────────────────────────────────────────────
        rel_qos = QoSProfile(
            reliability=ReliabilityPolicy.RELIABLE,
            history=HistoryPolicy.KEEP_LAST, depth=10,
        )

        # ── Publishers ────────────────────────────────────────────────────────
        self._battery_pub     = self.create_publisher(BatteryState, "/saltybot/battery",       rel_qos)
        self._alert_pub       = self.create_publisher(String,       "/saltybot/battery/alert",  rel_qos)
        self._speed_limit_pub = self.create_publisher(Float32,      "/saltybot/speed_limit",    rel_qos)
        self._cmd_vel_pub     = self.create_publisher(Twist,        "/cmd_vel",                 rel_qos)

        # ── Subscribers ───────────────────────────────────────────────────────
        self._telem_sub = self.create_subscription(
            String, "/saltybot/telemetry/battery",
            self._on_battery_telem, rel_qos,
        )

        # ── Arm service client (for emergency stop) ───────────────────────────
        self._arm_client = self.create_client(SetBool, "/saltybot/arm")

        # ── State ─────────────────────────────────────────────────────────────
        self._last_voltage_v  = 0.0
        self._last_current_ma = 0
        self._last_soc_pct    = 100
        self._last_alert      = _ALERT_NONE
        self._charging        = False
        self._emergency_fired = False     # prevent repeated disarm calls

        # ── SQLite history ─────────────────────────────────────────────────────
        self._db: Optional[sqlite3.Connection] = None
        self._init_db()

        # ── Timers ────────────────────────────────────────────────────────────
        self._publish_timer  = self.create_timer(1.0 / publish_rate, self._publish_cb)
        self._prune_timer    = self.create_timer(3600.0, self._prune_history)  # hourly

        self.get_logger().info(
            f"battery_node started — warn={self._warn_pct}% "
            f"critical={self._critical_pct}% "
            f"emergency={self._emergency_pct}% "
            f"cells={self._cell_count}"
        )

    # ── SQLite ────────────────────────────────────────────────────────────────

    def _init_db(self) -> None:
        """Create the SQLite DB and battery_history table if not present."""
        try:
            db_dir = os.path.dirname(self._db_path)
            if db_dir:
                os.makedirs(db_dir, exist_ok=True)
            self._db = sqlite3.connect(
                self._db_path, check_same_thread=False, timeout=5.0
            )
            self._db.execute("""
                CREATE TABLE IF NOT EXISTS battery_history (
                    id          INTEGER PRIMARY KEY AUTOINCREMENT,
                    ts          REAL    NOT NULL,
                    voltage_v   REAL    NOT NULL,
                    current_ma  INTEGER NOT NULL,
                    soc_pct     INTEGER NOT NULL,
                    charging    INTEGER NOT NULL,
                    alert_level TEXT    NOT NULL
                )
            """)
            self._db.execute(
                "CREATE INDEX IF NOT EXISTS idx_ts ON battery_history(ts)"
            )
            self._db.commit()
            self.get_logger().info(f"Battery history DB: {self._db_path}")
        except Exception as exc:
            self.get_logger().warn(f"Cannot open battery DB ({exc}) — logging disabled")
            self._db = None

    def _log_to_db(self, voltage_v: float, current_ma: int,
                   soc_pct: int, charging: bool, alert: str) -> None:
        if self._db is None:
            return
        try:
            self._db.execute(
                "INSERT INTO battery_history(ts,voltage_v,current_ma,soc_pct,charging,alert_level) "
                "VALUES (?,?,?,?,?,?)",
                (time.time(), voltage_v, current_ma, soc_pct, int(charging), alert),
            )
            self._db.commit()
        except Exception as exc:
            self.get_logger().warn(f"Battery DB write error: {exc}")

    def _prune_history(self) -> None:
        """Delete records older than history_retention_days."""
        if self._db is None:
            return
        try:
            cutoff = time.time() - self._retention_days * 86400
            self._db.execute("DELETE FROM battery_history WHERE ts < ?", (cutoff,))
            self._db.commit()
        except Exception as exc:
            self.get_logger().warn(f"Battery DB prune error: {exc}")

    # ── Telemetry receive ─────────────────────────────────────────────────────

    def _on_battery_telem(self, msg: String) -> None:
        """Parse JSON battery telemetry and update internal state."""
        try:
            data = json.loads(msg.data)
            voltage_mv  = int(data["voltage_mv"])
            current_ma  = int(data["current_ma"])
            soc_from_fw = int(data.get("soc_pct", 0))
        except (json.JSONDecodeError, KeyError, ValueError) as exc:
            self.get_logger().warn(f"Bad battery telemetry: {exc}")
            return

        voltage_v = voltage_mv / 1000.0
        charging  = current_ma < self._charge_detect_ma

        # SoC: prefer firmware value; fall back to voltage curve
        if soc_from_fw > 0:
            soc_pct = soc_from_fw
        else:
            soc_pct = _soc_from_voltage(voltage_v, self._cell_count)

        self._last_voltage_v  = voltage_v
        self._last_current_ma = current_ma
        self._last_soc_pct    = soc_pct
        self._charging        = charging

        # Check thresholds and act
        new_alert = self._compute_alert(soc_pct)
        if new_alert != self._last_alert:
            self._handle_alert_change(new_alert, soc_pct, voltage_v)

        # Log to SQLite
        self._log_to_db(voltage_v, current_ma, soc_pct, charging, new_alert)

    # ── Alert handling ────────────────────────────────────────────────────────

    def _compute_alert(self, soc_pct: int) -> str:
        """Return the current alert level string for a given SoC."""
        if soc_pct <= self._emergency_pct:
            return _ALERT_EMERGENCY
        if soc_pct <= self._critical_pct:
            return _ALERT_CRITICAL
        if soc_pct <= self._warn_pct:
            return _ALERT_WARN
        return _ALERT_NONE

    def _handle_alert_change(self, new_alert: str, soc_pct: int, voltage_v: float) -> None:
        """Act on threshold transition."""
        prev = self._last_alert
        self._last_alert = new_alert

        msg_text = (
            f"Battery {new_alert.upper()}: {soc_pct}% SoC / {voltage_v:.2f}V "
            f"(was {prev})"
        )

        if new_alert == _ALERT_NONE:
            # Recovery (e.g., plugged in for charging)
            self.get_logger().info(f"Battery OK: {soc_pct}% SoC")
            self._publish_speed_limit(1.0)

        elif new_alert == _ALERT_WARN:
            self.get_logger().warn(msg_text)
            self._publish_speed_limit(self._warn_speed_factor)

        elif new_alert == _ALERT_CRITICAL:
            self.get_logger().error(msg_text)
            self._publish_speed_limit(self._critical_speed_factor)

        elif new_alert == _ALERT_EMERGENCY:
            self.get_logger().fatal(msg_text)
            self._publish_speed_limit(0.0)
            if not self._emergency_fired:
                self._emergency_fired = True
                self._execute_safe_stop()

        # Publish alert message
        alert_payload = json.dumps({
            "level":       new_alert,
            "soc_pct":     soc_pct,
            "voltage_v":   round(voltage_v, 3),
            "prev_level":  prev,
            "ts":          time.time(),
        })
        alert_msg      = String()
        alert_msg.data = alert_payload
        self._alert_pub.publish(alert_msg)

    def _publish_speed_limit(self, factor: float) -> None:
        """Publish normalised speed limit factor 0.0–1.0."""
        msg       = Float32()
        msg.data  = float(factor)
        self._speed_limit_pub.publish(msg)

    def _execute_safe_stop(self) -> None:
<<<<<<< HEAD
        """Send zero /cmd_vel and disarm the ESP32 BALANCE."""
=======
        """Send zero /cmd_vel and disarm the ESP32-S3."""
>>>>>>> 291dd68 (feat: remove all STM32/Mamba/BlackPill references — ESP32-S3 only)
        self.get_logger().fatal("EMERGENCY: publishing zero /cmd_vel and disarming")
        # Publish zero velocity
        zero_twist = Twist()
        self._cmd_vel_pub.publish(zero_twist)

        # Call /saltybot/arm(False) if service is available
        if self._arm_client.service_is_ready():
            req          = SetBool.Request()
            req.data     = False
            self._arm_client.call_async(req)
        else:
            self.get_logger().warn(
                "/saltybot/arm service not available for emergency disarm"
            )

    # ── Periodic publish ──────────────────────────────────────────────────────

    def _publish_cb(self) -> None:
        """Publish sensor_msgs/BatteryState at configured rate."""
        if self._last_voltage_v == 0.0:
            return    # no telemetry received yet

        msg = BatteryState()
        msg.header.stamp    = self.get_clock().now().to_msg()
        msg.header.frame_id = "battery"

        msg.voltage         = self._last_voltage_v
        msg.current         = self._last_current_ma / 1000.0    # mA → A
        msg.percentage      = self._last_soc_pct / 100.0        # 0.0–1.0

        # Capacity: 3S 5000 mAh pack (configurable in future)
        msg.design_capacity = 5.0   # Ah
        msg.capacity        = float("nan")
        msg.charge          = float("nan")
        msg.temperature     = float("nan")

        if self._charging:
            msg.power_supply_status = BatteryState.POWER_SUPPLY_STATUS_CHARGING
        elif self._last_soc_pct >= 99:
            msg.power_supply_status = BatteryState.POWER_SUPPLY_STATUS_FULL
        else:
            msg.power_supply_status = BatteryState.POWER_SUPPLY_STATUS_DISCHARGING

        msg.power_supply_technology = BatteryState.POWER_SUPPLY_TECHNOLOGY_LIPO
        msg.power_supply_health     = (
            BatteryState.POWER_SUPPLY_HEALTH_GOOD
            if self._last_alert in (_ALERT_NONE, _ALERT_WARN)
            else BatteryState.POWER_SUPPLY_HEALTH_UNSPEC_FAILURE
        )
        msg.present = True

        # Cell voltages (estimated, equal distribution assumption)
        if self._cell_count > 0:
            cell_v = self._last_voltage_v / self._cell_count
            msg.cell_voltage = [cell_v] * self._cell_count

        # Location field: alert level for easy monitoring
        msg.location = self._last_alert

        self._battery_pub.publish(msg)

    # ── Lifecycle ─────────────────────────────────────────────────────────────

    def destroy_node(self) -> None:
        if self._db:
            try:
                self._db.close()
            except Exception:
                pass
        super().destroy_node()


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


if __name__ == "__main__":
    main()
