1use axum::extract::State;
2use axum::{Extension, Json, http::StatusCode, middleware};
3use chrono::{DateTime, NaiveDate, Utc};
4use regex::Regex;
5use serde::{Deserialize, Serialize};
6use tracing::instrument;
7
8use crate::auth::{AuthState, AuthUser};
9use crate::middleware::require_auth;
10use crate::repository::{InMemoryTaskRepository, TaskRepository};
11use oxinbox_core::Task;
12
13#[derive(Debug, Deserialize)]
14pub struct QueryRequest {
15 pub query: String,
16}
17
18#[derive(Debug, Serialize)]
19pub struct QueryResponse {
20 pub sql: String,
21 pub results: Vec<Task>,
22 pub answer: String,
23}
24
25const SQL_GEN_PROMPT: &str = concat!(
26 "Eres un asistente que convierte preguntas en lenguaje natural a SQL para PostgreSQL.\n\n",
27 "Esquema de la base de datos:\n",
28 "CREATE TABLE tasks (\n",
29 " id UUID PRIMARY KEY,\n",
30 " completed BOOLEAN,\n",
31 " priority CHAR(1),\n",
32 " description TEXT,\n",
33 " projects TEXT[],\n",
34 " contexts TEXT[],\n",
35 " status TEXT CHECK (status IN ('inbox', 'todo', 'doing', 'done', 'someday')),\n",
36 " created_at TIMESTAMPTZ,\n",
37 " updated_at TIMESTAMPTZ,\n",
38 " completed_at TIMESTAMPTZ,\n",
39 " due_date DATE\n",
40 ");\n\n",
41 "Reglas:\n",
42 "- Usa `ILIKE` para búsquedas de texto (ej: description ILIKE '%term%')\n",
43 "- projects y contexts son TEXT[], usa `array['val']` o la sintaxis `'{val}'`\n",
44 "- Para buscar en arrays usa: `'valor' = ANY(projects)`\n",
45 "- status puede ser: 'inbox', 'todo', 'doing', 'done', 'someday'\n",
46 "- dates están en formato ISO 8601 (YYYY-MM-DD)\n",
47 "- priority es un CHAR(1): 'A', 'B', 'C', o NULL\n\n",
48 "Responde ÚNICAMENTE con la sentencia SQL, sin explicaciones ni formato markdown.",
49);
50
51const NL_PROMPT_PREFIX: &str = "Eres un asistente que responde preguntas sobre tareas en español.\n\nLa pregunta del usuario era: \"";
52
53fn extract_sql(text: &str) -> String {
54 let text = text.trim();
55 if let Some(stripped) = text
56 .strip_prefix("```sql")
57 .or_else(|| text.strip_prefix("```"))
58 {
59 if let Some(end) = stripped.find("```") {
60 return stripped[..end].trim().to_string();
61 }
62 return stripped.trim().to_string();
63 }
64 text.to_string()
65}
66
67#[allow(clippy::too_many_lines)]
68fn execute_sql(sql: &str, tasks: &[Task]) -> Vec<Task> {
69 let sql_lower = sql.trim().to_lowercase();
70 if !sql_lower.starts_with("select") {
71 return Vec::new();
72 }
73
74 let where_clause = sql_lower
75 .split("order by")
76 .next()
77 .unwrap_or("")
78 .split("limit")
79 .next()
80 .unwrap_or("")
81 .split("select * from tasks")
82 .last()
83 .unwrap_or("")
84 .trim();
85 let where_str = where_clause
86 .strip_prefix("where")
87 .unwrap_or(where_clause)
88 .trim();
89 let where_str = where_str
90 .strip_prefix('(')
91 .unwrap_or(where_str)
92 .strip_suffix(')')
93 .unwrap_or(where_str);
94
95 let order_clause = sql_lower
96 .split("order by")
97 .nth(1)
98 .and_then(|s| s.split("limit").next())
99 .map_or("", str::trim);
100 let limit_clause = sql_lower
101 .split("limit")
102 .nth(1)
103 .and_then(|s| s.split_whitespace().next())
104 .and_then(|s| s.parse::<usize>().ok());
105
106 let mut filtered: Vec<&Task> = tasks.iter().collect();
107
108 if !where_str.is_empty() {
109 let conditions = split_conditions(where_str);
110 filtered.retain(|task| conditions.iter().all(|cond| evaluate_condition(task, cond)));
111 }
112
113 if !order_clause.is_empty() {
114 let parts: Vec<&str> = order_clause.split_whitespace().collect();
115 let col = parts.first().copied().unwrap_or("");
116 let desc = parts.get(1).copied().unwrap_or("asc") == "desc";
117 sort_tasks(&mut filtered, col, desc);
118 }
119
120 if let Some(limit) = limit_clause {
121 filtered.truncate(limit);
122 }
123
124 filtered.into_iter().cloned().collect()
125}
126
127fn split_conditions(where_str: &str) -> Vec<String> {
128 let mut conditions = Vec::new();
129 let mut depth = 0;
130 let mut current = String::new();
131
132 for ch in where_str.chars() {
133 match ch {
134 '(' => {
135 depth += 1;
136 current.push(ch);
137 }
138 ')' => {
139 depth -= 1;
140 current.push(ch);
141 }
142 'a' | 'n' | 'd' if depth == 0 && current.trim().ends_with(" an") && ch == 'd' => {
143 let trimmed = current.trim_end();
144 if let Some(stripped) = trimmed.strip_suffix(" an") {
145 let cond = stripped.trim().to_string();
146 if !cond.is_empty() {
147 conditions.push(cond);
148 }
149 current.clear();
150 } else {
151 current.push(ch);
152 }
153 }
154 'o' | 'r' if depth == 0 && current.trim().ends_with(" o") && ch == 'r' => {
155 let trimmed = current.trim_end();
156 if let Some(stripped) = trimmed.strip_suffix(" o") {
157 let cond = stripped.trim().to_string();
158 if !cond.is_empty() {
159 conditions.push(cond);
160 }
161 current.clear();
162 } else {
163 current.push(ch);
164 }
165 }
166 _ => {
167 current.push(ch);
168 }
169 }
170 }
171
172 let remaining = current.trim().to_string();
173 if !remaining.is_empty() && remaining != "and" && remaining != "or" {
174 conditions.push(remaining);
175 }
176
177 conditions
178}
179
180fn evaluate_condition(task: &Task, condition: &str) -> bool {
181 let condition = condition.trim();
182
183 if let Some(not_cond) = condition.strip_prefix("not ") {
184 return !evaluate_condition(task, not_cond.trim());
185 }
186
187 let re = Regex::new(r"^(\w+)\s*(=|!=|>=|<=|>|<|like|ilike|in)\s*(.+)$").unwrap();
188 let Some(caps) = re.captures(condition) else {
189 return true;
190 };
191
192 let column = caps.get(1).unwrap().as_str();
193 let op = caps.get(2).unwrap().as_str();
194 let raw_value = caps.get(3).unwrap().as_str().trim().trim_matches('\'');
195
196 let task_val = get_column_value(task, column);
197 let cond_val = raw_value.to_lowercase();
198
199 match (op.to_lowercase().as_str(), &task_val) {
200 ("=", ColumnValue::String(s)) => s.to_lowercase() == cond_val,
201 ("=", ColumnValue::Bool(b)) => {
202 let wanted = cond_val == "true" || cond_val == "1";
203 *b == wanted
204 }
205 ("=", ColumnValue::Char(c)) => c.is_some_and(|c| c.to_lowercase().to_string() == cond_val),
206 ("=", ColumnValue::Date(d)) => d.is_some_and(|d| d.to_string() == cond_val),
207 ("=", ColumnValue::DateTime(d)) => {
208 d.is_some_and(|d| d.format("%Y-%m-%d").to_string() == cond_val)
209 }
210 ("=", ColumnValue::StringArray(arr)) => arr.contains(&cond_val),
211
212 ("!=", ColumnValue::String(s)) => s.to_lowercase() != cond_val,
213 ("!=", ColumnValue::Bool(b)) => {
214 let wanted = cond_val == "true" || cond_val == "1";
215 *b != wanted
216 }
217 ("!=", ColumnValue::Char(c)) => c.is_none_or(|c| c.to_lowercase().to_string() != cond_val),
218 ("!=", ColumnValue::Date(d)) => d.is_none_or(|d| d.to_string() != cond_val),
219
220 (">", ColumnValue::Date(d)) => {
221 let parsed = NaiveDate::parse_from_str(&cond_val, "%Y-%m-%d").ok();
222 d.is_some() && parsed.is_some() && d.unwrap() > parsed.unwrap()
223 }
224 (">", ColumnValue::DateTime(d)) => {
225 let parsed = NaiveDate::parse_from_str(&cond_val, "%Y-%m-%d")
226 .ok()
227 .map(|d| d.and_hms_opt(0, 0, 0).unwrap().and_utc());
228 d.is_some() && parsed.is_some() && d.unwrap() > parsed.unwrap()
229 }
230 ("<", ColumnValue::Date(d)) => {
231 let parsed = NaiveDate::parse_from_str(&cond_val, "%Y-%m-%d").ok();
232 d.is_some() && parsed.is_some() && d.unwrap() < parsed.unwrap()
233 }
234 ("<", ColumnValue::DateTime(d)) => {
235 let parsed = NaiveDate::parse_from_str(&cond_val, "%Y-%m-%d")
236 .ok()
237 .map(|d| d.and_hms_opt(0, 0, 0).unwrap().and_utc());
238 d.is_some() && parsed.is_some() && d.unwrap() < parsed.unwrap()
239 }
240 (">=", ColumnValue::Date(d)) => {
241 let parsed = NaiveDate::parse_from_str(&cond_val, "%Y-%m-%d").ok();
242 d.is_some() && parsed.is_some() && d.unwrap() >= parsed.unwrap()
243 }
244 ("<=", ColumnValue::Date(d)) => {
245 let parsed = NaiveDate::parse_from_str(&cond_val, "%Y-%m-%d").ok();
246 d.is_some() && parsed.is_some() && d.unwrap() <= parsed.unwrap()
247 }
248
249 ("like" | "ilike", ColumnValue::String(s)) => {
250 let pattern = cond_val.replace('%', ".*").replace('_', ".");
251 let re = Regex::new(&format!("(?i)^{pattern}$")).unwrap();
252 re.is_match(s)
253 }
254
255 _ => true,
256 }
257}
258
259enum ColumnValue {
260 String(String),
261 Bool(bool),
262 Char(Option<char>),
263 Date(Option<NaiveDate>),
264 DateTime(Option<DateTime<Utc>>),
265 StringArray(Vec<String>),
266}
267
268fn get_column_value(task: &Task, column: &str) -> ColumnValue {
269 match column {
270 "description" => ColumnValue::String(task.description.clone()),
271 "status" => ColumnValue::String(format!("{:?}", task.status).to_lowercase()),
272 "completed" => ColumnValue::Bool(task.completed),
273 "priority" => ColumnValue::Char(task.priority),
274 "created_at" => ColumnValue::DateTime(Some(task.created_at)),
275 "updated_at" => ColumnValue::DateTime(Some(task.updated_at)),
276 "completed_at" => ColumnValue::DateTime(task.completed_at),
277 "due_date" => ColumnValue::Date(task.due_date),
278 "projects" => ColumnValue::StringArray(task.projects.clone()),
279 "contexts" => ColumnValue::StringArray(task.contexts.clone()),
280 _ => ColumnValue::String(String::new()),
281 }
282}
283
284fn sort_tasks(tasks: &mut Vec<&Task>, column: &str, desc: bool) {
285 let cmp = |a: &Task, b: &Task| -> std::cmp::Ordering {
286 match column {
287 "created_at" => a.created_at.cmp(&b.created_at),
288 "updated_at" => a.updated_at.cmp(&b.updated_at),
289 "due_date" => a.due_date.cmp(&b.due_date),
290 "priority" => a.priority.cmp(&b.priority),
291 "description" => a.description.cmp(&b.description),
292 "status" => format!("{:?}", a.status).cmp(&format!("{:?}", b.status)),
293 _ => std::cmp::Ordering::Equal,
294 }
295 };
296 if desc {
297 tasks.sort_by(|a, b| cmp(b, a));
298 } else {
299 tasks.sort_by(|a, b| cmp(a, b));
300 }
301}
302
303#[instrument(skip(state, _user), fields(query = %req.query))]
304pub async fn handle_query(
305 State(state): State<AuthState>,
306 Extension(_user): Extension<AuthUser>,
307 Json(req): Json<QueryRequest>,
308) -> Result<Json<QueryResponse>, (StatusCode, String)> {
309 let ai = state.ai_provider.as_ref().ok_or_else(|| {
310 tracing::warn!("AI not configured for query endpoint");
311 (StatusCode::SERVICE_UNAVAILABLE, "AI not configured".into())
312 })?;
313
314 let sql_response = ai
316 .chat(
317 SQL_GEN_PROMPT,
318 &[crate::ai::ChatMessage {
319 role: crate::ai::ChatRole::User,
320 content: req.query.clone(),
321 }],
322 )
323 .await
324 .map_err(|e| {
325 tracing::error!(error = %e, "SQL generation failed");
326 (
327 StatusCode::INTERNAL_SERVER_ERROR,
328 format!("SQL generation failed: {e}"),
329 )
330 })?;
331
332 let sql = extract_sql(&sql_response);
333 tracing::info!(sql = %sql, "generated SQL");
334
335 let repo = InMemoryTaskRepository::shared();
337 let all_tasks = repo.list(0).await.map_err(|e| {
338 tracing::error!(error = %e, "list tasks for query failed");
339 (StatusCode::INTERNAL_SERVER_ERROR, e.to_string())
340 })?;
341
342 let results = execute_sql(&sql, &all_tasks);
343 tracing::info!(count = results.len(), "SQL executed");
344
345 let answer_prompt = format!(
347 "{}{}\"\n\nLa consulta SQL ejecutada fue:\n```sql\n{}\n```\n\nResultados ({} filas):\n{}\n\nResponde en español de forma natural y concisa.",
348 NL_PROMPT_PREFIX,
349 req.query,
350 sql,
351 results.len(),
352 format_results(&results),
353 );
354
355 let answer = ai
356 .chat(
357 "Eres un asistente útil que responde en español.",
358 &[crate::ai::ChatMessage {
359 role: crate::ai::ChatRole::User,
360 content: answer_prompt,
361 }],
362 )
363 .await
364 .map_err(|e| {
365 tracing::error!(error = %e, "answer generation failed");
366 (
367 StatusCode::INTERNAL_SERVER_ERROR,
368 format!("answer generation failed: {e}"),
369 )
370 })?;
371
372 Ok(Json(QueryResponse {
373 sql,
374 results,
375 answer,
376 }))
377}
378
379fn format_results(tasks: &[Task]) -> String {
380 if tasks.is_empty() {
381 return "(sin resultados)".into();
382 }
383 tasks
384 .iter()
385 .map(|t| {
386 let status = format!("{:?}", t.status).to_lowercase();
387 let prio = t.priority.map_or(String::new(), |p| format!(" ({p})"));
388 let projects = if t.projects.is_empty() {
389 String::new()
390 } else {
391 format!(" +{}", t.projects.join(" +"))
392 };
393 let contexts = if t.contexts.is_empty() {
394 String::new()
395 } else {
396 format!(" @{}", t.contexts.join(" @"))
397 };
398 let due = t
399 .due_date
400 .map_or(String::new(), |d| format!(" [vence: {d}]"));
401 format!(
402 "- [{status}]{prio} {}{projects}{contexts}{due}",
403 t.description
404 )
405 })
406 .collect::<Vec<_>>()
407 .join("\n")
408}
409
410pub fn query_routes(state: AuthState) -> axum::Router<AuthState> {
411 axum::Router::new()
412 .route("/api/query", axum::routing::post(handle_query))
413 .layer(middleware::from_fn_with_state(state, require_auth))
414}