luma-home/src/ha_client.rs

//! Home Assistant WebSocket client running on a dedicated background thread.
//! Follows the same mpsc/thread pattern as weather-service.rs — no async runtime.
use serde_json::{Value, json};
use std::{
    collections::HashMap,
    sync::{
        Arc,
        atomic::{AtomicBool, Ordering},
        mpsc,
    },
    thread,
    time::{Duration, Instant},
};

const RECONNECT_DELAY: Duration = Duration::from_secs(5);
const PING_INTERVAL: Duration = Duration::from_secs(30);
const READ_TIMEOUT: Duration = Duration::from_secs(45);

// ── Public types ──────────────────────────────────────────────────────

#[derive(Clone, Debug)]
pub struct EntityState {
    pub entity_id: String,
    pub state: String,
    pub attributes: Value,
    pub last_changed: String,
}

#[derive(Clone, Debug)]
pub struct HistoryPoint {
    pub timestamp: f64,
    pub value: f64,
}

pub enum HaResponse {
    Connected,
    Disconnected(String),
    StateChanged(EntityState),
    AllStates(Vec<EntityState>),
    History {
        entity_id: String,
        points: Vec<HistoryPoint>,
    },
    ServiceResult {
        id: u64,
        success: bool,
    },
}

pub enum HaCommand {
    CallService {
        id: u64,
        domain: String,
        service: String,
        target: Value,
        data: Value,
    },
    FetchHistory {
        entity_id: String,
        hours: u32,
    },
}

// ── Config ────────────────────────────────────────────────────────────

#[derive(Clone, Debug)]
pub struct HaConfig {
    pub url: String,
    pub token: String,
    pub entity_ids: Vec<String>,
}

// ── Service (main-thread handle) ──────────────────────────────────────

pub struct HaService {
    pub receiver: mpsc::Receiver<HaResponse>,
    commands: mpsc::Sender<HaCommand>,
    closed: Arc<AtomicBool>,
    worker: Option<thread::JoinHandle<()>>,
}

impl HaService {
    pub fn new(config: HaConfig) -> Self {
        let (resp_tx, resp_rx) = mpsc::channel();
        let (cmd_tx, cmd_rx) = mpsc::channel();
        let closed = Arc::new(AtomicBool::new(false));
        let closed_flag = closed.clone();

        let worker = thread::spawn(move || {
            worker_loop(config, resp_tx, cmd_rx, closed_flag);
        });

        Self {
            receiver: resp_rx,
            commands: cmd_tx,
            closed,
            worker: Some(worker),
        }
    }

    pub fn send(&self, command: HaCommand) {
        let _ = self.commands.send(command);
    }
}

impl Drop for HaService {
    fn drop(&mut self) {
        self.closed.store(true, Ordering::SeqCst);
        if let Some(worker) = self.worker.take() {
            let _ = worker.join();
        }
    }
}

// ── Worker thread ─────────────────────────────────────────────────────

fn worker_loop(
    config: HaConfig,
    responses: mpsc::Sender<HaResponse>,
    commands: mpsc::Receiver<HaCommand>,
    closed: Arc<AtomicBool>,
) {
    while !closed.load(Ordering::SeqCst) {
        match connect_and_run(&config, &responses, &commands, &closed) {
            Ok(()) => break,
            Err(e) => {
                eprintln!("HA WebSocket: {e}");
                let _ = responses.send(HaResponse::Disconnected(e));
                if closed.load(Ordering::SeqCst) {
                    break;
                }
                thread::sleep(RECONNECT_DELAY);
            }
        }
    }
}

