ImageApi/src/ai/nl_query.rs

//! Natural-language → structured-query translation for unified photo search.
//!
//! The unified search endpoint (`/photos/search/unified`, Phase 2) needs to
//! turn a free-text query like *"sunset photos in Italy from last summer"*
//! into the structured filter the existing `/photos` engine understands plus
//! a semantic term for CLIP ranking. That translation is a single grounded
//! LLM call, isolated here so it can be unit-tested without a network or the
//! full `InsightGenerator`.
//!
//! Two-stage design:
//!  1. The LLM emits a [`RawNlQuery`] — references are by *name* (tags) and
//!     dates as ISO strings, never numeric ids it could hallucinate.
//!  2. [`resolve_raw_query`] maps names against the real tag vocabulary and
//!     converts ISO dates to unix seconds, producing a [`StructuredQuery`].
//!     A tag the model invents that isn't in the vocab is surfaced in
//!     `unmatched_tags` (the caller folds it back into the semantic term)
//!     rather than silently dropped — this is the anti-noise guard.
//!
//! Geocoding of `place` and person filtering are intentionally *not* handled
//! here: `place` stays as text for the caller to forward-geocode (async, see
//! `geo::forward_geocode`), and person filtering is deferred until a
//! person→photos resolver exists.

use crate::ai::llm_client::{ChatMessage, LlmClient, Tool, strip_think_blocks};
use anyhow::{Result, anyhow};
use serde::{Deserialize, Serialize};

/// Raw query object as emitted by the LLM. Tag references are by name
/// (resolved against the real vocab in Rust); dates are ISO `YYYY-MM-DD`.
/// Every field is optional so a partial / minimal model response still
/// deserializes.
#[derive(Debug, Clone, Default, Deserialize, PartialEq)]
pub struct RawNlQuery {
    /// Visual/scene description handed to CLIP for ranking. The descriptive
    /// remainder after structured filters are peeled off.
    #[serde(default)]
    pub semantic: Option<String>,
    /// Tag names the photos must have. Matched case-insensitively against
    /// the supplied vocabulary; non-matches land in `unmatched_tags`.
    #[serde(default)]
    pub tags: Vec<String>,
    /// Tag names the photos must NOT have.
    #[serde(default)]
    pub exclude_tags: Vec<String>,
    #[serde(default)]
    pub camera_make: Option<String>,
    #[serde(default)]
    pub camera_model: Option<String>,
    #[serde(default)]
    pub lens_model: Option<String>,
    /// Free-text place/location name to forward-geocode (e.g. "Italy").
    #[serde(default)]
    pub place: Option<String>,
    /// Inclusive start date, ISO `YYYY-MM-DD`.
    #[serde(default)]
    pub date_from: Option<String>,
    /// Inclusive end date, ISO `YYYY-MM-DD`.
    #[serde(default)]
    pub date_to: Option<String>,
    /// "photo" | "video" — normalized in [`resolve_raw_query`].
    #[serde(default)]
    pub media_type: Option<String>,
}

/// Resolved structured query: tag names mapped to ids against the real
/// vocab, ISO dates converted to unix seconds. `place` stays as text for the
/// caller to forward-geocode into a gps circle. Serializable so the endpoint
/// can echo it back to the client as "this is how I read your query"
/// (editable filter chips).
#[derive(Debug, Clone, Default, PartialEq, Serialize)]
pub struct StructuredQuery {
    pub semantic: Option<String>,
    pub tag_ids: Vec<i32>,
    pub exclude_tag_ids: Vec<i32>,
    /// Tag names the model produced that don't exist in the vocabulary.
    /// The caller folds these back into the semantic term so the concept
    /// isn't lost — and surfacing them keeps a hallucinated tag from
    /// silently filtering the whole library to nothing.
    pub unmatched_tags: Vec<String>,
    pub camera_make: Option<String>,
    pub camera_model: Option<String>,
    pub lens_model: Option<String>,
    /// Raw place name awaiting forward-geocoding by the caller.
    pub place: Option<String>,
    pub date_from: Option<i64>,
    pub date_to: Option<i64>,
    /// Normalized to "photo" | "video"; `None` means no media-type filter.
    pub media_type: Option<String>,
}

/// Convert an ISO `YYYY-MM-DD` date to the true instant (unix seconds) that
/// day starts at in the owner's display `zone`. With `end_of_day`, returns
/// the last second of that day instead, so a `date_to` filter is inclusive
/// of the whole day. The photo filters compare true instants (ADR-0001,
/// #152), so "June" means the owner's June — the same days `/memories`
/// groups by. Returns `None` for any unparseable input (the filter is
/// simply omitted rather than erroring).
pub fn iso_to_unix(date: &str, end_of_day: bool, zone: chrono_tz::Tz) -> Option<i64> {
    let d = chrono::NaiveDate::parse_from_str(date.trim(), "%Y-%m-%d").ok()?;
    if end_of_day {
        Some(crate::date_view::day_start_instant(d.succ_opt()?, zone) - 1)
    } else {
        Some(crate::date_view::day_start_instant(d, zone))
    }
}

