use std::collections::HashMap;
use std::path::{Path, PathBuf};
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use crate::error::TelegramResult;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TelegramState {
pub chat_id: Option<i64>,
pub last_seen: Option<DateTime<Utc>>,
#[serde(default)]
pub last_update_id: Option<i32>,
#[serde(default)]
pub pending_questions: HashMap<String, PendingQuestion>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PendingQuestion {
pub asked_at: DateTime<Utc>,
pub message_id: i32,
}
pub struct StateManager {
path: PathBuf,
}
impl StateManager {
pub fn new(path: impl Into<PathBuf>) -> Self {
Self { path: path.into() }
}
pub fn load(&self) -> TelegramResult<Option<TelegramState>> {
if !self.path.exists() {
return Ok(None);
}
let contents = std::fs::read_to_string(&self.path)?;
let state: TelegramState = serde_json::from_str(&contents)?;
Ok(Some(state))
}
pub fn save(&self, state: &TelegramState) -> TelegramResult<()> {
let json = serde_json::to_string_pretty(state)?;
let tmp_path = self.path.with_extension("json.tmp");
if let Some(parent) = self.path.parent() {
std::fs::create_dir_all(parent)?;
}
std::fs::write(&tmp_path, &json)?;
std::fs::rename(&tmp_path, &self.path)?;
Ok(())
}
pub fn load_or_default(&self) -> TelegramResult<TelegramState> {
Ok(self.load()?.unwrap_or_else(|| TelegramState {
chat_id: None,
last_seen: None,
last_update_id: None,
pending_questions: HashMap::new(),
}))
}
pub fn add_pending_question(
&self,
state: &mut TelegramState,
loop_id: &str,
message_id: i32,
) -> TelegramResult<()> {
state.pending_questions.insert(
loop_id.to_string(),
PendingQuestion {
asked_at: Utc::now(),
message_id,
},
);
self.save(state)
}
pub fn remove_pending_question(
&self,
state: &mut TelegramState,
loop_id: &str,
) -> TelegramResult<()> {
state.pending_questions.remove(loop_id);
self.save(state)
}
pub fn get_loop_for_reply(
&self,
state: &TelegramState,
reply_message_id: i32,
) -> Option<String> {
state
.pending_questions
.iter()
.find(|(_, q)| q.message_id == reply_message_id)
.map(|(loop_id, _)| loop_id.clone())
}
pub fn path(&self) -> &Path {
&self.path
}
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::TempDir;
fn test_manager() -> (StateManager, TempDir) {
let dir = TempDir::new().unwrap();
let path = dir.path().join("telegram-state.json");
(StateManager::new(path), dir)
}
#[test]
fn load_missing_file_returns_none() {
let (mgr, _dir) = test_manager();
assert!(mgr.load().unwrap().is_none());
}
#[test]
fn save_and_load_round_trip() {
let (mgr, _dir) = test_manager();
let state = TelegramState {
chat_id: Some(123_456),
last_seen: Some(Utc::now()),
last_update_id: Some(101),
pending_questions: HashMap::new(),
};
mgr.save(&state).unwrap();
let loaded = mgr.load().unwrap().unwrap();
assert_eq!(loaded.chat_id, Some(123_456));
assert_eq!(loaded.last_update_id, Some(101));
}
#[test]
fn corrupted_json_returns_error() {
let (mgr, _dir) = test_manager();
std::fs::write(mgr.path(), "not json").unwrap();
assert!(mgr.load().is_err());
}
#[test]
fn pending_question_tracking() {
let (mgr, _dir) = test_manager();
let mut state = mgr.load_or_default().unwrap();
mgr.add_pending_question(&mut state, "main", 42).unwrap();
assert!(state.pending_questions.contains_key("main"));
assert_eq!(state.pending_questions["main"].message_id, 42);
mgr.remove_pending_question(&mut state, "main").unwrap();
assert!(!state.pending_questions.contains_key("main"));
}
#[test]
fn reply_routing_lookup() {
let (mgr, _dir) = test_manager();
let mut state = mgr.load_or_default().unwrap();
mgr.add_pending_question(&mut state, "main", 10).unwrap();
mgr.add_pending_question(&mut state, "feature-auth", 20)
.unwrap();
assert_eq!(mgr.get_loop_for_reply(&state, 10), Some("main".to_string()));
assert_eq!(
mgr.get_loop_for_reply(&state, 20),
Some("feature-auth".to_string())
);
assert_eq!(mgr.get_loop_for_reply(&state, 99), None);
}
}