1use std::path::Path;
2use std::time::SystemTime;
3
4use anyhow::Result;
5use rusqlite::OptionalExtension;
6use serde::{Deserialize, Serialize};
7use uuid::Uuid;
8
9use crate::store::open_connection;
10
11#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
12#[serde(rename_all = "snake_case")]
13pub enum GoalStatus {
14 Active,
15 Paused,
16 Complete,
17 Blocked,
18 UsageLimited,
20 BudgetLimited,
23}
24
25impl GoalStatus {
26 pub fn label(self) -> &'static str {
27 match self {
28 Self::Active => "active",
29 Self::Paused => "paused",
30 Self::Complete => "complete",
31 Self::Blocked => "blocked",
32 Self::UsageLimited => "usage_limited",
33 Self::BudgetLimited => "budget_limited",
34 }
35 }
36
37 pub fn is_continuable(self) -> bool {
38 matches!(self, Self::Active)
39 }
40
41 pub fn is_terminal(self) -> bool {
49 matches!(self, Self::Complete | Self::BudgetLimited)
50 }
51
52 pub fn from_str(s: &str) -> Option<Self> {
53 match s {
54 "active" => Some(Self::Active),
55 "paused" => Some(Self::Paused),
56 "complete" => Some(Self::Complete),
57 "blocked" => Some(Self::Blocked),
58 "usage_limited" => Some(Self::UsageLimited),
59 "budget_limited" => Some(Self::BudgetLimited),
60 _ => None,
61 }
62 }
63}
64
65impl std::fmt::Display for GoalStatus {
66 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
67 f.write_str(self.label())
68 }
69}
70
71#[derive(Debug, Clone, Serialize, Deserialize)]
72pub struct GoalState {
73 pub goal_id: String,
78 pub objective: String,
79 pub status: GoalStatus,
80 pub tokens_used: i64,
81 pub time_used_seconds: i64,
82 pub token_budget: Option<i64>,
83 pub created_at: String,
84 pub updated_at: String,
85}
86
87pub fn new_goal_id() -> String {
89 Uuid::new_v4().to_string()
90}
91
92pub fn save_goal(store_path: &Path, session_id: &str, goal: &GoalState) -> Result<()> {
93 let conn = open_connection(store_path)?;
94 conn.execute(
95 "INSERT INTO goals (session_id, goal_id, objective, status, tokens_used, time_used_seconds, token_budget, created_at, updated_at)
96 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9)
97 ON CONFLICT(session_id) DO UPDATE SET
98 goal_id = excluded.goal_id,
99 objective = excluded.objective,
100 status = excluded.status,
101 tokens_used = excluded.tokens_used,
102 time_used_seconds = excluded.time_used_seconds,
103 token_budget = excluded.token_budget,
104 updated_at = excluded.updated_at",
105 rusqlite::params![
106 session_id,
107 goal.goal_id,
108 goal.objective,
109 goal.status.label(),
110 goal.tokens_used,
111 goal.time_used_seconds,
112 goal.token_budget,
113 goal.created_at,
114 goal.updated_at,
115 ],
116 )?;
117 Ok(())
118}
119
120pub fn load_goal(store_path: &Path, session_id: &str) -> Result<Option<GoalState>> {
121 let conn = open_connection(store_path)?;
122 let result = conn
123 .query_row(
124 "SELECT goal_id, objective, status, tokens_used, time_used_seconds, token_budget, created_at, updated_at
125 FROM goals WHERE session_id = ?1",
126 [session_id],
127 |row| {
128 Ok(GoalState {
129 goal_id: row.get::<_, String>(0).unwrap_or_default(),
130 objective: row.get(1)?,
131 status: GoalStatus::from_str(&row.get::<_, String>(2)?)
132 .unwrap_or(GoalStatus::Active),
133 tokens_used: row.get(3)?,
134 time_used_seconds: row.get(4)?,
135 token_budget: row.get(5)?,
136 created_at: row.get(6)?,
137 updated_at: row.get(7)?,
138 })
139 },
140 )
141 .optional()?;
142 Ok(result)
143}
144
145#[derive(Debug, Clone, Copy, PartialEq, Eq)]
147pub enum AccountingOutcome {
148 WithinBudget,
150 BudgetExceeded,
152 Skipped,
156}
157
158pub fn account_goal_usage(
168 store_path: &Path,
169 session_id: &str,
170 token_delta: i64,
171 time_delta_seconds: i64,
172 expected_goal_id: Option<&str>,
173) -> Result<AccountingOutcome> {
174 let conn = open_connection(store_path)?;
175 let now = now_utc();
176
177 let rows_affected = if let Some(expected_id) = expected_goal_id {
178 conn.execute(
179 "UPDATE goals
180 SET tokens_used = tokens_used + ?1,
181 time_used_seconds = time_used_seconds + ?2,
182 updated_at = ?3
183 WHERE session_id = ?4 AND goal_id = ?5",
184 rusqlite::params![token_delta, time_delta_seconds, now, session_id, expected_id],
185 )?
186 } else {
187 conn.execute(
188 "UPDATE goals
189 SET tokens_used = tokens_used + ?1,
190 time_used_seconds = time_used_seconds + ?2,
191 updated_at = ?3
192 WHERE session_id = ?4",
193 rusqlite::params![token_delta, time_delta_seconds, now, session_id],
194 )?
195 };
196
197 if rows_affected == 0 {
198 return Ok(AccountingOutcome::Skipped);
201 }
202
203 let (tokens_used, token_budget): (i64, Option<i64>) = conn.query_row(
205 "SELECT tokens_used, token_budget FROM goals WHERE session_id = ?1",
206 [session_id],
207 |row| Ok((row.get(0)?, row.get(1)?)),
208 )?;
209 match token_budget {
210 Some(budget) if tokens_used >= budget => Ok(AccountingOutcome::BudgetExceeded),
211 _ => Ok(AccountingOutcome::WithinBudget),
212 }
213}
214
215pub fn delete_goal(store_path: &Path, session_id: &str) -> Result<bool> {
216 let conn = open_connection(store_path)?;
217 let rows_deleted = conn.execute("DELETE FROM goals WHERE session_id = ?1", [session_id])?;
218 Ok(rows_deleted > 0)
219}
220
221pub fn now_utc() -> String {
222 let now = SystemTime::now()
223 .duration_since(SystemTime::UNIX_EPOCH)
224 .unwrap_or_default();
225 let secs = now.as_secs();
226 let time_of_day = secs % 86400;
227 let hours = time_of_day / 3600;
228 let minutes = (time_of_day % 3600) / 60;
229 let seconds = time_of_day % 60;
230
231 let days = secs / 86400;
233 let (year, month, day) = days_to_ymd(days);
234 format!(
235 "{:04}-{:02}-{:02} {:02}:{:02}:{:02}",
236 year, month, day, hours, minutes, seconds
237 )
238}
239
240fn days_to_ymd(days: u64) -> (u64, u64, u64) {
241 let z = days + 719468;
243 let era = z / 146097;
244 let doe = z - era * 146097;
245 let yoe = (doe - doe / 1460 + doe / 36524 - doe / 146096) / 365;
246 let y = yoe + era * 400;
247 let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
248 let mp = (5 * doy + 2) / 153;
249 let d = doy - (153 * mp + 2) / 5 + 1;
250 let m = if mp < 10 { mp + 3 } else { mp - 9 };
251 let y = if m <= 2 { y + 1 } else { y };
252 (y, m, d)
253}