NewStrcuture_Backup/rust-ai/src/main.rs

use axum::{
    extract::State,
    http::{HeaderMap, Method, StatusCode},
    routing::{get, post},
    Json, Router,
};
use tower_http::cors::{CorsLayer, AllowOrigin, Any};
use jsonwebtoken::{decode, DecodingKey, Validation, Algorithm};
use serde::{Deserialize, Serialize};
use sqlx::postgres::PgPoolOptions;
use std::collections::HashMap;
use std::fs;
use std::sync::Arc;
use tokio::sync::Mutex;
use tracing::{error, info, warn};
use uuid::Uuid;
use rand::Rng;
use std::time::Duration;
use std::time::{SystemTime, UNIX_EPOCH};

// ── JWT Claims ────────────────────────────────────────────────

#[derive(Debug, Deserialize)]
struct Claims {
    #[serde(rename = "userId")]
    user_id: String,
    role: String,
}

fn verify_jwt(token: &str, secret: &str) -> Option<Claims> {
    let key = DecodingKey::from_secret(secret.as_bytes());
    let validation = Validation::new(Algorithm::HS256);
    decode::<Claims>(token, &key, &validation).ok().map(|d| d.claims)
}

// ── Rate limiter ──────────────────────────────────────────────

struct RateLimiter {
    buckets: Mutex<HashMap<String, Vec<u64>>>,
    max_requests: usize,
    window_secs: u64,
}

impl RateLimiter {
    fn new(max_requests: usize, window_secs: u64) -> Self {
        Self {
            buckets: Mutex::new(HashMap::new()),
            max_requests,
            window_secs,
        }
    }

    async fn check(&self, key: &str) -> bool {
        let now = SystemTime::now().duration_since(UNIX_EPOCH).unwrap_or_default().as_secs();
        let mut buckets = self.buckets.lock().await;
        let timestamps = buckets.entry(key.to_string()).or_default();
        timestamps.retain(|t| now - *t <= self.window_secs);
        if timestamps.len() >= self.max_requests {
            false
        } else {
            timestamps.push(now);
            true
        }
    }
}

// ── Shared state ───────────────────────────────────────────────

struct AppState {
    db: sqlx::PgPool,
    ollama_url: String,
    model: String,
    jobs: Vec<Job>,
    leads: Arc<Mutex<LeadStore>>,
    http_client: reqwest::Client,
    jwt_secret: String,
    rate_limiter: RateLimiter,
}

#[derive(Debug, Clone, Serialize)]
struct Lead {
    title: String,
    url: String,
    source: String,
    found_at: u64,
    author: String,
    date: String,
    content: String,
}

struct LeadStore {
    leads: Vec<Lead>,
    max_size: usize,
}

impl LeadStore {
    fn new(max_size: usize) -> Self {
        Self { leads: Vec::new(), max_size }
    }

    fn push(&mut self, lead: Lead) {
        if !self.leads.iter().any(|l| l.url == lead.url) {
            self.leads.insert(0, lead);
            self.leads.truncate(self.max_size);
        }
    }