/// Normalize a free-form media-type string to the engine's vocabulary.
/// Anything that isn't clearly photo or video (including "all") yields
/// `None` — no filter.
fn normalize_media_type(raw: &str) -> Option<String> {
    match raw.trim().to_lowercase().as_str() {
        "photo" | "photos" | "image" | "images" | "picture" | "pictures" => {
            Some("photo".to_string())
        }
        "video" | "videos" | "movie" | "movies" | "clip" | "clips" => Some("video".to_string()),
        _ => None,
    }
}

/// Resolve a raw LLM query against the real tag vocabulary, producing the
/// structured filter. Pure — no network, no LLM — so it carries the
/// correctness-critical mapping logic under unit test.
///
/// `tag_vocab` is `(tag_id, tag_name)` pairs (the shape `TagDao::get_all_tags`
/// yields once the count is dropped). Matching is case-insensitive and exact
/// on the trimmed name.
pub fn resolve_raw_query(
    raw: RawNlQuery,
    tag_vocab: &[(i32, String)],
    display_zone: chrono_tz::Tz,
) -> StructuredQuery {
    // Case-insensitive name → id lookup. Built once per call.
    let lookup: std::collections::HashMap<String, i32> = tag_vocab
        .iter()
        .map(|(id, name)| (name.trim().to_lowercase(), *id))
        .collect();

    let resolve_names = |names: &[String], ids: &mut Vec<i32>, unmatched: &mut Vec<String>| {
        for name in names {
            let key = name.trim().to_lowercase();
            if key.is_empty() {
                continue;
            }
            match lookup.get(&key) {
                Some(id) if !ids.contains(id) => ids.push(*id),
                Some(_) => {} // duplicate, already collected
                None => {
                    if !unmatched.iter().any(|u| u.eq_ignore_ascii_case(name)) {
                        unmatched.push(name.trim().to_string());
                    }
                }
            }
        }
    };

    let mut tag_ids = Vec::new();
    let mut unmatched_tags = Vec::new();
    resolve_names(&raw.tags, &mut tag_ids, &mut unmatched_tags);

    // Excluded tags that don't match a real tag are simply ignored — you
    // can't exclude a tag that doesn't exist, and folding them into
    // `semantic` would make no sense.
    let mut exclude_tag_ids = Vec::new();
    let mut exclude_unmatched = Vec::new();
    resolve_names(
        &raw.exclude_tags,
        &mut exclude_tag_ids,
        &mut exclude_unmatched,
    );

    let clean = |s: Option<String>| s.map(|v| v.trim().to_string()).filter(|v| !v.is_empty());

    StructuredQuery {
        semantic: clean(raw.semantic),
        tag_ids,
        exclude_tag_ids,
        unmatched_tags,
        camera_make: clean(raw.camera_make),
        camera_model: clean(raw.camera_model),
        lens_model: clean(raw.lens_model),
        place: clean(raw.place),
        date_from: raw
            .date_from
            .as_deref()
            .and_then(|d| iso_to_unix(d, false, display_zone)),
        date_to: raw
            .date_to
            .as_deref()
            .and_then(|d| iso_to_unix(d, true, display_zone)),
        media_type: raw.media_type.as_deref().and_then(normalize_media_type),
    }
}

