ImageApi/src/ai/sms_client.rs

use anyhow::Result;
use reqwest::Client;
use serde::Deserialize;

use super::ollama::OllamaClient;

#[derive(Clone)]
pub struct SmsApiClient {
    client: Client,
    base_url: String,
    token: Option<String>,
}

impl SmsApiClient {
    pub fn new(base_url: String, token: Option<String>) -> Self {
        Self {
            client: Client::new(),
            base_url,
            token,
        }
    }

    /// Compute a `[start, end]` unix-second window of `2 * radius_days`
    /// centered on `center_ts`. `radius_days < 1` is clamped to 1 to avoid
    /// degenerate zero-width windows.
    pub(crate) fn window_for_radius(center_ts: i64, radius_days: i64) -> (i64, i64) {
        let r = radius_days.max(1);
        let span = r * 86400;
        (center_ts - span, center_ts + span)
    }

    /// Fetch messages for a specific contact within ±`radius_days` of the
    /// given timestamp. Falls back to all contacts when no messages found
    /// for the named contact. Sorted by proximity to the center timestamp.
    pub async fn fetch_messages_for_contact(
        &self,
        contact: Option<&str>,
        center_timestamp: i64,
        radius_days: i64,
    ) -> Result<Vec<SmsMessage>> {
        let effective_radius = radius_days.max(1);
        let (start_ts, end_ts) = Self::window_for_radius(center_timestamp, radius_days);

        let center_dt = chrono::DateTime::from_timestamp(center_timestamp, 0)
            .ok_or_else(|| anyhow::anyhow!("Invalid timestamp"))?;

        // If contact specified, try fetching for that contact first
        if let Some(contact_name) = contact {
            log::info!(
                "Fetching SMS for contact: {} (±{} days from {})",
                contact_name,
                effective_radius,
                center_dt.format("%Y-%m-%d %H:%M:%S")
            );
            let messages = self
                .fetch_messages(start_ts, end_ts, Some(contact_name), Some(center_timestamp))
                .await?;

            if !messages.is_empty() {
                log::info!(
                    "Found {} messages for contact {}",
                    messages.len(),
                    contact_name
                );
                return Ok(messages);
            }

            log::info!(
                "No messages found for contact {}, falling back to all contacts",
                contact_name
            );
        }

        // Fallback to all contacts
        log::info!(
            "Fetching all SMS messages (±{} days from {})",
            effective_radius,
            center_dt.format("%Y-%m-%d %H:%M:%S")
        );
        self.fetch_messages(start_ts, end_ts, None, Some(center_timestamp))
            .await
    }

    /// Fetch all messages for a specific contact across all time
    /// Used for embedding generation - retrieves complete message history
    /// Handles pagination automatically if the API returns a limited number of results
    pub async fn fetch_all_messages_for_contact(&self, contact: &str) -> Result<Vec<SmsMessage>> {
        let start_ts = chrono::DateTime::parse_from_rfc3339("2000-01-01T00:00:00Z")
            .unwrap()
            .timestamp();
        let end_ts = chrono::Utc::now().timestamp();

        log::info!("Fetching all historical messages for contact: {}", contact);

        let mut all_messages = Vec::new();
        let mut offset = 0;
        let limit = 1000; // Fetch in batches of 1000

        loop {
            log::debug!(
                "Fetching batch at offset {} for contact {}",
                offset,
                contact
            );

            let batch = self
                .fetch_messages_paginated(start_ts, end_ts, Some(contact), None, limit, offset)
                .await?;

            let batch_size = batch.len();
            all_messages.extend(batch);

            log::debug!(
                "Fetched {} messages (total so far: {})",
                batch_size,
                all_messages.len()
            );

            // If we got fewer messages than the limit, we've reached the end
            if batch_size < limit {
                break;
            }

            offset += limit;
        }

        log::info!(
            "Fetched {} total messages for contact {}",
            all_messages.len(),
            contact
        );

        Ok(all_messages)
    }

