#!/usr/bin/env python3
"""
companion_agent.py - companion-computer side of solar_hop.lua v2.

Runs on the Raspberry Pi (or any Linux companion) as a systemd service.
Responsibilities
  1. Graceful shutdown when the flight controller asks for it, by either
     - a GPIO line driven by the FC relay output (RELAY_COMP_SHDN), or
     - a MAVLink STATUSTEXT "HOP: SHUTDOWN" (belt and braces, no wiring needed).
  2. Optional perch selection: on STATUSTEXT "HOP: REQ_SITE" call select_site(),
     fly to the returned lat/lon in GUIDED, then command LAND. If select_site()
     returns None the script lands in place after its own timeout.
  3. Encrypt every output file at rest with a public key (age). The private key
     lives only on the recovery workstation, so an aircraft found by someone else yields nothing readable.

Dependencies: pymavlink, gpiozero (optional), age (https://github.com/FiloSottile/age)
Config via environment: MAV_PORT (default /dev/serial0), MAV_BAUD (921600),
SHDN_GPIO (BCM pin, default 17), AGE_RECIPIENT (public key), DATA_DIR.
"""
import os, sys, time, subprocess, threading, logging
from pathlib import Path

from pymavlink import mavutil

log = logging.getLogger("hop-agent")
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")

MAV_PORT      = os.environ.get("MAV_PORT", "/dev/serial0")
MAV_BAUD      = int(os.environ.get("MAV_BAUD", "921600"))
SHDN_GPIO     = int(os.environ.get("SHDN_GPIO", "17"))
AGE_RECIPIENT = os.environ.get("AGE_RECIPIENT")          # age1... public key
DATA_DIR      = Path(os.environ.get("DATA_DIR", "/data/hop"))
LAND_ALT_M    = 15.0

# ----------------------------------------------------------------------------
# Hooks you implement
# ----------------------------------------------------------------------------
def select_site(current_lat, current_lon):
    """Return (lat, lon) of a safe perch within ~1 km, or None.
    Plug the onboard model (e.g. Hailo flat/clear-surface classifier on the
    downward camera) in here. Must return quickly; the FC times out."""
    return None


def encrypt_outputs():
    """Encrypt any unencrypted files in DATA_DIR in place with the public key."""
    if not AGE_RECIPIENT:
        log.warning("AGE_RECIPIENT not set; outputs left in clear")
        return
    for f in DATA_DIR.rglob("*"):
        if f.is_file() and f.suffix != ".age":
            out = f.with_suffix(f.suffix + ".age")
            r = subprocess.run(["age", "-r", AGE_RECIPIENT, "-o", str(out), str(f)])
            if r.returncode == 0:
                f.unlink()
            else:
                log.error("age failed on %s", f)


# ----------------------------------------------------------------------------
def graceful_shutdown(reason):
    log.info("shutdown requested: %s", reason)
    try:
        encrypt_outputs()
        subprocess.run(["sync"])
    finally:
        subprocess.run(["systemctl", "poweroff"])


def gpio_watch():
    try:
        from gpiozero import Button
    except Exception as e:   # gpiozero absent or not a Pi
        log.warning("GPIO watch disabled: %s", e)
        return
    line = Button(SHDN_GPIO, pull_up=False, bounce_time=0.2)   # FC relay drives 3.3 V high
    line.wait_for_press()
    graceful_shutdown("GPIO")


def guided_goto(m, lat, lon, alt_rel):
    m.mav.set_position_target_global_int_send(
        0, m.target_system, m.target_component,
        mavutil.mavlink.MAV_FRAME_GLOBAL_RELATIVE_ALT_INT,
        0b0000111111111000,          # position only
        int(lat * 1e7), int(lon * 1e7), alt_rel,
        0, 0, 0, 0, 0, 0, 0, 0)


def set_mode(m, name):
    mode_id = m.mode_mapping()[name]
    m.mav.set_mode_send(m.target_system, mavutil.mavlink.MAV_MODE_FLAG_CUSTOM_MODE_ENABLED, mode_id)


def handle_req_site(m):
    pos = m.recv_match(type="GLOBAL_POSITION_INT", blocking=True, timeout=5)
    if pos is None:
        return
    site = select_site(pos.lat / 1e7, pos.lon / 1e7)
    if site is None:
        log.info("no site; FC will land in place")
        return
    lat, lon = site
    log.info("site %.6f %.6f", lat, lon)
    t0 = time.time()
    while time.time() - t0 < 150:
        guided_goto(m, lat, lon, pos.relative_alt / 1000.0)
        p = m.recv_match(type="GLOBAL_POSITION_INT", blocking=True, timeout=2)
        if p and abs(p.lat / 1e7 - lat) < 3e-5 and abs(p.lon / 1e7 - lon) < 4e-5:
            break
        time.sleep(1)
    set_mode(m, "LAND")


def main():
    DATA_DIR.mkdir(parents=True, exist_ok=True)
    threading.Thread(target=gpio_watch, daemon=True).start()
    m = mavutil.mavlink_connection(MAV_PORT, baud=MAV_BAUD)
    m.wait_heartbeat()
    log.info("heartbeat from sys %d", m.target_system)
    while True:
        msg = m.recv_match(type="STATUSTEXT", blocking=True, timeout=30)
        if msg is None:
            continue
        text = msg.text.strip()
        if text.startswith("HOP: SHUTDOWN"):
            graceful_shutdown("STATUSTEXT")
        elif text.startswith("HOP: REQ_SITE"):
            handle_req_site(m)


if __name__ == "__main__":
    main()