/// Build the grounded system prompt. The model is told the current date (so
/// "last summer" resolves) and the exact tag vocabulary (so it uses real
/// tags or routes the concept to `semantic` instead of inventing one).
fn build_system_prompt(tag_vocab: &[(i32, String)], today: chrono::NaiveDate) -> String {
    // Cap the vocab dump so a huge library doesn't blow the context window;
    // the most-used tags are the ones a query is likely to reference.
    const MAX_TAGS: usize = 400;
    let mut names: Vec<&str> = tag_vocab.iter().map(|(_, n)| n.as_str()).collect();
    names.sort_unstable();
    names.dedup();
    let shown = names.len().min(MAX_TAGS);
    let vocab = names[..shown].join(", ");
    let truncation = if names.len() > MAX_TAGS {
        format!(" (showing {MAX_TAGS} of {} tags)", names.len())
    } else {
        String::new()
    };

    format!(
        "You translate a user's natural-language photo-search request into a JSON \
filter. Today's date is {today}. Respond with ONLY a JSON object, no prose, no \
code fences.\n\n\
Schema (all fields optional):\n\
{{\n  \
\"semantic\": string|null,        // visual scene/subject for image similarity search\n  \
\"tags\": string[],               // ONLY names from the tag list below\n  \
\"exclude_tags\": string[],       // ONLY names from the tag list below\n  \
\"camera_make\": string|null,\n  \
\"camera_model\": string|null,\n  \
\"lens_model\": string|null,\n  \
\"place\": string|null,           // a location name to look up (city, country, landmark)\n  \
\"date_from\": \"YYYY-MM-DD\"|null,  // inclusive\n  \
\"date_to\": \"YYYY-MM-DD\"|null,    // inclusive\n  \
\"media_type\": \"photo\"|\"video\"|null\n\
}}\n\n\
Rules:\n\
- Put descriptive/visual concepts (\"sunset\", \"crowded beach\", \"red car\") in \"semantic\".\n\
- Only use \"tags\"/\"exclude_tags\" values that appear EXACTLY in the tag list. If a \
concept isn't a listed tag, put it in \"semantic\" instead — never invent a tag.\n\
- Resolve relative dates against today's date (\"last summer\", \"2023\", \"last month\").\n\
- Put place/location names in \"place\" (not \"semantic\").\n\
- Omit (use null / empty array) anything the request doesn't mention.\n\n\
Available tags{truncation}: {vocab}"
    )
}

/// Extract the JSON object from a model response that may include a leading
/// `<think>` block, code fences, or trailing prose. Strips the think block
/// first (so reasoning that mentions braces can't fool the scan), then
/// returns the substring from the first `{` to the last `}` inclusive — or
/// the trimmed text if no braces are found (which then fails to parse with a
/// clear error).
fn extract_json(raw: &str) -> String {
    let s = strip_think_blocks(raw);
    let start = s.find('{');
    let end = s.rfind('}');
    match (start, end) {
        (Some(a), Some(b)) if b >= a => s[a..=b].to_string(),
        _ => s.trim().to_string(),
    }
}

/// Parse a model response string into a [`StructuredQuery`], resolving names
/// against the vocab. Separated from the LLM call so it's unit-testable.
pub fn parse_response(
    response: &str,
    tag_vocab: &[(i32, String)],
    display_zone: chrono_tz::Tz,
) -> Result<StructuredQuery> {
    let json = extract_json(response);
    let raw: RawNlQuery = serde_json::from_str(&json)
        .map_err(|e| anyhow!("failed to parse NL query JSON: {e}; raw response: {response:?}"))?;
    Ok(resolve_raw_query(raw, tag_vocab, display_zone))
}

