lesavka/server/src/main.rs

//! lesavka-server - gadget cycle guarded by env
// server/src/main.rs
#![forbid(unsafe_code)]

use anyhow::Context as _;
use futures_util::{Stream, StreamExt};
use std::path::Path;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::time::Duration;
use std::{backtrace::Backtrace, panic, pin::Pin, sync::Arc};
use tokio::{fs::OpenOptions, io::AsyncWriteExt, process::Command, sync::Mutex};
use tokio_stream::wrappers::ReceiverStream;
use tonic::transport::Server;
use tonic::{Request, Response, Status};
use tonic_reflection::server::Builder as ReflBuilder;
use tracing::{debug, error, info, trace, warn};
use tracing_appender::non_blocking::WorkerGuard;
use tracing_subscriber::{filter::EnvFilter, fmt, prelude::*};

use lesavka_common::lesavka::{
    AudioPacket, Empty, KeyboardReport, MonitorRequest, MouseReport, PasteReply, PasteRequest,
    ResetUsbReply, VideoPacket,
    relay_server::{Relay, RelayServer},
};

use lesavka_server::{audio, camera, gadget::UsbGadget, handshake::HandshakeSvc, paste, video};

/*──────────────── constants ────────────────*/
const VERSION: &str = env!("CARGO_PKG_VERSION");
const PKG_NAME: &str = env!("CARGO_PKG_NAME");

/*──────────────── logging ───────────────────*/
fn init_tracing() -> anyhow::Result<WorkerGuard> {
    let file = std::fs::OpenOptions::new()
        .create(true)
        .truncate(true)
        .write(true)
        .open("/tmp/lesavka-server.log")?;
    let (file_writer, guard) = tracing_appender::non_blocking(file);

    let env_filter = EnvFilter::try_from_default_env()
        .unwrap_or_else(|_| EnvFilter::new("lesavka_server=info,lesavka_server::video=warn"));
    let filter_str = env_filter.to_string();

    tracing_subscriber::registry()
        .with(env_filter)
        .with(fmt::layer().with_target(true).with_thread_ids(true))
        .with(
            fmt::layer()
                .with_writer(file_writer)
                .with_ansi(false)
                .with_target(true)
                .with_level(true),
        )
        .init();
    tracing::info!("📜 effective RUST_LOG = \"{}\"", filter_str);
    Ok(guard)
}

/*──────────────── helpers ───────────────────*/
async fn open_with_retry(path: &str) -> anyhow::Result<tokio::fs::File> {
    for attempt in 1..=200 {
        // ≈10 s
        match OpenOptions::new()
            .write(true)
            .custom_flags(libc::O_NONBLOCK)
            .open(path)
            .await
        {
            Ok(f) => {
                info!("✅ {path} opened on attempt #{attempt}");
                return Ok(f);
            }
            Err(e) if e.raw_os_error() == Some(libc::EBUSY) => {
                trace!("⏳ {path} busy… retry #{attempt}");
                tokio::time::sleep(Duration::from_millis(50)).await;
            }
            Err(e) => return Err(e).with_context(|| format!("opening {path}")),
        }
    }
    Err(anyhow::anyhow!("timeout waiting for {path}"))
}

fn allow_gadget_cycle() -> bool {
    std::env::var("LESAVKA_ALLOW_GADGET_CYCLE").is_ok()
}

