// SPDX-FileCopyrightText: 2024 Phil Höfer <phil@suma-ev.de>
// SPDX-License-Identifier: AGPL-3.0-only
use core::str;
use std::collections::HashMap;
use std::fs::File;
use std::io::{self, BufRead, BufReader};
use std::error::Error;
use std::str::FromStr;
use std::fs;
use std::time::SystemTime;
use importers::Importer;
use importers::file::FileImporter;
use predictors::basic_set::SetPredictor;
use predictors::composite::CompositePredictor;
use toml::Table;
use csv::ReaderBuilder;
use tiny_http::{Server, Response};
mod importers;
mod predictors;
use predictors::{basic_markov, basic_set, composite, Predictor};
use predictors::basic_markov::MarkovChain;
use postgres::{Client, NoTls};
fn main() -> Result<(), io::Error> {
let mut last_update = std::time::SystemTime::now();
let config_toml = read_config("data/config.toml").or(read_config("config.toml")).unwrap().parse::<Table>().unwrap();
let mut config = HashMap::new();
for (key, value) in config_toml {
match value {
toml::Value::String(s) => {
config.insert(key, s);
}
toml::Value::Integer(i) => {
config.insert(key, i.to_string());
}
_ => {
// Ignore other types
}
}
}
let mut markov_chain = CompositePredictor::new_from_config(config.clone());
for path in vec!["../../data/data.csv","data/data.csv","data.csv","data_full.csv"].into_iter() {
match &mut FileImporter::new_from_path(path) {
Ok(imp) => {
match markov_chain.update_from_importer(imp, 10000) {
Ok(count) => {
println!("Read {} queries from file {}.",count, path);
}
Err(e) => {
println!("Error while reading from file: {}", e)
}
};
},
_ => {}
}
}
// data locations: vec!["../../data/data.csv","data/data.csv","data.csv","data_full.csv"]
read_from_db(config.clone(), &mut markov_chain);
// let term_frequency_threshold = match config.get("term_frequency_threshold") {
// Some(toml::Value::Integer(n)) if *n >= 0 => *n as usize,
// _ => 2
// };
//let filtered_markov_chain = filter_markov_chain(&markov_chain,term_frequency_threshold);
// Print the Markov Chain for verification
// for (key, values) in &markov_chain {
// //println!("{}: {:?}", key, values);
// }
// Test the function
let word = "wie"; // replace with the word you want to check
let top_words = &markov_chain.predict(&word, 3);
println!("Example prediction for '{}': {:?}", word, top_words);
let server = Server::http("0.0.0.0:8000").unwrap();
for request in server.incoming_requests() {
process_request(request, config.clone(), &mut markov_chain, last_update)
}
Ok(())
}
fn process_request(request: tiny_http::Request, config: HashMap<String, String>, markov_chain: &mut CompositePredictor, mut last_update: SystemTime) {
// println!("received request! method: {:?}, url: {:?}, headers: {:?}",
// request.method(),
// request.url(),
// request.headers()
// );
let query = get_query(request.url());
//println!("got query:{}", query.clone().unwrap());
match query {
Ok(query) => {
match config.get("auth") {
Some(server_auth) => {
match get_authkey(request.url().clone()) {
Ok(client_auth) => {
if client_auth != server_auth.clone() {
println!("invalid auth:{}, server auth: {}", client_auth, server_auth);
request.respond(Response::from_string(""));
return;
}
},
_ => {}
}
},
_ => {}
}
let predict_count = match config.get("max_predict_count") {
Some(n) if n.parse::<usize>().unwrap_or_default() > 0 => n.parse::<usize>().unwrap_or(5),
_ => 5
};
let predictions = &markov_chain.predict(&query, predict_count);
//println!("Query: {}, Prediction:{}", query, prediction);
let response = Response::from_string(
format!("[\"{}\",[{}]]", query, predictions.join(",")));
request.respond(response);
let now = std::time::SystemTime::now();
let elapsed = now.duration_since(last_update).expect("Time went backwards");
if elapsed >= std::time::Duration::from_secs(24 * 60 * 60) {
markov_chain.decay();
read_from_db(config.clone(), markov_chain);
last_update = now;
}
},
Err(e) => {
//println!("Error: {}",e);
}
}
}
fn read_from_db(config: HashMap<String, String>, predictor: &mut CompositePredictor) {
let mut raw_list = String::new();
let blocklist: Vec<&str> = match config.get("blocklist") {
Some(list) => {
match fs::read_to_string(list) {
Ok(l) => {
raw_list = l.clone();
raw_list.lines().collect()
},
_ => Vec::new()
}
},
_ => Vec::new()
};
let default_password = &String::from("");
match (config.get("db_host"), config.get("db_password")) {
(Some(db_host), Some(db_password)) => {
let mut client = Client::connect(format!("host={} user=metager password={} dbname=postgres",db_host, db_password).as_str(), NoTls);
match client {
Ok(mut c) => {
let mut count = 0;
match c.query("SELECT query FROM public.logs_partitioned ORDER BY time DESC LIMIT 100000", &[]) {
Ok(rows) => {
for row in rows {
let query: &str = row.get(0);
predictor.update_from_query(query, &blocklist);
count += 1;
}
println!("{} queries read from DB", count);
},
Err(e) => {
println!("Error while reading from DB: {}", e);
}
}
},
Err(e) => {
println!("Error while connecting to DB: {}", e);
}
}
},
_ => {
println!("No DB config found. Skipping...");
}
}
}
fn read_config(file_path: &str) -> Result<String, io::Error> {
fs::read_to_string(file_path)
}
//TODO: unify get_X() functions
fn get_query(request_url: &str) -> Result<String, url::ParseError> {
let parsed_url = request_url.split_once('?').map_or(request_url, |(_, after)| after);
//println!("parsed_url:{}", parsed_url);
let query_pairs = url::form_urlencoded::parse(parsed_url.as_bytes());
for (key, value) in query_pairs {
//println!("key:{}, value: {}", key, value);
if key == "q" {
return Ok(value.into_owned());
}
}
Ok(String::from_str("").unwrap())
}
fn get_authkey(request_url: &str) -> Result<String, url::ParseError> {
let parsed_url = request_url.split_once('?').map_or(request_url, |(_, after)| after);
//println!("parsed_url:{}", parsed_url);
let query_pairs = url::form_urlencoded::parse(parsed_url.as_bytes());
for (key, value) in query_pairs {
//println!("key:{}, value: {}", key, value);
if key == "auth" {
return Ok(value.into_owned());
}
}
Ok(String::from_str("").unwrap())
}
fn filter_markov_chain(markov_chain: &MarkovChain, min_count: usize) -> MarkovChain {
let mut filtered_chain: MarkovChain = HashMap::new();
for (key, followers) in markov_chain {
let filtered_followers: HashMap<String, usize> = followers
.iter()
.filter(|&(_, &count)| count >= min_count)
.map(|(word, &count)| (word.clone(), count))
.collect();
if !filtered_followers.is_empty() {
filtered_chain.insert(key.clone(), filtered_followers);
}
}
filtered_chain
}