    /// Internal method to fetch messages with pagination support
    async fn fetch_messages_paginated(
        &self,
        start_ts: i64,
        end_ts: i64,
        contact: Option<&str>,
        center_timestamp: Option<i64>,
        limit: usize,
        offset: usize,
    ) -> Result<Vec<SmsMessage>> {
        let mut url = format!(
            "{}/api/messages/by-date-range/?start_date={}&end_date={}&limit={}&offset={}",
            self.base_url, start_ts, end_ts, limit, offset
        );

        if let Some(contact_name) = contact {
            url.push_str(&format!("&contact={}", urlencoding::encode(contact_name)));
        }

        if let Some(ts) = center_timestamp {
            url.push_str(&format!("&timestamp={}", ts));
        }

        log::debug!("Fetching SMS messages from: {}", url);

        let mut request = self.client.get(&url);

        if let Some(token) = &self.token {
            request = request.header("Authorization", format!("Bearer {}", token));
        }

        let response = request.send().await?;

        log::debug!("SMS API response status: {}", response.status());

        if !response.status().is_success() {
            let status = response.status();
            let error_body = response.text().await.unwrap_or_default();
            log::error!("SMS API request failed: {} - {}", status, error_body);
            return Err(anyhow::anyhow!(
                "SMS API request failed: {} - {}",
                status,
                error_body
            ));
        }

        let data: SmsApiResponse = response.json().await?;

        Ok(data
            .messages
            .into_iter()
            .map(|m| SmsMessage {
                contact: m.contact_name,
                body: m.body,
                timestamp: m.date,
                is_sent: m.type_ == 2,
            })
            .collect())
    }

    /// Internal method to fetch messages with optional contact filter and timestamp sorting
    async fn fetch_messages(
        &self,
        start_ts: i64,
        end_ts: i64,
        contact: Option<&str>,
        center_timestamp: Option<i64>,
    ) -> Result<Vec<SmsMessage>> {
        // Call Django endpoint
        let mut url = format!(
            "{}/api/messages/by-date-range/?start_date={}&end_date={}",
            self.base_url, start_ts, end_ts
        );

        // Add contact filter if provided
        if let Some(contact_name) = contact {
            url.push_str(&format!("&contact={}", urlencoding::encode(contact_name)));
        }

        // Add timestamp for proximity sorting if provided
        if let Some(ts) = center_timestamp {
            url.push_str(&format!("&timestamp={}", ts));
        }

        log::debug!("Fetching SMS messages from: {}", url);

        let mut request = self.client.get(&url);

        // Add authorization header if token exists
        if let Some(token) = &self.token {
            request = request.header("Authorization", format!("Bearer {}", token));
        }

        let response = request.send().await?;

        log::debug!("SMS API response status: {}", response.status());

        if !response.status().is_success() {
            let status = response.status();
            let error_body = response.text().await.unwrap_or_default();
            log::error!("SMS API request failed: {} - {}", status, error_body);
            return Err(anyhow::anyhow!(
                "SMS API request failed: {} - {}",
                status,
                error_body
            ));
        }

        let data: SmsApiResponse = response.json().await?;

        // Convert to internal format
        Ok(data
            .messages
            .into_iter()
            .map(|m| SmsMessage {
                contact: m.contact_name,
                body: m.body,
                timestamp: m.date,
                is_sent: m.type_ == 2, // type 2 = sent
            })
            .collect())
    }

    /// Search message bodies via the Django side's FTS5 / semantic / hybrid
    /// endpoint. `mode` selects the ranking strategy:
    ///   - "fts5"     keyword-only, supports phrase / prefix / boolean / NEAR
    ///   - "semantic" embedding similarity
    ///   - "hybrid"   both merged via reciprocal rank fusion (recommended)
    ///
    /// Equivalent to `search_messages_with_contact(query, mode, limit, None)`;
    /// kept as a convenience for callers that don't filter by contact.
    pub async fn search_messages(
        &self,
        query: &str,
        mode: &str,
        limit: usize,
    ) -> Result<Vec<SmsSearchHit>> {
        self.search_messages_with_contact(query, mode, limit, None).await
    }