async fn recover_hid_if_needed(
    err: &std::io::Error,
    gadget: UsbGadget,
    kb: Arc<Mutex<tokio::fs::File>>,
    ms: Arc<Mutex<tokio::fs::File>>,
    did_cycle: Arc<AtomicBool>,
) {
    let code = err.raw_os_error();
    let should_recover = matches!(
        code,
        Some(libc::ENOTCONN) | Some(libc::ESHUTDOWN) | Some(libc::EPIPE)
    );
    if !should_recover {
        return;
    }
    if did_cycle
        .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
        .is_err()
    {
        return;
    }

    let allow_cycle = allow_gadget_cycle();
    tokio::spawn(async move {
        if allow_cycle {
            warn!("🔁 HID transport down (errno={code:?}) - cycling gadget");
            match tokio::task::spawn_blocking(move || gadget.cycle()).await {
                Ok(Ok(())) => info!("✅ USB gadget cycle complete (auto-recover)"),
                Ok(Err(e)) => error!("💥 USB gadget cycle failed: {e:#}"),
                Err(e) => error!("💥 USB gadget cycle task panicked: {e:#}"),
            }
        } else {
            warn!(
                "🔒 HID transport down (errno={code:?}) - gadget cycle disabled; set LESAVKA_ALLOW_GADGET_CYCLE=1 to enable"
            );
        }

        if let Err(e) = async {
            let kb_new = open_with_retry("/dev/hidg0").await?;
            let ms_new = open_with_retry("/dev/hidg1").await?;
            *kb.lock().await = kb_new;
            *ms.lock().await = ms_new;
            Ok::<(), anyhow::Error>(())
        }
        .await
        {
            error!("💥 HID reopen failed: {e:#}");
        }

        tokio::time::sleep(Duration::from_secs(2)).await;
        did_cycle.store(false, Ordering::SeqCst);
    });
}

async fn open_voice_with_retry(uac_dev: &str) -> anyhow::Result<audio::Voice> {
    let attempts = std::env::var("LESAVKA_MIC_INIT_ATTEMPTS")
        .ok()
        .and_then(|v| v.parse::<u32>().ok())
        .unwrap_or(5)
        .max(1);
    let delay_ms = std::env::var("LESAVKA_MIC_INIT_DELAY_MS")
        .ok()
        .and_then(|v| v.parse::<u64>().ok())
        .unwrap_or(250);
    let mut last_err: Option<anyhow::Error> = None;
    for attempt in 1..=attempts {
        match audio::Voice::new(uac_dev).await {
            Ok(v) => {
                if attempt > 1 {
                    info!(%uac_dev, attempt, "🎤 microphone sink recovered");
                }
                return Ok(v);
            }
            Err(e) => {
                warn!(%uac_dev, attempt, "⚠️  microphone sink init failed: {e:#}");
                last_err = Some(e);
                tokio::time::sleep(Duration::from_millis(delay_ms)).await;
            }
        }
    }
    Err(last_err.unwrap_or_else(|| anyhow::anyhow!("microphone sink init failed")))
}

/// Pick the UVC gadget video node.
/// Priority: 1) `LESAVKA_UVC_DEV` override; 2) first `video_output` node.
/// Returns an error when nothing matches instead of guessing a capture card.
fn pick_uvc_device() -> anyhow::Result<String> {
    if let Ok(path) = std::env::var("LESAVKA_UVC_DEV") {
        return Ok(path);
    }

    let ctrl = UsbGadget::find_controller().ok();
    if let Some(ctrl) = ctrl.as_deref() {
        let by_path = format!("/dev/v4l/by-path/platform-{ctrl}-video-index0");
        if Path::new(&by_path).exists() {
            return Ok(by_path);
        }
    }

    // walk /dev/video* via udev and look for an output‑capable node (gadget exposes one)
    let mut fallback: Option<String> = None;
    if let Ok(mut en) = udev::Enumerator::new() {
        let _ = en.match_subsystem("video4linux");
        if let Ok(devs) = en.scan_devices() {
            for dev in devs {
                let caps = dev
                    .property_value("ID_V4L_CAPABILITIES")
                    .and_then(|v| v.to_str())
                    .unwrap_or_default();
                if !caps.contains(":video_output:") {
                    continue;
                }
                let Some(node) = dev.devnode() else { continue };
                let node = node.to_string_lossy().into_owned();
                let product = dev
                    .property_value("ID_V4L_PRODUCT")
                    .and_then(|v| v.to_str())
                    .unwrap_or_default();
                let path = dev
                    .property_value("ID_PATH")
                    .and_then(|v| v.to_str())
                    .unwrap_or_default();
                if let Some(ctrl) = ctrl.as_deref() {
                    if product == ctrl || path.contains(ctrl) {
                        return Ok(node);
                    }
                }
                if fallback.is_none() {
                    fallback = Some(node);
                }
            }
        }
    }
    if let Some(node) = fallback {
        return Ok(node);
    }

    Err(anyhow::anyhow!(
        "no video_output v4l2 node found; set LESAVKA_UVC_DEV"
    ))
}