    fn recent(&self, max_age_secs: u64, limit: usize) -> Vec<Lead> {
        let now = SystemTime::now().duration_since(UNIX_EPOCH).unwrap_or_default().as_secs();
        self.leads.iter()
            .filter(|l| now.saturating_sub(l.found_at) <= max_age_secs)
            .take(limit)
            .cloned()
            .collect()
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
struct Job {
    job_title: String,
    keywords: Vec<String>,
    industry: String,
    description: String,
}

#[derive(Debug, Deserialize)]
struct ChatRequest {
    message: String,
}

#[derive(Debug, Serialize)]
struct ChatResponse {
    response: String,
}

#[derive(Debug, Serialize)]
struct JobsResponse {
    jobs: Vec<Job>,
}

#[derive(Debug, Serialize)]
struct HealthResponse {
    status: String,
    model: String,
}

// ── Ollama API types ───────────────────────────────────────────

#[derive(Debug, Serialize)]
struct OllamaChatMessage {
    role: String,
    content: String,
}

#[derive(Debug, Serialize)]
struct OllamaRequest {
    model: String,
    messages: Vec<OllamaChatMessage>,
    stream: bool,
    options: OllamaOptions,
}

#[derive(Debug, Serialize)]
struct OllamaOptions {
    temperature: f32,
    num_predict: u32,
}

#[derive(Debug, Deserialize)]
struct OllamaResponse {
    message: Option<OllamaResponseMessage>,
}

#[derive(Debug, Deserialize)]
struct OllamaResponseMessage {
    content: String,
}

// ── Helpers ────────────────────────────────────────────────────

fn truncate(s: &str, max: usize) -> String {
    s.chars().take(max).collect()
}

fn extract_claims(headers: &HeaderMap, state: &AppState) -> Result<Claims, (StatusCode, String)> {
    let auth_header = headers.get("Authorization").and_then(|v| v.to_str().ok()).unwrap_or("");
    let token = auth_header.strip_prefix("Bearer ").unwrap_or("");
    let claims = verify_jwt(token, &state.jwt_secret).ok_or_else(|| {
        (StatusCode::UNAUTHORIZED, "Unauthorized".to_string())
    })?;
    match claims.role.as_str() {
        "sales" | "admin" | "super_admin" => Ok(claims),
        _ => Err((StatusCode::FORBIDDEN, "Forbidden".to_string())),
    }
}

fn format_leads_output(leads: &[Lead]) -> String {
    if leads.is_empty() {
        return "No new requests found yet.".to_string();
    }
    leads
        .iter()
        .enumerate()
        .map(|(i, l)| {
            let author = if l.author.is_empty() { "Unknown" } else { &l.author };
            let date = truncate(&l.date, 10);
            format!(
                "{}. {}\n   {}\n   {}\n   {}",
                i + 1,
                author,
                date,
                l.title,
                l.url
            )
        })
        .collect::<Vec<_>>()
        .join("\n")
}

fn build_system_prompt(jobs: &[Job], leads: &[Lead]) -> String {
    let job_list: Vec<String> = jobs
        .iter()
        .map(|j| format!("- {} ({}): {}", j.job_title, j.industry, j.description))
        .collect();
    let job_list_str = job_list.join("\n");

    let lead_summary: Vec<String> = leads
        .iter()
        .map(|l| format!("{} | {} | {}", l.author, l.title, l.url))
        .collect();
    let lead_summary_str = if lead_summary.is_empty() {
        "None yet.".to_string()
    } else {
        lead_summary.join("\n")
    };

    format!(
        "You are a Sales AI Assistant for Coast IT CRM.\n\n\
         Available job categories to target:\n{}\n\n\
         Recent leads context:\n{}\n\n\
         Rules:\n\
         - When asked about leads, answer concisely under 150 words.\n\
         - If asked to suggest a sales strategy, give brief actionable advice.\n\
         - Be direct and professional. No fluff.",
        job_list_str, lead_summary_str
    )
}

// ── Chat handler ───────────────────────────────────────────────

async fn handle_chat(
    State(state): State<Arc<AppState>>,
    headers: HeaderMap,
    Json(req): Json<ChatRequest>,
) -> Result<Json<ChatResponse>, (StatusCode, String)> {
    let claims = extract_claims(&headers, &state)?;

    if !state.rate_limiter.check(&claims.user_id).await {
        return Err((StatusCode::TOO_MANY_REQUESTS, "Rate limit exceeded".to_string()));
    }

    let msg_lower = req.message.to_lowercase();
    let msg_words: Vec<&str> = msg_lower.split_whitespace().collect();
    let has_listing = msg_words.iter().any(|w| ["listings", "listing", "leads", "links", "lists"].contains(w));
    let has_show = msg_lower.contains("show me") || msg_lower.contains("give me") || msg_lower.contains("pull");
    let has_job = msg_words.contains(&"jobs") || msg_words.contains(&"job");
    if has_listing || (has_show && has_job) || (has_show && msg_lower.contains("links")) || msg_lower.contains("recent leads") {
        let now = SystemTime::now().duration_since(UNIX_EPOCH).unwrap_or_default().as_secs();
        let service_url = "http://localhost:3008/scrape/all".to_string();
        if let Ok(resp) = state.http_client
            .post(&service_url)
            .timeout(Duration::from_secs(120))
            .send()
            .await
        {
            if let Ok(leads_data) = resp.json::<Vec<serde_json::Value>>().await {
                info!("Scraped {} leads from {}", leads_data.len(), service_url);
                let mut store = state.leads.lock().await;
                for item in &leads_data {
                    store.push(Lead {
                        title: truncate(item["title"].as_str().unwrap_or(""), 120),
                        url: item["url"].as_str().unwrap_or("").to_string(),
                        source: item["source"].as_str().unwrap_or("unknown").to_string(),
                        found_at: now,
                        author: truncate(item["author"].as_str().unwrap_or(""), 60),
                        date: truncate(item["date"].as_str().unwrap_or(""), 30),
                        content: truncate(item["content"].as_str().unwrap_or(""), 300),
                    });
                }
            }
        } else {
            error!("Scraper service unreachable: {}", service_url);
        }

        let recent_leads = state.leads.lock().await.recent(604800, 20);
        let response = format_leads_output(&recent_leads);
        let _ = sqlx::query(
            "INSERT INTO ai_conversations (id, user_id, role, message, response) VALUES ($1, $2, $3, $4, $5)",
        )
        .bind(Uuid::new_v4())
        .bind(Uuid::parse_str(&claims.user_id).unwrap_or(Uuid::nil()))
        .bind(&claims.role)
        .bind(&req.message)
        .bind(&response)
        .execute(&state.db)
        .await;
        return Ok(Json(ChatResponse { response }));
    }

    let recent_leads = state.leads.lock().await.recent(604800, 20);
    let system_prompt = build_system_prompt(&state.jobs, &recent_leads);
    let message_text = req.message.clone();

    let ollama_req = OllamaRequest {
        model: state.model.clone(),
        messages: vec![
            OllamaChatMessage { role: "system".to_string(), content: system_prompt },
            OllamaChatMessage { role: "user".to_string(), content: req.message.clone() },
        ],
        stream: false,
        options: OllamaOptions { temperature: 0.7, num_predict: 1024 },
    };

    let resp = state.http_client
        .post(format!("{}/api/chat", state.ollama_url))
        .json(&ollama_req)
        .send()
        .await
        .map_err(|e| {
            error!("Ollama request failed: {}", e);
            (StatusCode::SERVICE_UNAVAILABLE, "AI service unavailable".to_string())
        })?;

    let ollama_resp: OllamaResponse = resp.json().await.map_err(|e| {
        error!("Failed to parse Ollama response: {}", e);
        (StatusCode::SERVICE_UNAVAILABLE, "AI response parse error".to_string())
    })?;

    let response_text = ollama_resp.message.map(|m| m.content).unwrap_or_default();

    let _ = sqlx::query(
        "INSERT INTO ai_conversations (id, user_id, role, message, response) VALUES ($1, $2, $3, $4, $5)",
    )
    .bind(Uuid::new_v4())
    .bind(Uuid::parse_str(&claims.user_id).unwrap_or(Uuid::nil()))
    .bind(&claims.role)
    .bind(&message_text)
    .bind(&response_text)
    .execute(&state.db)
    .await;

    Ok(Json(ChatResponse { response: response_text }))
}

// ── Jobs handler ───────────────────────────────────────────────

async fn handle_jobs(
    State(state): State<Arc<AppState>>,
    headers: HeaderMap,
) -> Result<Json<JobsResponse>, (StatusCode, String)> {
    let _claims = extract_claims(&headers, &state)?;
    if !state.rate_limiter.check(&_claims.user_id).await {
        return Err((StatusCode::TOO_MANY_REQUESTS, "Rate limit exceeded".to_string()));
    }
    Ok(Json(JobsResponse { jobs: state.jobs.clone() }))
}

// ── Health handler ─────────────────────────────────────────────

async fn handle_health(
    State(state): State<Arc<AppState>>,
) -> Json<HealthResponse> {
    Json(HealthResponse {
        status: "ok".to_string(),
        model: state.model.clone(),
    })
}

// ── Main ───────────────────────────────────────────────────────

#[tokio::main]
async fn main() {
    tracing_subscriber::fmt()
        .with_env_filter(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| "crm_ai=info,tower_http=info".into()),
        )
        .init();

    dotenvy::dotenv().ok();

    let database_url = std::env::var("DATABASE_URL").expect("DATABASE_URL must be set");
    let jwt_secret = std::env::var("JWT_SECRET").expect("JWT_SECRET must be set");
    let ollama_url = std::env::var("OLLAMA_BASE_URL").unwrap_or_else(|_| "http://localhost:11434".to_string());
    let model = std::env::var("AI_MODEL").unwrap_or_else(|_| "dolphin-phi".to_string());
    let host = std::env::var("AI_HOST").unwrap_or_else(|_| "127.0.0.1".to_string());
    let port: u16 = std::env::var("AI_PORT").unwrap_or_else(|_| "3001".to_string()).parse().expect("AI_PORT must be a number");

    let jobs_path = std::env::var("JOBS_PATH").unwrap_or_else(|_| "data/ai/jobs.jsonl".to_string());
    let jobs_content = fs::read_to_string(&jobs_path).unwrap_or_default();
    let jobs: Vec<Job> = jobs_content.lines().filter(|l| !l.trim().is_empty()).filter_map(|l| serde_json::from_str(l).ok()).collect();

    info!("Loaded {} job categories, model: {}, Ollama: {}", jobs.len(), model, ollama_url);

    let db = PgPoolOptions::new()
        .max_connections(20)
        .connect(&database_url)
        .await
        .expect("Failed to connect to database");

    info!("Connected to PostgreSQL");

    let http_client = reqwest::Client::builder()
        .timeout(Duration::from_secs(120))
        .build()
        .expect("Failed to build HTTP client");

    let lead_store = Arc::new(Mutex::new(LeadStore::new(100)));

    let state = Arc::new(AppState {
        db,
        ollama_url,
        model,
        jobs,
        leads: lead_store.clone(),
        http_client,
        jwt_secret,
        rate_limiter: RateLimiter::new(30, 60),
    });

    let cors = CorsLayer::new()
        .allow_origin(AllowOrigin::list([
            "http://localhost:3006".parse().unwrap(),
            "http://127.0.0.1:3006".parse().unwrap(),
        ]))
        .allow_methods([Method::GET, Method::POST])
        .allow_headers(Any);

    let app = Router::new()
        .route("/ai/chat", post(handle_chat))
        .route("/ai/jobs", get(handle_jobs))
        .route("/health", get(handle_health))
        .layer(cors)
        .with_state(state);

    let addr = format!("{}:{}", host, port);
    info!("CRM AI server listening on {}", addr);

    let listener = tokio::net::TcpListener::bind(&addr)
        .await
        .expect("Failed to bind address");

    let bg_leads = lead_store.clone();
    let bg_url = "http://localhost:3008/scrape/all".to_string();
    tokio::spawn(async move {
        let client = match reqwest::Client::builder()
            .timeout(Duration::from_secs(120))
            .build()
        {
            Ok(c) => c,
            Err(e) => {
                error!("Failed to build background HTTP client: {} — scraper disabled", e);
                return;
            }
        };
        loop {
            let now = SystemTime::now().duration_since(UNIX_EPOCH).unwrap_or_default().as_secs();
            match client.post(&bg_url).send().await {
                Ok(resp) => {
                    if resp.status().is_success() {
                        match resp.json::<Vec<serde_json::Value>>().await {
                            Ok(data) => {
                                let mut store = bg_leads.lock().await;
                                for item in &data {
                                    store.push(Lead {
                                        title: truncate(item["title"].as_str().unwrap_or(""), 120),
                                        url: item["url"].as_str().unwrap_or("").to_string(),
                                        source: item["source"].as_str().unwrap_or("unknown").to_string(),
                                        found_at: now,
                                        author: truncate(item["author"].as_str().unwrap_or(""), 60),
                                        date: truncate(item["date"].as_str().unwrap_or(""), 30),
                                        content: truncate(item["content"].as_str().unwrap_or(""), 300),
                                    });
                                }
                            }
                            Err(e) => {
                                warn!("Failed to parse scraper JSON: {}", e);
                            }
                        }
                    } else {
                        warn!("Scraper returned status: {}", resp.status());
                    }
                }
                Err(e) => {
                    warn!("Scraper request failed: {}", e);
                }
            }
            let delay = rand::thread_rng().gen_range(120..300);
            tokio::time::sleep(Duration::from_secs(delay)).await;
        }
    });

    axum::serve(listener, app)
        .await
        .expect("Server failed");
}