fn connect_and_run(
    config: &HaConfig,
    responses: &mpsc::Sender<HaResponse>,
    commands: &mpsc::Receiver<HaCommand>,
    closed: &Arc<AtomicBool>,
) -> Result<(), String> {
    use tungstenite::{Message, connect};

    let (mut ws, _) =
        connect(&config.url).map_err(|e| format!("connect failed: {e}"))?;

    match ws.get_ref() {
        tungstenite::stream::MaybeTlsStream::Plain(tcp) => {
            tcp.set_read_timeout(Some(READ_TIMEOUT)).ok();
        }
        tungstenite::stream::MaybeTlsStream::Rustls(tls) => {
            tls.get_ref().set_read_timeout(Some(READ_TIMEOUT)).ok();
        }
        _ => {}
    }

    let mut msg_id: u64 = 0;
    let mut next_id = || {
        msg_id += 1;
        msg_id
    };

    // Phase 1: auth_required
    let auth_required = read_json(&mut ws)?;
    if auth_required.get("type").and_then(Value::as_str) != Some("auth_required") {
        return Err(format!("expected auth_required, got: {auth_required}"));
    }

    // Phase 2: authenticate
    ws_send(
        &mut ws,
        &json!({ "type": "auth", "access_token": config.token }),
    )?;

    let auth_result = read_json(&mut ws)?;
    match auth_result.get("type").and_then(Value::as_str) {
        Some("auth_ok") => {}
        Some("auth_invalid") => {
            return Err(format!(
                "auth failed: {}",
                auth_result
                    .get("message")
                    .and_then(Value::as_str)
                    .unwrap_or("invalid token")
            ));
        }
        _ => return Err(format!("unexpected auth response: {auth_result}")),
    }

    let _ = responses.send(HaResponse::Connected);

    // Phase 3: get all states
    let states_id = next_id();
    ws_send(&mut ws, &json!({ "id": states_id, "type": "get_states" }))?;

    // Phase 4: subscribe to state changes
    let sub_id = next_id();
    ws_send(
        &mut ws,
        &json!({
            "id": sub_id,
            "type": "subscribe_events",
            "event_type": "state_changed"
        }),
    )?;

    // Pending service call / history requests mapped by HA message id
    let mut pending_service: HashMap<u64, u64> = HashMap::new(); // ha_id -> caller_id
    let mut pending_history: HashMap<u64, String> = HashMap::new(); // ha_id -> entity_id

    let mut last_ping = Instant::now();

    loop {
        if closed.load(Ordering::SeqCst) {
            let _ = ws.close(None);
            return Ok(());
        }

        // Drain commands from the main thread
        while let Ok(cmd) = commands.try_recv() {
            match cmd {
                HaCommand::CallService {
                    id: caller_id,
                    domain,
                    service,
                    target,
                    data,
                } => {
                    let ha_id = next_id();
                    pending_service.insert(ha_id, caller_id);
                    ws_send(
                        &mut ws,
                        &json!({
                            "id": ha_id,
                            "type": "call_service",
                            "domain": domain,
                            "service": service,
                            "target": target,
                            "service_data": data,
                        }),
                    )?;
                }
                HaCommand::FetchHistory { entity_id, hours } => {
                    let ha_id = next_id();
                    pending_history.insert(ha_id, entity_id.clone());
                    let now = time::OffsetDateTime::now_utc();
                    let start = now - time::Duration::hours(hours as i64);
                    let format =
                        time::format_description::well_known::Rfc3339;
                    let start_str = start.format(&format).unwrap_or_default();
                    let end_str = now.format(&format).unwrap_or_default();
                    ws_send(
                        &mut ws,
                        &json!({
                            "id": ha_id,
                            "type": "recorder/statistics_during_period",
                            "start_time": start_str,
                            "end_time": end_str,
                            "statistic_ids": [entity_id],
                            "period": "5minute",
                            "types": ["mean"],
                        }),
                    )?;
                }
            }
        }

        // Send ping if needed
        if last_ping.elapsed() >= PING_INTERVAL {
            let ping_id = next_id();
            ws_send(&mut ws, &json!({ "id": ping_id, "type": "ping" }))?;
            last_ping = Instant::now();
        }

        // Read one message (with timeout set on the socket)
        let msg = match ws.read() {
            Ok(msg) => msg,
            Err(tungstenite::Error::Io(ref e))
                if e.kind() == std::io::ErrorKind::WouldBlock
                    || e.kind() == std::io::ErrorKind::TimedOut =>
            {
                continue;
            }
            Err(e) => return Err(format!("read error: {e}")),
        };

        let text = match msg {
            Message::Text(t) => t,
            Message::Ping(data) => {
                let _ = ws.send(Message::Pong(data));
                continue;
            }
            Message::Close(_) => return Err("server closed connection".into()),
            _ => continue,
        };

        let value: Value = match serde_json::from_str(&text) {
            Ok(v) => v,
            Err(_) => continue,
        };

        let msg_type = value.get("type").and_then(Value::as_str).unwrap_or("");
        let id = value.get("id").and_then(Value::as_u64).unwrap_or(0);

        match msg_type {
            "result" if id == states_id => {
                if let Some(result) = value.get("result").and_then(Value::as_array) {
                    let states: Vec<EntityState> = result
                        .iter()
                        .filter_map(parse_entity_state)
                        .filter(|e| config.entity_ids.contains(&e.entity_id))
                        .collect();
                    let _ = responses.send(HaResponse::AllStates(states));
                }
            }
            "event" => {
                if let Some(event_data) = value
                    .get("event")
                    .and_then(|e| e.get("data"))
                    .and_then(|d| d.get("new_state"))
                {
                    if let Some(entity) = parse_entity_state(event_data) {
                        if config.entity_ids.contains(&entity.entity_id) {
                            let _ = responses.send(HaResponse::StateChanged(entity));
                        }
                    }
                }
            }
            "result" if pending_service.contains_key(&id) => {
                if let Some(caller_id) = pending_service.remove(&id) {
                    let success = value
                        .get("success")
                        .and_then(Value::as_bool)
                        .unwrap_or(false);
                    let _ = responses.send(HaResponse::ServiceResult {
                        id: caller_id,
                        success,
                    });
                }
            }
            "result" if pending_history.contains_key(&id) => {
                if let Some(entity_id) = pending_history.remove(&id) {
                    let points = parse_statistics(&value, &entity_id);
                    let _ = responses.send(HaResponse::History {
                        entity_id,
                        points,
                    });
                }
            }
            "pong" => {}
            _ => {}
        }
    }
}