fn uvc_ctrl_bin() -> String {
    std::env::var("LESAVKA_UVC_CTRL_BIN")
        .unwrap_or_else(|_| "/usr/local/bin/lesavka-uvc".to_string())
}

fn spawn_uvc_control(bin: &str, uvc_dev: &str) -> anyhow::Result<tokio::process::Child> {
    Command::new(bin)
        .arg("--device")
        .arg(uvc_dev)
        .spawn()
        .context("spawning lesavka-uvc")
}

async fn supervise_uvc_control(bin: String) {
    let mut waiting_logged = false;
    loop {
        let uvc_dev = match pick_uvc_device() {
            Ok(dev) => {
                if waiting_logged {
                    info!(%dev, "📷 UVC device discovered");
                    waiting_logged = false;
                }
                dev
            }
            Err(e) => {
                if !waiting_logged {
                    warn!("⚠️  UVC device not ready: {e:#}");
                    waiting_logged = true;
                }
                tokio::time::sleep(Duration::from_secs(2)).await;
                continue;
            }
        };

        match spawn_uvc_control(&bin, &uvc_dev) {
            Ok(mut child) => {
                info!(%uvc_dev, "📷 UVC control helper started");
                match child.wait().await {
                    Ok(status) => {
                        warn!(%uvc_dev, "⚠️  lesavka-uvc exited: {status}");
                    }
                    Err(e) => {
                        warn!(%uvc_dev, "⚠️  lesavka-uvc wait failed: {e:#}");
                    }
                }
            }
            Err(e) => {
                warn!(%uvc_dev, "⚠️  failed to start lesavka-uvc: {e:#}");
            }
        }

        tokio::time::sleep(Duration::from_secs(2)).await;
    }
}

/*──────────────── Handler ───────────────────*/
struct Handler {
    kb: Arc<Mutex<tokio::fs::File>>,
    ms: Arc<Mutex<tokio::fs::File>>,
    gadget: UsbGadget,
    did_cycle: Arc<AtomicBool>,
    camera_rt: Arc<CameraRuntime>,
}

struct CameraRelaySlot {
    cfg: camera::CameraConfig,
    relay: Arc<video::CameraRelay>,
}

struct CameraRuntime {
    generation: AtomicU64,
    slot: Mutex<Option<CameraRelaySlot>>,
}

impl CameraRuntime {
    fn new() -> Self {
        Self {
            generation: AtomicU64::new(0),
            slot: Mutex::new(None),
        }
    }

    async fn activate(
        &self,
        cfg: &camera::CameraConfig,
    ) -> Result<(u64, Arc<video::CameraRelay>), Status> {
        let session_id = self.generation.fetch_add(1, Ordering::SeqCst) + 1;
        let mut slot = self.slot.lock().await;
        let mut reused = false;

        let relay = if let Some(existing) = slot.as_ref() {
            if camera_cfg_eq(&existing.cfg, cfg) {
                reused = true;
                existing.relay.clone()
            } else {
                self.make_relay(cfg)?
            }
        } else {
            self.make_relay(cfg)?
        };

        if !reused {
            *slot = Some(CameraRelaySlot {
                cfg: cfg.clone(),
                relay: relay.clone(),
            });
            info!(
                session_id,
                output = cfg.output.as_str(),
                codec = cfg.codec.as_str(),
                width = cfg.width,
                height = cfg.height,
                fps = cfg.fps,
                "🎥 camera relay (re)created"
            );
        } else {
            info!(session_id, "🎥 camera relay reused");
        }

        Ok((session_id, relay))
    }

    fn is_active(&self, session_id: u64) -> bool {
        self.generation.load(Ordering::Relaxed) == session_id
    }