    /// Same shape as `search_messages` but with optional `contact_id`. The
    /// SMS-API endpoint accepts contact_id natively; date filtering is the
    /// caller's responsibility (post-filter on the returned rows).
    pub async fn search_messages_with_contact(
        &self,
        query: &str,
        mode: &str,
        limit: usize,
        contact_id: Option<i64>,
    ) -> Result<Vec<SmsSearchHit>> {
        let mut url = format!(
            "{}/api/messages/search/?q={}&mode={}&limit={}",
            self.base_url,
            urlencoding::encode(query),
            urlencoding::encode(mode),
            limit
        );
        if let Some(cid) = contact_id {
            url.push_str(&format!("&contact_id={}", cid));
        }

        let mut request = self.client.get(&url);
        if let Some(token) = &self.token {
            request = request.header("Authorization", format!("Bearer {}", token));
        }

        let response = request.send().await?;
        if !response.status().is_success() {
            let status = response.status();
            let body = response.text().await.unwrap_or_default();
            return Err(anyhow::anyhow!(
                "SMS search request failed: {} - {}",
                status,
                body
            ));
        }

        let data: SmsSearchResponse = response.json().await?;
        Ok(data.results)
    }

    pub async fn summarize_context(
        &self,
        messages: &[SmsMessage],
        ollama: &OllamaClient,
    ) -> Result<String> {
        if messages.is_empty() {
            return Ok(String::from("No messages on this day"));
        }

        // Create prompt for Ollama with sender/receiver distinction
        let user_name = crate::ai::user_display_name();
        let messages_text: String = messages
            .iter()
            .take(60) // Limit to avoid token overflow
            .map(|m| {
                if m.is_sent {
                    format!("{}: {}", user_name, m.body)
                } else {
                    format!("{}: {}", m.contact, m.body)
                }
            })
            .collect::<Vec<_>>()
            .join("\n");

        let prompt = format!(
            r#"Summarize these messages in up to 4-5 sentences. Focus on key topics, places, people mentioned, and the overall context of the conversations.

Messages:
{}

Summary:"#,
            messages_text
        );

        ollama
            .generate(
                &prompt,
                // Some("You are a summarizer for the purposes of jogging my memory and highlighting events and situations."),
                Some("You are the keeper of memories, ingest the context and give me a casual summary of the moment."),
            )
            .await
    }
}

#[derive(Debug, Clone)]
pub struct SmsMessage {
    pub contact: String,
    pub body: String,
    pub timestamp: i64,
    pub is_sent: bool,
}

#[derive(Deserialize)]
struct SmsApiResponse {
    messages: Vec<SmsApiMessage>,
}

#[derive(Deserialize)]
struct SmsApiMessage {
    contact_name: String,
    body: String,
    date: i64,
    #[serde(rename = "type")]
    type_: i32,
}

#[derive(Debug, Clone, Deserialize)]
pub struct SmsSearchHit {
    #[allow(dead_code)]
    pub message_id: i64,
    pub contact_name: String,
    #[allow(dead_code)]
    pub contact_address: String,
    pub body: String,
    pub date: i64,
    /// Message direction code: 1 = received, 2 = sent.
    #[serde(rename = "type")]
    pub type_: i32,
    /// Present for semantic / hybrid modes; absent for fts5.
    #[serde(default)]
    pub similarity_score: Option<f32>,
}

#[derive(Deserialize)]
struct SmsSearchResponse {
    results: Vec<SmsSearchHit>,
    #[allow(dead_code)]
    #[serde(default)]
    search_method: String,
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn window_for_radius_produces_2n_day_span() {
        let center: i64 = 1_700_000_000;
        let (start, end) = SmsApiClient::window_for_radius(center, 7);
        assert_eq!(end - start, 14 * 86400);
        assert_eq!(start + 7 * 86400, center);
        assert_eq!(end - 7 * 86400, center);
    }

    #[test]
    fn window_for_radius_clamps_zero_to_one() {
        let (start, end) = SmsApiClient::window_for_radius(100_000, 0);
        assert_eq!(end - start, 2 * 86400);
    }

    #[test]
    fn window_for_radius_clamps_negative_to_one() {
        let (start, end) = SmsApiClient::window_for_radius(100_000, -7);
        assert_eq!(end - start, 2 * 86400);
    }
}