// ── Helpers ───────────────────────────────────────────────────────────

type WsStream = tungstenite::WebSocket<tungstenite::stream::MaybeTlsStream<std::net::TcpStream>>;

fn ws_send(ws: &mut WsStream, value: &Value) -> Result<(), String> {
    let text = serde_json::to_string(value).map_err(|e| e.to_string())?;
    ws.send(tungstenite::Message::Text(text))
        .map_err(|e| format!("send failed: {e}"))
}

fn read_json(ws: &mut WsStream) -> Result<Value, String> {
    loop {
        match ws.read() {
            Ok(tungstenite::Message::Text(text)) => {
                return serde_json::from_str(&text)
                    .map_err(|e| format!("invalid JSON: {e}"));
            }
            Ok(tungstenite::Message::Ping(data)) => {
                let _ = ws.send(tungstenite::Message::Pong(data));
            }
            Ok(tungstenite::Message::Close(_)) => {
                return Err("connection closed during handshake".into());
            }
            Ok(_) => continue,
            Err(e) => return Err(format!("read failed: {e}")),
        }
    }
}

fn parse_entity_state(value: &Value) -> Option<EntityState> {
    Some(EntityState {
        entity_id: value.get("entity_id")?.as_str()?.to_owned(),
        state: value.get("state")?.as_str()?.to_owned(),
        attributes: value.get("attributes").cloned().unwrap_or(Value::Object(Default::default())),
        last_changed: value
            .get("last_changed")
            .and_then(Value::as_str)
            .unwrap_or("")
            .to_owned(),
    })
}

fn parse_statistics(value: &Value, entity_id: &str) -> Vec<HistoryPoint> {
    let result = match value.get("result") {
        Some(r) => r,
        None => return vec![],
    };
    let stats = match result.get(entity_id).and_then(Value::as_array) {
        Some(a) => a,
        None => return vec![],
    };
    stats
        .iter()
        .filter_map(|entry| {
            let ts = entry.get("start")?.as_f64()?;
            let val = entry.get("mean")?.as_f64()?;
            Some(HistoryPoint {
                timestamp: ts,
                value: val,
            })
        })
        .collect()
}

/// Collect all entity IDs referenced in the HA config JSON.
pub fn entity_ids_from_config(config: &Value) -> Vec<String> {
    let mut ids = Vec::new();

    if let Some(weather_id) = config
        .get("weather")
        .and_then(|w| w.get("entity_id"))
        .and_then(Value::as_str)
    {
        ids.push(weather_id.to_owned());
    }

    if let Some(vacuum) = config.get("vacuum") {
        if let Some(id) = vacuum.get("entity_id").and_then(Value::as_str) {
            ids.push(id.to_owned());
        }
        if let Some(id) = vacuum.get("cleaning_mode_entity").and_then(Value::as_str) {
            ids.push(id.to_owned());
        }
    }

    if let Some(media) = config.get("media").and_then(|m| m.get("entity_id")).and_then(Value::as_str) {
        ids.push(media.to_owned());
    }

    if let Some(rooms) = config.get("rooms").and_then(Value::as_array) {
        for room in rooms {
            if let Some(sensors) = room.get("sensors").and_then(Value::as_array) {
                for sensor in sensors {
                    if let Some(id) = sensor.as_str() {
                        ids.push(id.to_owned());
                    }
                }
            }
            if let Some(cards) = room.get("cards").and_then(Value::as_array) {
                for card in cards {
                    if let Some(id) = card.get("entity_id").and_then(Value::as_str) {
                        ids.push(id.to_owned());
                    }
                }
            }
        }
    }

    ids.sort();
    ids.dedup();
    ids
}