1use async_trait::async_trait;
4use serde::Deserialize;
5
6use super::base::{SearchOptions, SearchProvider, SearchResult};
7use crate::error::SearchError;
8
9#[derive(Debug, Deserialize)]
11struct GoogleApiResponse {
12 items: Option<Vec<GoogleApiItem>>,
13}
14
15#[derive(Debug, Deserialize)]
16struct GoogleApiItem {
17 title: String,
18 link: String,
19 snippet: Option<String>,
20}
21
22#[derive(Debug, Clone, Default)]
24pub struct GoogleConfig {
25 pub api_key: Option<String>,
27 pub search_engine_id: Option<String>,
29}
30
31pub struct GoogleProvider {
33 name: String,
34 enabled: bool,
35 weight: f64,
36 client: reqwest::Client,
37 config: GoogleConfig,
38 api_url: String,
39}
40
41impl GoogleProvider {
42 pub fn new(config: GoogleConfig) -> Self {
44 Self {
45 name: "google".to_string(),
46 enabled: true,
47 weight: 1.0,
48 client: reqwest::Client::builder()
49 .user_agent("Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36")
50 .build()
51 .expect("Failed to create HTTP client"),
52 config,
53 api_url: "https://www.googleapis.com/customsearch/v1".to_string(),
54 }
55 }
56
57 pub fn from_env() -> Self {
59 Self::new(GoogleConfig {
60 api_key: std::env::var("GOOGLE_API_KEY").ok(),
61 search_engine_id: std::env::var("GOOGLE_CX").ok(),
62 })
63 }
64
65 pub fn has_api_credentials(&self) -> bool {
67 self.config.api_key.is_some() && self.config.search_engine_id.is_some()
68 }
69
70 async fn search_with_api(
71 &self,
72 query: &str,
73 options: &SearchOptions,
74 ) -> Result<Vec<SearchResult>, SearchError> {
75 let api_key = self.config.api_key.as_ref().unwrap();
76 let cx = self.config.search_engine_id.as_ref().unwrap();
77 let limit = options.limit.unwrap_or(10).min(10);
78
79 let mut url = format!(
80 "{}?key={}&cx={}&q={}&num={}",
81 self.api_url,
82 api_key,
83 cx,
84 urlencoding::encode(query),
85 limit
86 );
87
88 if let Some(ref lang) = options.language {
89 url.push_str(&format!("&lr=lang_{}", lang));
90 }
91
92 if let Some(ref region) = options.region {
93 url.push_str(&format!("&gl={}", region));
94 }
95
96 if let Some(safe) = options.safe_search {
97 url.push_str(&format!("&safe={}", if safe { "active" } else { "off" }));
98 }
99
100 let response = self.client.get(&url).send().await?;
101
102 if !response.status().is_success() {
103 let error_text = response.text().await.unwrap_or_default();
104 return Err(SearchError::ApiError {
105 provider: self.name.clone(),
106 message: error_text,
107 });
108 }
109
110 let api_response: GoogleApiResponse = response.json().await?;
111
112 let results = api_response
113 .items
114 .unwrap_or_default()
115 .into_iter()
116 .enumerate()
117 .map(|(i, item)| SearchResult {
118 title: item.title,
119 url: item.link,
120 snippet: item.snippet.unwrap_or_default(),
121 source: self.name.clone(),
122 rank: i + 1,
123 score: None,
124 sources: None,
125 })
126 .collect();
127
128 Ok(results)
129 }
130
131 async fn search_with_scraping(
132 &self,
133 query: &str,
134 options: &SearchOptions,
135 ) -> Result<Vec<SearchResult>, SearchError> {
136 let limit = options.limit.unwrap_or(10);
137 let mut url = format!(
138 "https://www.google.com/search?q={}&num={}",
139 urlencoding::encode(query),
140 limit.min(20)
141 );
142
143 if let Some(ref lang) = options.language {
144 url.push_str(&format!("&hl={}", lang));
145 }
146
147 if let Some(ref region) = options.region {
148 url.push_str(&format!("&gl={}", region));
149 }
150
151 let response = self
152 .client
153 .get(&url)
154 .header(
155 "Accept",
156 "text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8",
157 )
158 .header("Accept-Language", "en-US,en;q=0.5")
159 .send()
160 .await?;
161
162 if !response.status().is_success() {
163 return Err(SearchError::ApiError {
164 provider: self.name.clone(),
165 message: format!("HTTP {}", response.status()),
166 });
167 }
168
169 let html = response.text().await?;
170 Ok(self.parse_scraped_results(&html, limit))
171 }
172
173 fn parse_scraped_results(&self, html: &str, limit: usize) -> Vec<SearchResult> {
174 use scraper::{Html, Selector};
175
176 let document = Html::parse_document(html);
177 let mut results = Vec::new();
178 let mut seen_urls = std::collections::HashSet::new();
179
180 let link_selector = Selector::parse("a").unwrap();
181 let h3_selector = Selector::parse("h3").unwrap();
182
183 for element in document.select(&link_selector) {
184 if results.len() >= limit {
185 break;
186 }
187
188 let href = element.value().attr("href").unwrap_or_default();
189 let url = if href.starts_with("/url?q=") {
190 href.strip_prefix("/url?q=")
191 .and_then(|u| u.split('&').next())
192 .map(|u| {
193 urlencoding::decode(u)
194 .unwrap_or_else(|_| u.into())
195 .to_string()
196 })
197 } else if href.starts_with("http") && !href.contains("google.com") {
198 Some(href.to_string())
199 } else {
200 None
201 };
202
203 if let Some(url) = url {
204 if seen_urls.contains(&url) || url.contains("google.com") {
205 continue;
206 }
207
208 let title = element
209 .select(&h3_selector)
210 .next()
211 .map(|h3| h3.text().collect::<String>())
212 .unwrap_or_default()
213 .trim()
214 .to_string();
215
216 if title.is_empty() {
217 continue;
218 }
219
220 seen_urls.insert(url.clone());
221 results.push(SearchResult {
222 title,
223 url,
224 snippet: String::new(),
225 source: self.name.clone(),
226 rank: results.len() + 1,
227 score: None,
228 sources: None,
229 });
230 }
231 }
232
233 results
234 }
235}
236
237impl Default for GoogleProvider {
238 fn default() -> Self {
239 Self::from_env()
240 }
241}
242
243#[async_trait]
244impl SearchProvider for GoogleProvider {
245 fn name(&self) -> &str {
246 &self.name
247 }
248
249 fn is_available(&self) -> bool {
250 self.enabled
251 }
252
253 fn weight(&self) -> f64 {
254 self.weight
255 }
256
257 fn set_weight(&mut self, weight: f64) {
258 self.weight = weight.clamp(0.0, 1.0);
259 }
260
261 fn set_enabled(&mut self, enabled: bool) {
262 self.enabled = enabled;
263 }
264
265 async fn search(
266 &self,
267 query: &str,
268 options: &SearchOptions,
269 ) -> Result<Vec<SearchResult>, SearchError> {
270 if query.is_empty() {
271 return Ok(Vec::new());
272 }
273
274 if self.has_api_credentials() {
275 match self.search_with_api(query, options).await {
276 Ok(results) => return Ok(results),
277 Err(e) => {
278 tracing::warn!("Google API search failed, falling back to scraping: {}", e);
279 }
280 }
281 }
282
283 self.search_with_scraping(query, options).await
284 }
285}