    fn make_relay(&self, cfg: &camera::CameraConfig) -> Result<Arc<video::CameraRelay>, Status> {
        let relay = match cfg.output {
            camera::CameraOutput::Uvc => {
                if std::env::var("LESAVKA_DISABLE_UVC").is_ok() {
                    return Err(Status::failed_precondition(
                        "UVC output disabled (LESAVKA_DISABLE_UVC set)",
                    ));
                }
                let uvc = pick_uvc_device().map_err(|e| Status::internal(format!("{e:#}")))?;
                info!(%uvc, "🎥 stream_camera using UVC sink");
                video::CameraRelay::new_uvc(0, &uvc, cfg)
                    .map_err(|e| Status::internal(format!("{e:#}")))?
            }
            camera::CameraOutput::Hdmi => video::CameraRelay::new_hdmi(0, cfg)
                .map_err(|e| Status::internal(format!("{e:#}")))?,
        };
        Ok(Arc::new(relay))
    }
}

fn camera_cfg_eq(a: &camera::CameraConfig, b: &camera::CameraConfig) -> bool {
    if a.output != b.output
        || a.codec != b.codec
        || a.width != b.width
        || a.height != b.height
        || a.fps != b.fps
    {
        return false;
    }
    match (&a.hdmi, &b.hdmi) {
        (Some(ha), Some(hb)) => ha.name == hb.name && ha.id == hb.id,
        (None, None) => true,
        _ => false,
    }
}

impl Handler {
    async fn new(gadget: UsbGadget) -> anyhow::Result<Self> {
        if allow_gadget_cycle() {
            info!("🛠️  Initial USB reset…");
            let _ = gadget.cycle(); // ignore failure - may boot without host
        } else {
            info!(
                "🔒 gadget cycle disabled at startup (set LESAVKA_ALLOW_GADGET_CYCLE=1 to enable)"
            );
        }

        info!("🛠️  opening HID endpoints …");
        let kb = open_with_retry("/dev/hidg0").await?;
        let ms = open_with_retry("/dev/hidg1").await?;
        info!("✅ HID endpoints ready");

        Ok(Self {
            kb: Arc::new(Mutex::new(kb)),
            ms: Arc::new(Mutex::new(ms)),
            gadget,
            did_cycle: Arc::new(AtomicBool::new(false)),
            camera_rt: Arc::new(CameraRuntime::new()),
        })
    }

    async fn reopen_hid(&self) -> anyhow::Result<()> {
        let kb_new = open_with_retry("/dev/hidg0").await?;
        let ms_new = open_with_retry("/dev/hidg1").await?;
        *self.kb.lock().await = kb_new;
        *self.ms.lock().await = ms_new;
        Ok(())
    }
}

/*──────────────── gRPC service ─────────────*/
#[tonic::async_trait]
impl Relay for Handler {
    /* existing streams ─ unchanged, except: no more auto-reset */
    type StreamKeyboardStream = ReceiverStream<Result<KeyboardReport, Status>>;
    type StreamMouseStream = ReceiverStream<Result<MouseReport, Status>>;
    type CaptureVideoStream = Pin<Box<dyn Stream<Item = Result<VideoPacket, Status>> + Send>>;
    type CaptureAudioStream = Pin<Box<dyn Stream<Item = Result<AudioPacket, Status>> + Send>>;
    type StreamMicrophoneStream = ReceiverStream<Result<Empty, Status>>;
    type StreamCameraStream = ReceiverStream<Result<Empty, Status>>;

    async fn stream_keyboard(
        &self,
        req: Request<tonic::Streaming<KeyboardReport>>,
    ) -> Result<Response<Self::StreamKeyboardStream>, Status> {
        let (tx, rx) = tokio::sync::mpsc::channel(32);
        let kb = self.kb.clone();
        let ms = self.ms.clone();
        let gadget = self.gadget.clone();
        let did_cycle = self.did_cycle.clone();

        tokio::spawn(async move {
            let mut s = req.into_inner();
            while let Some(pkt) = s.next().await.transpose()? {
                if let Err(e) = kb.lock().await.write_all(&pkt.data).await {
                    warn!("⌨️  write failed: {e} (dropped)");
                    recover_hid_if_needed(
                        &e,
                        gadget.clone(),
                        kb.clone(),
                        ms.clone(),
                        did_cycle.clone(),
                    )
                    .await;
                }
                tx.send(Ok(pkt)).await.ok();
            }
            Ok::<(), Status>(())
        });

        Ok(Response::new(ReceiverStream::new(rx)))
    }