/// Translate a natural-language query into a [`StructuredQuery`] via one
/// grounded LLM call. The `client` is any configured backend (the unified
/// endpoint passes the resolved chat backend); `tag_vocab` grounds the tag
/// mapping; `today` (the owner's day, in `display_zone`) anchors
/// relative-date resolution, and `display_zone` turns the model's ISO
/// dates into day-bound instants.
pub async fn translate_nl_query(
    client: &dyn LlmClient,
    nl: &str,
    tag_vocab: &[(i32, String)],
    today: chrono::NaiveDate,
    display_zone: chrono_tz::Tz,
) -> Result<StructuredQuery> {
    let system = build_system_prompt(tag_vocab, today);
    let messages = vec![ChatMessage::system(system), ChatMessage::user(nl)];
    let (msg, _, _) = client.chat_with_tools(messages, Vec::<Tool>::new()).await?;
    parse_response(&msg.content, tag_vocab, display_zone)
}

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

    fn vocab() -> Vec<(i32, String)> {
        vec![
            (1, "beach".to_string()),
            (2, "Sunset".to_string()), // mixed case to exercise case-insensitivity
            (3, "family".to_string()),
        ]
    }

    #[test]
    fn iso_to_unix_start_and_end_of_day() {
        // 2023-01-01 UTC midnight = 1672531200.
        assert_eq!(
            iso_to_unix("2023-01-01", false, chrono_tz::UTC),
            Some(1_672_531_200)
        );
        // End of that day is 86399 seconds later.
        assert_eq!(
            iso_to_unix("2023-01-01", true, chrono_tz::UTC),
            Some(1_672_531_200 + 86_399)
        );
    }

    /// #152: ISO dates become the owner's day bounds in the display zone,
    /// not UTC days, so "June" lines up with the true-instant filters.
    #[test]
    fn iso_to_unix_uses_display_zone_day_bounds() {
        let ny = chrono_tz::America::New_York;
        // 2023-06-01 00:00 EDT = 04:00Z.
        assert_eq!(iso_to_unix("2023-06-01", false, ny), Some(1_685_592_000));
        // Last second of 2023-06-30 EDT = 2023-07-01 03:59:59Z.
        assert_eq!(iso_to_unix("2023-06-30", true, ny), Some(1_688_183_999));
    }

    #[test]
    fn iso_to_unix_rejects_garbage() {
        assert_eq!(iso_to_unix("last summer", false, chrono_tz::UTC), None);
        assert_eq!(iso_to_unix("2023-13-99", false, chrono_tz::UTC), None);
        assert_eq!(iso_to_unix("", false, chrono_tz::UTC), None);
    }

    #[test]
    fn resolve_matches_tags_case_insensitively() {
        let raw = RawNlQuery {
            tags: vec!["BEACH".to_string(), "sunset".to_string()],
            ..Default::default()
        };
        let q = resolve_raw_query(raw, &vocab(), chrono_tz::UTC);
        assert_eq!(q.tag_ids, vec![1, 2]);
        assert!(q.unmatched_tags.is_empty());
    }

    #[test]
    fn resolve_surfaces_unmatched_tags_not_silently_dropped() {
        // A hallucinated / non-vocab tag must be surfaced so the caller can
        // fold it into semantic — never silently used as a hard filter.
        let raw = RawNlQuery {
            tags: vec!["beach".to_string(), "golden hour".to_string()],
            ..Default::default()
        };
        let q = resolve_raw_query(raw, &vocab(), chrono_tz::UTC);
        assert_eq!(q.tag_ids, vec![1]);
        assert_eq!(q.unmatched_tags, vec!["golden hour".to_string()]);
    }

    #[test]
    fn resolve_dedups_repeated_tags() {
        let raw = RawNlQuery {
            tags: vec![
                "beach".to_string(),
                "Beach".to_string(),
                "beach".to_string(),
            ],
            ..Default::default()
        };
        let q = resolve_raw_query(raw, &vocab(), chrono_tz::UTC);
        assert_eq!(q.tag_ids, vec![1]);
    }

    #[test]
    fn resolve_normalizes_media_type_and_dates() {
        let raw = RawNlQuery {
            media_type: Some("Videos".to_string()),
            date_from: Some("2023-06-01".to_string()),
            date_to: Some("2023-06-30".to_string()),
            ..Default::default()
        };
        let q = resolve_raw_query(raw, &vocab(), chrono_tz::UTC);
        assert_eq!(q.media_type.as_deref(), Some("video"));
        assert_eq!(
            q.date_from,
            iso_to_unix("2023-06-01", false, chrono_tz::UTC)
        );
        assert_eq!(q.date_to, iso_to_unix("2023-06-30", true, chrono_tz::UTC));
    }

    #[test]
    fn resolve_media_type_all_is_no_filter() {
        let raw = RawNlQuery {
            media_type: Some("all".to_string()),
            ..Default::default()
        };
        assert_eq!(
            resolve_raw_query(raw, &vocab(), chrono_tz::UTC).media_type,
            None
        );
    }

    #[test]
    fn resolve_trims_and_empties_to_none() {
        let raw = RawNlQuery {
            semantic: Some("   ".to_string()),
            camera_make: Some("  Fujifilm  ".to_string()),
            place: Some("".to_string()),
            ..Default::default()
        };
        let q = resolve_raw_query(raw, &vocab(), chrono_tz::UTC);
        assert_eq!(q.semantic, None);
        assert_eq!(q.camera_make.as_deref(), Some("Fujifilm"));
        assert_eq!(q.place, None);
    }

    #[test]
    fn parse_response_handles_code_fences_and_prose() {
        let resp = "Here is the filter:\n```json\n{\"semantic\":\"sunset\",\"tags\":[\"beach\"]}\n```\nDone.";
        let q = parse_response(resp, &vocab(), chrono_tz::UTC).expect("parse");
        assert_eq!(q.semantic.as_deref(), Some("sunset"));
        assert_eq!(q.tag_ids, vec![1]);
    }

    #[test]
    fn parse_response_handles_think_block_then_json() {
        let resp = "<think>user wants beach sunsets</think>{\"tags\":[\"beach\",\"sunset\"]}";
        let q = parse_response(resp, &vocab(), chrono_tz::UTC).expect("parse");
        assert_eq!(q.tag_ids, vec![1, 2]);
    }

    #[test]
    fn parse_response_errors_on_non_json() {
        assert!(parse_response("I cannot help with that.", &vocab(), chrono_tz::UTC).is_err());
    }

    #[test]
    fn build_system_prompt_includes_date_and_vocab() {
        let today = chrono::NaiveDate::from_ymd_opt(2026, 6, 14).unwrap();
        let prompt = build_system_prompt(&vocab(), today);
        assert!(
            prompt.contains("2026-06-14"),
            "prompt should state today's date"
        );
        assert!(prompt.contains("beach"), "prompt should list the vocab");
        assert!(
            prompt.contains("never invent a tag"),
            "prompt should warn against inventing tags"
        );
    }
}