1use anyhow::{Context, Result};
2use chrono::{DateTime, Utc};
3use rusqlite::{params, Connection};
4use uuid::Uuid;
5
6#[derive(Debug, Clone, PartialEq, Eq)]
7pub enum AnnotationType {
8 Comment,
9 Classification,
10 Note,
11}
12
13impl AnnotationType {
14 pub fn as_str(&self) -> &str {
15 match self {
16 AnnotationType::Comment => "Comment",
17 AnnotationType::Classification => "Classification",
18 AnnotationType::Note => "Note",
19 }
20 }
21
22 #[allow(clippy::should_implement_trait)]
23 pub fn from_str(s: &str) -> Result<Self> {
24 match s {
25 "Comment" => Ok(AnnotationType::Comment),
26 "Classification" => Ok(AnnotationType::Classification),
27 "Note" => Ok(AnnotationType::Note),
28 _ => anyhow::bail!("Invalid annotation type: {}", s),
29 }
30 }
31}
32
33#[derive(Debug, Clone)]
34pub struct Annotation {
35 pub id: String,
36 pub sighting_id: i64,
37 pub user_id: String,
38 pub text: String,
39 pub annotation_type: AnnotationType,
40 pub created_at: DateTime<Utc>,
41 pub updated_at: DateTime<Utc>,
42}
43
44#[derive(Debug, Clone)]
45pub struct SightingRating {
46 pub id: String,
47 pub sighting_id: i64,
48 pub user_id: String,
49 pub confidence: i32,
50 pub notes: Option<String>,
51 pub created_at: DateTime<Utc>,
52}
53
54pub struct AnnotationStore<'a> {
55 conn: &'a Connection,
56}
57
58impl<'a> AnnotationStore<'a> {
59 pub fn new(conn: &'a Connection) -> Self {
60 AnnotationStore { conn }
61 }
62
63 pub fn init_annotation_schema(&self) -> Result<()> {
64 self.conn
65 .execute_batch(
66 "CREATE TABLE IF NOT EXISTS annotations (
67 id TEXT PRIMARY KEY,
68 sighting_id INTEGER NOT NULL,
69 user_id TEXT NOT NULL,
70 text TEXT NOT NULL,
71 annotation_type TEXT NOT NULL DEFAULT 'Comment',
72 created_at TEXT NOT NULL,
73 updated_at TEXT NOT NULL,
74 FOREIGN KEY (sighting_id) REFERENCES sightings(id),
75 FOREIGN KEY (user_id) REFERENCES users(id)
76 );
77
78 CREATE TABLE IF NOT EXISTS sighting_ratings (
79 id TEXT PRIMARY KEY,
80 sighting_id INTEGER NOT NULL,
81 user_id TEXT NOT NULL,
82 confidence INTEGER NOT NULL CHECK(confidence >= 1 AND confidence <= 5),
83 notes TEXT,
84 created_at TEXT NOT NULL,
85 FOREIGN KEY (sighting_id) REFERENCES sightings(id),
86 FOREIGN KEY (user_id) REFERENCES users(id),
87 UNIQUE(sighting_id, user_id)
88 );
89
90 CREATE INDEX IF NOT EXISTS idx_annotations_sighting ON annotations(sighting_id);
91 CREATE INDEX IF NOT EXISTS idx_annotations_user ON annotations(user_id);
92 CREATE INDEX IF NOT EXISTS idx_ratings_sighting ON sighting_ratings(sighting_id);",
93 )
94 .context("Failed to create annotation schema")?;
95 Ok(())
96 }
97
98 pub fn create_annotation(
99 &self,
100 sighting_id: i64,
101 user_id: &str,
102 text: &str,
103 annotation_type: AnnotationType,
104 ) -> Result<Annotation> {
105 let id = Uuid::new_v4().to_string();
106 let now = Utc::now();
107
108 self.conn
109 .execute(
110 "INSERT INTO annotations (id, sighting_id, user_id, text, annotation_type, created_at, updated_at)
111 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
112 params![
113 id,
114 sighting_id,
115 user_id,
116 text,
117 annotation_type.as_str(),
118 now.to_rfc3339(),
119 now.to_rfc3339(),
120 ],
121 )
122 .context("Failed to insert annotation")?;
123
124 Ok(Annotation {
125 id,
126 sighting_id,
127 user_id: user_id.to_string(),
128 text: text.to_string(),
129 annotation_type,
130 created_at: now,
131 updated_at: now,
132 })
133 }
134
135 pub fn get_annotations_for_sighting(&self, sighting_id: i64) -> Result<Vec<Annotation>> {
136 let mut stmt = self.conn.prepare(
137 "SELECT id, sighting_id, user_id, text, annotation_type, created_at, updated_at
138 FROM annotations WHERE sighting_id = ?1 ORDER BY created_at ASC",
139 )?;
140
141 let annotations = stmt
142 .query_map(params![sighting_id], |row| {
143 Ok(Annotation {
144 id: row.get(0)?,
145 sighting_id: row.get(1)?,
146 user_id: row.get(2)?,
147 text: row.get(3)?,
148 annotation_type: AnnotationType::from_str(&row.get::<_, String>(4)?)
149 .unwrap_or(AnnotationType::Comment),
150 created_at: {
151 let s: String = row.get(5)?;
152 DateTime::parse_from_rfc3339(&s)
153 .map(|dt| dt.with_timezone(&Utc))
154 .unwrap_or_else(|_| Utc::now())
155 },
156 updated_at: {
157 let s: String = row.get(6)?;
158 DateTime::parse_from_rfc3339(&s)
159 .map(|dt| dt.with_timezone(&Utc))
160 .unwrap_or_else(|_| Utc::now())
161 },
162 })
163 })?
164 .collect::<Result<Vec<_>, _>>()?;
165
166 Ok(annotations)
167 }
168
169 pub fn update_annotation(&self, annotation_id: &str, text: &str) -> Result<()> {
170 let now = Utc::now();
171 let affected = self.conn.execute(
172 "UPDATE annotations SET text = ?1, updated_at = ?2 WHERE id = ?3",
173 params![text, now.to_rfc3339(), annotation_id],
174 )?;
175
176 if affected == 0 {
177 anyhow::bail!("Annotation not found");
178 }
179 Ok(())
180 }
181
182 pub fn delete_annotation(&self, annotation_id: &str) -> Result<()> {
183 let affected = self.conn.execute(
184 "DELETE FROM annotations WHERE id = ?1",
185 params![annotation_id],
186 )?;
187
188 if affected == 0 {
189 anyhow::bail!("Annotation not found");
190 }
191 Ok(())
192 }
193
194 pub fn add_rating(
195 &self,
196 sighting_id: i64,
197 user_id: &str,
198 confidence: i32,
199 notes: Option<&str>,
200 ) -> Result<SightingRating> {
201 if !(1..=5).contains(&confidence) {
202 anyhow::bail!("Confidence must be between 1 and 5");
203 }
204
205 let id = Uuid::new_v4().to_string();
206 let now = Utc::now();
207
208 self.conn
209 .execute(
210 "INSERT OR REPLACE INTO sighting_ratings (id, sighting_id, user_id, confidence, notes, created_at)
211 VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
212 params![
213 id,
214 sighting_id,
215 user_id,
216 confidence,
217 notes,
218 now.to_rfc3339(),
219 ],
220 )
221 .context("Failed to insert rating")?;
222
223 Ok(SightingRating {
224 id,
225 sighting_id,
226 user_id: user_id.to_string(),
227 confidence,
228 notes: notes.map(|s| s.to_string()),
229 created_at: now,
230 })
231 }
232
233 pub fn get_ratings_for_sighting(&self, sighting_id: i64) -> Result<Vec<SightingRating>> {
234 let mut stmt = self.conn.prepare(
235 "SELECT id, sighting_id, user_id, confidence, notes, created_at
236 FROM sighting_ratings WHERE sighting_id = ?1",
237 )?;
238
239 let ratings = stmt
240 .query_map(params![sighting_id], |row| {
241 Ok(SightingRating {
242 id: row.get(0)?,
243 sighting_id: row.get(1)?,
244 user_id: row.get(2)?,
245 confidence: row.get(3)?,
246 notes: row.get(4)?,
247 created_at: {
248 let s: String = row.get(5)?;
249 DateTime::parse_from_rfc3339(&s)
250 .map(|dt| dt.with_timezone(&Utc))
251 .unwrap_or_else(|_| Utc::now())
252 },
253 })
254 })?
255 .collect::<Result<Vec<_>, _>>()?;
256
257 Ok(ratings)
258 }
259
260 pub fn get_average_rating(&self, sighting_id: i64) -> Result<Option<f64>> {
261 let result = self.conn.query_row(
262 "SELECT AVG(confidence) FROM sighting_ratings WHERE sighting_id = ?1",
263 params![sighting_id],
264 |row| row.get::<_, Option<f64>>(0),
265 )?;
266
267 Ok(result)
268 }
269}