    async fn stream_mouse(
        &self,
        req: Request<tonic::Streaming<MouseReport>>,
    ) -> Result<Response<Self::StreamMouseStream>, Status> {
        let (tx, rx) = tokio::sync::mpsc::channel(1024);
        let ms = self.ms.clone();
        let kb = self.kb.clone();
        let gadget = self.gadget.clone();
        let did_cycle = self.did_cycle.clone();

        tokio::spawn(async move {
            let mut s = req.into_inner();
            while let Some(pkt) = s.next().await.transpose()? {
                if let Err(e) = ms.lock().await.write_all(&pkt.data).await {
                    warn!("🖱️  write failed: {e} (dropped)");
                    recover_hid_if_needed(
                        &e,
                        gadget.clone(),
                        kb.clone(),
                        ms.clone(),
                        did_cycle.clone(),
                    )
                    .await;
                }
                tx.send(Ok(pkt)).await.ok();
            }
            Ok::<(), Status>(())
        });

        Ok(Response::new(ReceiverStream::new(rx)))
    }

    async fn stream_microphone(
        &self,
        req: Request<tonic::Streaming<AudioPacket>>,
    ) -> Result<Response<Self::StreamMicrophoneStream>, Status> {
        // 1 ─ build once, early
        let uac_dev = std::env::var("LESAVKA_UAC_DEV").unwrap_or_else(|_| "hw:UAC2Gadget,0".into());
        info!(%uac_dev, "🎤 stream_microphone using UAC sink");
        let mut sink = open_voice_with_retry(&uac_dev)
            .await
            .map_err(|e| Status::internal(format!("{e:#}")))?;

        // 2 ─ dummy outbound stream (same trick as before)
        let (tx, rx) = tokio::sync::mpsc::channel(1);

        // 3 ─ drive the sink in a background task
        tokio::spawn(async move {
            let mut inbound = req.into_inner();
            static CNT: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);

            while let Some(pkt) = inbound.next().await.transpose()? {
                let n = CNT.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
                if n < 10 || n % 300 == 0 {
                    tracing::trace!("🎤⬇ srv  pkt#{n}  {} bytes", pkt.data.len());
                }
                sink.push(&pkt);
            }
            sink.finish(); // flush on EOS
            let _ = tx.send(Ok(Empty {})).await;
            Ok::<(), Status>(())
        });

        Ok(Response::new(ReceiverStream::new(rx)))
    }

    async fn stream_camera(
        &self,
        req: Request<tonic::Streaming<VideoPacket>>,
    ) -> Result<Response<Self::StreamCameraStream>, Status> {
        let cfg = camera::current_camera_config();
        info!(
            output = cfg.output.as_str(),
            codec = cfg.codec.as_str(),
            width = cfg.width,
            height = cfg.height,
            fps = cfg.fps,
            hdmi = cfg.hdmi.as_ref().map(|h| h.name.as_str()).unwrap_or("none"),
            "🎥 stream_camera output selected"
        );

        let (session_id, relay) = self.camera_rt.activate(&cfg).await?;
        let camera_rt = self.camera_rt.clone();

        // dummy outbound (same pattern as other streams)
        let (tx, rx) = tokio::sync::mpsc::channel(1);

        tokio::spawn(async move {
            let mut s = req.into_inner();
            while let Some(pkt) = s.next().await.transpose()? {
                if !camera_rt.is_active(session_id) {
                    info!(session_id, "🎥 stream_camera session superseded");
                    break;
                }
                relay.feed(pkt); // ← all logging inside video.rs
            }
            tx.send(Ok(Empty {})).await.ok();
            Ok::<(), Status>(())
        });

        Ok(Response::new(ReceiverStream::new(rx)))
    }

    async fn capture_video(
        &self,
        req: Request<MonitorRequest>,
    ) -> Result<Response<Self::CaptureVideoStream>, Status> {
        let req = req.into_inner();
        let id = req.id;
        let dev = match id {
            0 => "/dev/lesavka_l_eye",
            1 => "/dev/lesavka_r_eye",
            _ => return Err(Status::invalid_argument("monitor id must be 0 or 1")),
        };
        debug!("🎥 streaming {dev}");
        let s = video::eye_ball(dev, id, req.max_bitrate)
            .await
            .map_err(|e| Status::internal(format!("{e:#}")))?;
        Ok(Response::new(Box::pin(s)))
    }

    async fn capture_audio(
        &self,
        req: Request<MonitorRequest>,
    ) -> Result<Response<Self::CaptureAudioStream>, Status> {
        // Only one speaker stream for now; both 0/1 → same ALSA dev.
        let _id = req.into_inner().id;
        // Allow override (`LESAVKA_ALSA_DEV=hw:2,0` for debugging).
        let dev = std::env::var("LESAVKA_ALSA_DEV").unwrap_or_else(|_| "hw:UAC2Gadget,0".into());

        let s = audio::ear(&dev, 0)
            .await
            .map_err(|e| Status::internal(format!("{e:#}")))?;

        Ok(Response::new(Box::pin(s)))
    }

    async fn paste_text(&self, req: Request<PasteRequest>) -> Result<Response<PasteReply>, Status> {
        let req = req.into_inner();
        let text = paste::decrypt(&req).map_err(|e| Status::unauthenticated(format!("{e}")))?;
        if let Err(e) = paste::type_text(self.kb.as_ref(), &text).await {
            return Ok(Response::new(PasteReply {
                ok: false,
                error: format!("{e}"),
            }));
        }
        Ok(Response::new(PasteReply {
            ok: true,
            error: String::new(),
        }))
    }

    /*────────────── USB-reset RPC ────────────*/
    async fn reset_usb(&self, _req: Request<Empty>) -> Result<Response<ResetUsbReply>, Status> {
        info!("🔴 explicit ResetUsb() called");
        match self.gadget.cycle() {
            Ok(_) => {
                if let Err(e) = self.reopen_hid().await {
                    error!("💥 reopen HID failed: {e:#}");
                    return Err(Status::internal(e.to_string()));
                }
                Ok(Response::new(ResetUsbReply { ok: true }))
            }
            Err(e) => {
                error!("💥 cycle failed: {e:#}");
                Err(Status::internal(e.to_string()))
            }
        }
    }
}

/*──────────────── main ───────────────────────*/
#[tokio::main(worker_threads = 4)]
async fn main() -> anyhow::Result<()> {
    let _guard = init_tracing()?;
    info!("🚀 {} v{} starting up", PKG_NAME, VERSION);

    panic::set_hook(Box::new(|p| {
        let bt = Backtrace::force_capture();
        error!("💥 panic: {p}\n{bt}");
    }));

    let gadget = UsbGadget::new("lesavka");
    if std::env::var("LESAVKA_DISABLE_UVC").is_err() {
        if std::env::var("LESAVKA_UVC_EXTERNAL").is_ok() {
            info!("📷 UVC control helper external; not spawning");
        } else {
            let bin = uvc_ctrl_bin();
            tokio::spawn(supervise_uvc_control(bin));
        }
    } else {
        info!("📷 UVC disabled (LESAVKA_DISABLE_UVC set)");
    }
    let handler = Handler::new(gadget.clone()).await?;

    info!("🌐 lesavka-server listening on 0.0.0.0:50051");
    Server::builder()
        .tcp_nodelay(true)
        .max_frame_size(Some(2 * 1024 * 1024))
        .add_service(RelayServer::new(handler))
        .add_service(HandshakeSvc::server())
        .add_service(ReflBuilder::configure().build_v1().unwrap())
        .serve(([0, 0, 0, 0], 50051).into())
        .await?;
    Ok(())
}