mod query_builder;

use std::collections::{HashMap, HashSet};
use std::path::Path;
use std::sync::{Arc, Mutex};

use common::proto::komp_ac::search::searcher_server::Searcher;
pub use common::proto::komp_ac::search::searcher_server::SearcherServer;
use common::proto::komp_ac::search::{
    SearchOrderDirection, SearchRequest, SearchResponse, search_response::Hit,
};
use common::search::{SchemaFields, register_tokenizers, search_index_path};
use query_builder::{
    ConstraintMode, SearchConstraint, SearchConstraintTarget, build_master_query,
};
use sqlx::{AssertSqlSafe, PgPool, Row};
use tantivy::collector::TopDocs;
use tantivy::schema::Value;
use tantivy::{Index, IndexReader, ReloadPolicy, TantivyDocument};
use tonic::{Request, Response, Status};
use tracing::info;

const INDEX_ROOT: &str = "./tantivy_indexes";
const DEFAULT_RESULT_LIMIT: usize = 60;
const HARD_RESULT_LIMIT: usize = 200;
const DEFAULT_LIST_LIMIT: usize = 5;
const SEARCH_SCORE_RELATIVE_FLOOR: f32 = 0.25;
const SEARCH_SCORE_GROUP_WIDTH: f32 = 0.25;

pub struct SearcherService {
    pub pool: PgPool,
    profiles: Mutex<HashMap<String, Arc<ProfileIndex>>>,
}

impl SearcherService {
    pub fn new(pool: PgPool) -> Self {
        Self {
            pool,
            profiles: Mutex::new(HashMap::new()),
        }
    }

    async fn run_rpc(
        &self,
        request: Request<SearchRequest>,
    ) -> Result<Response<SearchResponse>, Status> {
        let req = request.into_inner();
        let normalized = normalize_request(req)?;

        if !profile_exists(&self.pool, &normalized.profile_name).await? {
            return Err(Status::not_found(format!(
                "Profile '{}' was not found",
                normalized.profile_name
            )));
        }

        if let Some(table_name) = normalized.table_name.as_deref() {
            if !table_exists(&self.pool, &normalized.profile_name, table_name).await? {
                return Err(Status::not_found(format!(
                    "Table '{}' was not found in profile '{}'",
                    table_name, normalized.profile_name
                )));
            }
        }

        if !normalized.has_input() {
            let Some(table_name) = normalized.table_name.as_deref() else {
                return Err(Status::invalid_argument(
                    "table_name is required when query is empty",
                ));
            };

            let hits = fetch_ordered_rows(
                &self.pool,
                &normalized.profile_name,
                table_name,
                normalized.order.as_ref(),
                normalized.limit.unwrap_or(DEFAULT_LIST_LIMIT),
                normalized.offset,
            )
            .await?;
            return Ok(Response::new(SearchResponse { hits }));
        }

        let index_path = search_index_path(Path::new(INDEX_ROOT), &normalized.profile_name);
        if !index_path.exists() {
            return Err(Status::not_found(format!(
                "No search index found for profile '{}'",
                normalized.profile_name
            )));
        }

        let resolved_must = resolve_constraints(
            &self.pool,
            &normalized.profile_name,
            normalized.table_name.as_deref(),
            &normalized.must,
        )
        .await?;

        let profile = profile_index(&self.profiles, &normalized.profile_name, &index_path)?;
        let mut hits = run_search(
            &self.pool,
            &profile,
            &normalized.profile_name,
            normalized.table_name.as_deref(),
            &normalized.free_query,
            &resolved_must,
            normalized.limit.unwrap_or(DEFAULT_RESULT_LIMIT),
            normalized.offset,
            normalized.order.as_ref(),
        )
        .await?;

        if normalized.order.is_none() {
            hits.sort_by(|left, right| right.score.total_cmp(&left.score));
            if let Some(limit) = normalized.limit {
                if hits.len() > limit {
                    hits.truncate(limit);
                }
            }
        }

        info!(
            "search: profile={} table={:?} free='{}' constraints={} hits={}",
            normalized.profile_name,
            normalized.table_name,
            normalized.free_query,
            resolved_must.len(),
            hits.len()
        );

        Ok(Response::new(SearchResponse { hits }))
    }
}

struct ProfileIndex {
    index: Index,
    reader: IndexReader,
    fields: SchemaFields,
}

impl ProfileIndex {
    fn open(path: &Path) -> Result<Self, Status> {
        let index = Index::open_in_dir(path)
            .map_err(|e| Status::internal(format!("Failed to open index: {}", e)))?;
        register_tokenizers(&index)
            .map_err(|e| Status::internal(format!("Failed to register tokenizers: {}", e)))?;
        let reader = index
            .reader_builder()
            .reload_policy(ReloadPolicy::OnCommitWithDelay)
            .try_into()
            .map_err(|e| Status::internal(format!("Failed to build index reader: {}", e)))?;
        let fields = SchemaFields::from(&index.schema()).map_err(|e| {
            Status::internal(format!(
                "Search index schema mismatch. Delete the stale index and create it again: {}",
                e
            ))
        })?;

        Ok(Self {
            index,
            reader,
            fields,
        })
    }
}

#[derive(Debug)]
struct NormalizedSearchRequest {
    profile_name: String,
    table_name: Option<String>,
    free_query: String,
    must: Vec<NormalizedColumnConstraint>,
    limit: Option<usize>,
    offset: usize,
    order: Option<NormalizedSearchOrder>,
}

#[derive(Debug)]
struct NormalizedSearchOrder {
    column: String,
    direction: SearchOrderDirection,
}

#[derive(Debug)]
struct NormalizedColumnConstraint {
    column: String,
    query: String,
    mode: ConstraintMode,
}

impl NormalizedSearchRequest {
    fn has_input(&self) -> bool {
        !self.free_query.is_empty() || !self.must.is_empty()
    }
}

fn profile_index(
    cache: &Mutex<HashMap<String, Arc<ProfileIndex>>>,
    profile_name: &str,
    path: &Path,
) -> Result<Arc<ProfileIndex>, Status> {
    {
        let cache_guard = cache
            .lock()
            .map_err(|_| Status::internal("Profile index cache lock poisoned"))?;
        if let Some(index) = cache_guard.get(profile_name) {
            return Ok(index.clone());
        }
    }

    let opened = Arc::new(ProfileIndex::open(path)?);
    let mut cache_guard = cache
        .lock()
        .map_err(|_| Status::internal("Profile index cache lock poisoned"))?;
    if let Some(index) = cache_guard.get(profile_name) {
        return Ok(index.clone());
    }
    cache_guard.insert(profile_name.to_string(), opened.clone());
    Ok(opened)
}

fn validate_identifier(value: &str, field_name: &str) -> Result<(), Status> {
    let mut chars = value.chars();
    let Some(first) = chars.next() else {
        return Err(Status::invalid_argument(format!(
            "{field_name} must not be empty"
        )));
    };

    if !(first.is_ascii_alphabetic() || first == '_')
        || !chars.all(|ch| ch.is_ascii_alphanumeric() || ch == '_')
    {
        return Err(Status::invalid_argument(format!(
            "{field_name} contains invalid characters"
        )));
    }

    Ok(())
}

fn validate_search_column(value: &str) -> Result<(), Status> {
    if value.is_empty() {
        return Err(Status::invalid_argument(
            "constraint.column must not be empty",
        ));
    }

    if value.chars().any(|ch| ch.is_control() || ch == '\0') {
        return Err(Status::invalid_argument(
            "constraint.column contains invalid characters",
        ));
    }

    Ok(())
}

fn qualify_profile_table(profile_name: &str, table_name: &str) -> String {
    format!("\"{}\".\"{}\"", profile_name, table_name)
}

async fn profile_exists(pool: &PgPool, profile_name: &str) -> Result<bool, Status> {
    let exists =
        sqlx::query_scalar::<_, bool>("SELECT EXISTS(SELECT 1 FROM schemas WHERE name = $1)")
            .bind(profile_name)
            .fetch_one(pool)
            .await
            .map_err(|e| Status::internal(format!("Profile lookup failed: {}", e)))?;
    Ok(exists)
}

async fn table_exists(pool: &PgPool, profile_name: &str, table_name: &str) -> Result<bool, Status> {
    let exists = sqlx::query_scalar::<_, bool>(
        r#"
        SELECT EXISTS(
            SELECT 1
            FROM table_definitions td
            JOIN schemas s ON td.schema_id = s.id
            WHERE s.name = $1 AND td.table_name = $2
        )
        "#,
    )
    .bind(profile_name)
    .bind(table_name)
    .fetch_one(pool)
    .await
    .map_err(|e| Status::internal(format!("Table lookup failed: {}", e)))?;
    Ok(exists)
}

fn normalize_request(req: SearchRequest) -> Result<NormalizedSearchRequest, Status> {
    let profile_name = req.profile_name.trim();
    if profile_name.is_empty() {
        return Err(Status::invalid_argument("profile_name is required"));
    }
    validate_identifier(profile_name, "profile_name")?;

    let table_name = match req.table_name.as_deref().map(str::trim) {
        Some(table_name) if !table_name.is_empty() => {
            validate_identifier(table_name, "table_name")?;
            Some(table_name.to_string())
        }
        _ => None,
    };

    let free_query = req.free_query.trim().to_string();
    let mut must = Vec::new();

    for constraint in req.must {
        let column = constraint.column.trim();
        validate_search_column(column)?;

        let query = constraint.query.trim();
        if query.is_empty() {
            return Err(Status::invalid_argument(
                "constraint.query must not be empty",
            ));
        }

        must.push(NormalizedColumnConstraint {
            column: column.to_string(),
            query: query.to_string(),
            mode: constraint_mode_from_proto(constraint.mode),
        });
    }

    let limit = req
        .limit
        .map(|value| (value as usize).min(HARD_RESULT_LIMIT));
    let offset = req.offset.unwrap_or_default() as usize;
    let order = req
        .order
        .map(|order| {
            let column = order.column.trim();
            if column.is_empty() {
                return Err(Status::invalid_argument("order.column must not be empty"));
            }
            validate_search_column(column)?;
            let direction = SearchOrderDirection::try_from(order.direction).map_err(|_| {
                Status::invalid_argument("order.direction must be ASC or DESC")
            })?;
            if direction == SearchOrderDirection::Unspecified {
                return Err(Status::invalid_argument(
                    "order.direction must be ASC or DESC",
                ));
            }
            Ok(NormalizedSearchOrder {
                column: column.to_string(),
                direction,
            })
        })
        .transpose()?;

    if order.is_some() && table_name.is_none() {
        return Err(Status::invalid_argument(
            "table_name is required when order is specified",
        ));
    }

    Ok(NormalizedSearchRequest {
        profile_name: profile_name.to_string(),
        table_name,
        free_query,
        must,
        limit,
        offset,
        order,
    })
}

fn constraint_mode_from_proto(raw_mode: i32) -> ConstraintMode {
    match raw_mode {
        2 => ConstraintMode::Exact,
        _ => ConstraintMode::Fuzzy,
    }
}

async fn resolve_constraints(
    pool: &PgPool,
    profile_name: &str,
    table_filter: Option<&str>,
    must: &[NormalizedColumnConstraint],
) -> Result<Vec<SearchConstraint>, Status> {
    let mut resolved = Vec::with_capacity(must.len());
    for constraint in must {
        let targets = resolve_constraint_targets(pool, profile_name, table_filter, &constraint.column)
            .await?;

        resolved.push(SearchConstraint {
            targets,
            query: constraint.query.clone(),
            mode: constraint.mode,
        });
    }
    Ok(resolved)
}

async fn resolve_constraint_targets(
    pool: &PgPool,
    profile_name: &str,
    table_filter: Option<&str>,
    column: &str,
) -> Result<Vec<SearchConstraintTarget>, Status> {
    let rows = if let Some(table_name) = table_filter {
        sqlx::query(
            r#"
            SELECT td.table_name, tdc.physical_name
            FROM schemas s
            JOIN table_definitions td ON td.schema_id = s.id
            JOIN table_definition_columns tdc ON tdc.table_definition_id = td.id
            WHERE s.name = $1
              AND td.table_name = $2
              AND tdc.display_name = $3
            "#,
        )
        .bind(profile_name)
        .bind(table_name)
        .bind(column)
        .fetch_all(pool)
        .await
    } else {
        sqlx::query(
            r#"
            SELECT td.table_name, tdc.physical_name
            FROM schemas s
            JOIN table_definitions td ON td.schema_id = s.id
            JOIN table_definition_columns tdc ON tdc.table_definition_id = td.id
            WHERE s.name = $1
              AND tdc.display_name = $2
            "#,
        )
        .bind(profile_name)
        .bind(column)
        .fetch_all(pool)
        .await
    }
    .map_err(|e| Status::internal(format!("Column mapping lookup failed: {}", e)))?;

    if rows.is_empty() {
        let target = SearchConstraintTarget {
            table_name: table_filter.map(str::to_string),
            column: column.to_string(),
        };
        return Ok(vec![target]);
    }

    let mut seen = HashSet::new();
    let mut targets = Vec::new();
    for row in rows {
        let table_name: String = row
            .try_get("table_name")
            .map_err(|e| Status::internal(format!("Column mapping table read failed: {}", e)))?;
        let physical_name: String = row
            .try_get("physical_name")
            .map_err(|e| Status::internal(format!("Column mapping physical read failed: {}", e)))?;
        let key = (table_name.clone(), physical_name.clone());
        if seen.insert(key) {
            targets.push(SearchConstraintTarget {
                table_name: if table_filter.is_some() {
                    None
                } else {
                    Some(table_name)
                },
                column: physical_name,
            });
        }
    }

    Ok(targets)
}

async fn table_physical_to_display_map(
    pool: &PgPool,
    profile_name: &str,
    table_name: &str,
) -> Result<HashMap<String, String>, Status> {
    let rows = sqlx::query(
        r#"
        SELECT tdc.physical_name, tdc.display_name
        FROM schemas s
        JOIN table_definitions td ON td.schema_id = s.id
        JOIN table_definition_columns tdc ON tdc.table_definition_id = td.id
        WHERE s.name = $1
          AND td.table_name = $2
        "#,
    )
    .bind(profile_name)
    .bind(table_name)
    .fetch_all(pool)
    .await
    .map_err(|e| Status::internal(format!("Column mapping lookup failed: {}", e)))?;

    let mut mapping = HashMap::with_capacity(rows.len());
    for row in rows {
        let physical_name: String = row
            .try_get("physical_name")
            .map_err(|e| Status::internal(format!("Column mapping physical read failed: {}", e)))?;
        let display_name: String = row
            .try_get("display_name")
            .map_err(|e| Status::internal(format!("Column mapping display read failed: {}", e)))?;
        mapping.insert(physical_name, display_name);
    }
    Ok(mapping)
}

async fn table_row_display_column(
    pool: &PgPool,
    profile_name: &str,
    table_name: &str,
) -> Result<String, Status> {
    sqlx::query_scalar(
        r#"
        SELECT td.row_display_column
        FROM schemas s
        JOIN table_definitions td ON td.schema_id = s.id
        WHERE s.name = $1 AND td.table_name = $2
        "#,
    )
    .bind(profile_name)
    .bind(table_name)
    .fetch_one(pool)
    .await
    .map_err(|e| Status::internal(format!("Row display column lookup failed: {}", e)))
}

fn remap_json_to_display_names(
    value: serde_json::Value,
    physical_to_display: &HashMap<String, String>,
) -> serde_json::Value {
    match value {
        serde_json::Value::Object(object) => {
            let mut remapped = serde_json::Map::with_capacity(object.len());
            for (key, value) in object {
                let final_key = physical_to_display.get(&key).cloned().unwrap_or(key);
                remapped.insert(final_key, value);
            }
            serde_json::Value::Object(remapped)
        }
        other => other,
    }
}

fn row_display_value(value: &serde_json::Value, column: &str) -> String {
    match value.get(column) {
        Some(serde_json::Value::String(value)) => value.clone(),
        Some(serde_json::Value::Number(value)) => value.to_string(),
        Some(serde_json::Value::Bool(value)) => value.to_string(),
        _ => String::new(),
    }
}

enum ResolvedOrderColumn {
    Position,
    Column(String),
}

fn order_direction_sql(direction: SearchOrderDirection) -> &'static str {
    match direction {
        SearchOrderDirection::Asc => "ASC",
        SearchOrderDirection::Desc => "DESC",
        SearchOrderDirection::Unspecified => unreachable!("order direction is normalized"),
    }
}

async fn resolve_order_column(
    pool: &PgPool,
    profile_name: &str,
    table_name: &str,
    requested_column: &str,
) -> Result<ResolvedOrderColumn, Status> {
    if requested_column.eq_ignore_ascii_case("position") {
        return Ok(ResolvedOrderColumn::Position);
    }

    let requested_column = if requested_column.eq_ignore_ascii_case("row_display_column") {
        table_row_display_column(pool, profile_name, table_name).await?
    } else {
        requested_column.to_string()
    };
    let physical_to_display =
        table_physical_to_display_map(pool, profile_name, table_name).await?;
    let physical_column = physical_to_display
        .iter()
        .find(|(physical, display)| {
            physical.eq_ignore_ascii_case(&requested_column)
                || display.eq_ignore_ascii_case(&requested_column)
        })
        .map(|(physical, _)| physical.clone())
        .unwrap_or_else(|| requested_column.clone());

    let physical_column = sqlx::query_scalar::<_, String>(
        r#"
        SELECT column_name
        FROM information_schema.columns
        WHERE table_schema = $1 AND table_name = $2 AND LOWER(column_name) = LOWER($3)
        "#,
    )
    .bind(profile_name)
    .bind(table_name)
    .bind(&physical_column)
    .fetch_optional(pool)
    .await
    .map_err(|e| Status::internal(format!("Order column lookup failed: {}", e)))?;
    let Some(physical_column) = physical_column else {
        return Err(Status::invalid_argument(format!(
            "Order column '{}' was not found in table '{}.{}'",
            requested_column, profile_name, table_name
        )));
    };

    Ok(ResolvedOrderColumn::Column(physical_column))
}

fn order_clause(column: &ResolvedOrderColumn, direction: SearchOrderDirection) -> String {
    let direction = order_direction_sql(direction);
    match column {
        ResolvedOrderColumn::Position => {
            format!("picker_position {}", direction)
        }
        ResolvedOrderColumn::Column(column) => {
            let quoted_column = format!("\"{}\"", column.replace('"', "\"\""));
            format!("{} {} NULLS LAST, id ASC", quoted_column, direction)
        }
    }
}

fn ranked_order_clause(
    column: &ResolvedOrderColumn,
    direction: SearchOrderDirection,
) -> String {
    let direction = order_direction_sql(direction);
    match column {
        ResolvedOrderColumn::Position => {
            format!("candidate.candidate_group DESC, positioned.picker_position {}", direction)
        }
        ResolvedOrderColumn::Column(column) => {
            let quoted_column = format!("\"{}\"", column.replace('"', "\"\""));
            format!(
                "candidate.candidate_group DESC, positioned.{} {} NULLS LAST, positioned.id ASC",
                quoted_column, direction
            )
        }
    }
}

fn relevance_group(score: f32, best_score: f32) -> i32 {
    if best_score <= 0.0 || !best_score.is_finite() || !score.is_finite() {
        return 0;
    }

    ((score.max(0.0) / best_score / SEARCH_SCORE_GROUP_WIDTH).floor() as i32).clamp(0, 3)
}

async fn fetch_ordered_rows(
    pool: &PgPool,
    profile_name: &str,
    table_name: &str,
    order: Option<&NormalizedSearchOrder>,
    limit: usize,
    offset: usize,
) -> Result<Vec<Hit>, Status> {
    let physical_to_display = table_physical_to_display_map(pool, profile_name, table_name).await?;
    let display_column = table_row_display_column(pool, profile_name, table_name).await?;
    let (resolved_order, direction) = match order {
        Some(order) => (
            resolve_order_column(pool, profile_name, table_name, &order.column).await?,
            order.direction,
        ),
        None => (ResolvedOrderColumn::Position, SearchOrderDirection::Desc),
    };
    let sql = format!(
        "WITH positioned AS (\
            SELECT t.*, ROW_NUMBER() OVER (ORDER BY id ASC) AS picker_position \
            FROM {} t WHERE deleted = FALSE\
        ) \
        SELECT id, to_jsonb(positioned) - 'picker_position' AS data, picker_position \
        FROM positioned ORDER BY {} LIMIT $1 OFFSET $2",
        qualify_profile_table(profile_name, table_name),
        order_clause(&resolved_order, direction),
    );

    let rows = sqlx::query(AssertSqlSafe(sql))
        .bind(limit as i64)
        .bind(offset as i64)
        .fetch_all(pool)
        .await
        .map_err(|e| Status::internal(format!("DB query for default results failed: {}", e)))?;

    Ok(rows
        .into_iter()
        .map(|row| {
            let id: i64 = row.try_get("id").unwrap_or_default();
            let json_data: serde_json::Value = row.try_get("data").unwrap_or_default();
            let position: i64 = row.try_get("picker_position").unwrap_or_default();
            let json_data = remap_json_to_display_names(json_data, &physical_to_display);
            let row_display_value = row_display_value(&json_data, &display_column);
            Hit {
                id,
                score: 0.0,
                content_json: json_data.to_string(),
                table_name: table_name.to_string(),
                row_display_value,
                row_display_column: display_column.clone(),
                position: u64::try_from(position).ok(),
            }
        })
        .collect())
}

async fn run_search(
    pool: &PgPool,
    profile: &ProfileIndex,
    profile_name: &str,
    table_filter: Option<&str>,
    free_query: &str,
    must: &[SearchConstraint],
    limit: usize,
    offset: usize,
    order: Option<&NormalizedSearchOrder>,
) -> Result<Vec<Hit>, Status> {
    let master_query = build_master_query(
        &profile.index,
        &profile.fields,
        free_query,
        must,
        table_filter,
    )?;

    let searcher = profile.reader.searcher();
    let window_limit = if order.is_some() {
        searcher.num_docs() as usize
    } else {
        offset.saturating_add(limit)
    };
    if window_limit == 0 {
        return Ok(Vec::new());
    }
    let top_docs = searcher
        .search(&*master_query, &TopDocs::with_limit(window_limit).order_by_score())
        .map_err(|e| Status::internal(format!("Search failed: {}", e)))?;

    if top_docs.is_empty() {
        return Ok(vec![]);
    }

    let best_score = top_docs.first().map(|(score, _)| *score).unwrap_or_default();
    let score_floor = if best_score > 0.0 {
        best_score * SEARCH_SCORE_RELATIVE_FLOOR
    } else {
        0.0
    };
    let eligible_docs = top_docs
        .into_iter()
        .filter(|(score, _)| *score >= score_floor)
        .collect::<Vec<_>>();
    let page_docs = if order.is_some() {
        eligible_docs
    } else {
        eligible_docs
            .into_iter()
            .skip(offset)
            .take(limit)
            .collect::<Vec<_>>()
    };

    let mut candidates: Vec<(f32, i64, String)> = Vec::with_capacity(page_docs.len());
    for (score, doc_address) in page_docs {
        let doc: TantivyDocument = searcher
            .doc(doc_address)
            .map_err(|e| Status::internal(format!("Failed to retrieve document: {}", e)))?;
        let Some(pg_id) = doc
            .get_first(profile.fields.pg_id)
            .and_then(|value| value.as_u64())
        else {
            continue;
        };
        let Some(table_name) = doc
            .get_first(profile.fields.table_name)
            .and_then(|value| value.as_str())
        else {
            continue;
        };
        candidates.push((score, pg_id as i64, table_name.to_string()));
    }

    if candidates.is_empty() {
        return Ok(vec![]);
    }

    if let Some(order) = order {
        let table_name = table_filter.expect("ordered searches require a normalized table filter");
        return fetch_ordered_candidate_rows(
            pool,
            profile_name,
            table_name,
            &candidates,
            order,
            limit,
            offset,
        )
        .await;
    }

    let mut ids_by_table: HashMap<String, Vec<i64>> = HashMap::new();
    for (_, pg_id, table_name) in &candidates {
        ids_by_table
            .entry(table_name.clone())
            .or_default()
            .push(*pg_id);
    }

    let mut content_map: HashMap<(String, i64), (String, String, String)> = HashMap::new();
    for (table_name, pg_ids) in ids_by_table {
        validate_identifier(&table_name, "table_name")?;
        let physical_to_display =
            table_physical_to_display_map(pool, profile_name, &table_name).await?;
        let display_column = table_row_display_column(pool, profile_name, &table_name).await?;
        let sql = format!(
            "SELECT id, to_jsonb(t) AS data FROM {} t WHERE deleted = FALSE AND id = ANY($1)",
            qualify_profile_table(profile_name, &table_name)
        );

        let rows = sqlx::query(AssertSqlSafe(sql))
            .bind(&pg_ids)
            .fetch_all(pool)
            .await
            .map_err(|e| Status::internal(format!("Database query failed: {}", e)))?;

        for row in rows {
            let id: i64 = row.try_get("id").unwrap_or_default();
            let json_data: serde_json::Value = row.try_get("data").unwrap_or_default();
            let json_data = remap_json_to_display_names(json_data, &physical_to_display);
            let display_value = row_display_value(&json_data, &display_column);
            content_map.insert(
                (table_name.clone(), id),
                (json_data.to_string(), display_value, display_column.clone()),
            );
        }
    }

    Ok(candidates
        .into_iter()
        .filter_map(|(score, pg_id, table_name)| {
            content_map
                .get(&(table_name.clone(), pg_id))
                .map(|(content_json, row_display_value, row_display_column)| Hit {
                    id: pg_id,
                    score,
                    content_json: content_json.clone(),
                    table_name,
                    row_display_value: row_display_value.clone(),
                    row_display_column: row_display_column.clone(),
                    position: None,
                })
        })
        .collect())
}

async fn fetch_ordered_candidate_rows(
    pool: &PgPool,
    profile_name: &str,
    table_name: &str,
    candidates: &[(f32, i64, String)],
    order: &NormalizedSearchOrder,
    limit: usize,
    offset: usize,
) -> Result<Vec<Hit>, Status> {
    let physical_to_display = table_physical_to_display_map(pool, profile_name, table_name).await?;
    let display_column = table_row_display_column(pool, profile_name, table_name).await?;
    let resolved_order =
        resolve_order_column(pool, profile_name, table_name, &order.column).await?;
    let sql = format!(
        "WITH positioned AS (\
            SELECT t.*, ROW_NUMBER() OVER (ORDER BY id ASC) AS picker_position \
            FROM {} t WHERE deleted = FALSE\
        ) \
        SELECT positioned.id, to_jsonb(positioned) - 'picker_position' AS data, \
            picker_position, candidate_score \
        FROM positioned \
        JOIN UNNEST($1::BIGINT[], $2::REAL[], $3::INTEGER[]) \
            AS candidate(candidate_id, candidate_score, candidate_group) \
            ON candidate_id = positioned.id \
        ORDER BY {} LIMIT $4 OFFSET $5",
        qualify_profile_table(profile_name, table_name),
        ranked_order_clause(&resolved_order, order.direction),
    );
    let ids = candidates
        .iter()
        .map(|(_, id, _)| *id)
        .collect::<Vec<_>>();
    let scores = candidates
        .iter()
        .map(|(score, _, _)| *score)
        .collect::<Vec<_>>();
    let best_score = scores
        .iter()
        .copied()
        .max_by(f32::total_cmp)
        .unwrap_or_default();
    let groups = scores
        .iter()
        .map(|score| relevance_group(*score, best_score))
        .collect::<Vec<_>>();
    let rows = sqlx::query(AssertSqlSafe(sql))
        .bind(&ids)
        .bind(&scores)
        .bind(&groups)
        .bind(limit as i64)
        .bind(offset as i64)
        .fetch_all(pool)
        .await
        .map_err(|e| Status::internal(format!("Ordered search query failed: {}", e)))?;

    Ok(rows
        .into_iter()
        .map(|row| {
            let id: i64 = row.try_get("id").unwrap_or_default();
            let json_data: serde_json::Value = row.try_get("data").unwrap_or_default();
            let position: i64 = row.try_get("picker_position").unwrap_or_default();
            let score: f32 = row.try_get("candidate_score").unwrap_or_default();
            let json_data = remap_json_to_display_names(json_data, &physical_to_display);
            let display_value = row_display_value(&json_data, &display_column);
            Hit {
                id,
                score,
                content_json: json_data.to_string(),
                table_name: table_name.to_string(),
                row_display_value: display_value,
                row_display_column: display_column.clone(),
                position: u64::try_from(position).ok(),
            }
        })
        .collect())
}

#[tonic::async_trait]
impl Searcher for SearcherService {
    async fn search(
        &self,
        request: Request<SearchRequest>,
    ) -> Result<Response<SearchResponse>, Status> {
        self.run_rpc(request).await
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use common::proto::komp_ac::search::{SearchOrder, SearchOrderDirection};

    fn request(order: Option<SearchOrder>, table_name: Option<&str>) -> SearchRequest {
        SearchRequest {
            profile_name: "finance".to_string(),
            table_name: table_name.map(str::to_string),
            free_query: String::new(),
            must: Vec::new(),
            limit: Some(20),
            offset: Some(40),
            order,
        }
    }

    #[test]
    fn normalizes_typed_column_order() {
        let normalized = normalize_request(request(
            Some(SearchOrder {
                column: "created_at".to_string(),
                direction: SearchOrderDirection::Desc as i32,
            }),
            Some("invoice"),
        ))
        .unwrap();
        let order = normalized.order.unwrap();

        assert_eq!(order.column, "created_at");
        assert_eq!(order.direction, SearchOrderDirection::Desc);
    }

    #[test]
    fn ordered_request_requires_a_table() {
        let result = normalize_request(request(
            Some(SearchOrder {
                column: "position".to_string(),
                direction: SearchOrderDirection::Desc as i32,
            }),
            None,
        ));

        assert!(result.is_err());
    }

    #[test]
    fn column_order_has_stable_id_tie_breaker() {
        assert_eq!(
            order_clause(
                &ResolvedOrderColumn::Column("1".to_string()),
                SearchOrderDirection::Asc,
            ),
            "\"1\" ASC NULLS LAST, id ASC"
        );
    }

    #[test]
    fn ranked_order_preserves_relevance_before_column_order() {
        assert_eq!(
            ranked_order_clause(
                &ResolvedOrderColumn::Column("created_at".to_string()),
                SearchOrderDirection::Desc,
            ),
            "candidate.candidate_group DESC, positioned.\"created_at\" DESC NULLS LAST, positioned.id ASC"
        );
    }

    #[test]
    fn relevance_groups_nearby_high_scores_and_separates_low_scores() {
        let best_score = 8.7;

        assert_eq!(relevance_group(8.7, best_score), 3);
        assert_eq!(relevance_group(8.6, best_score), 3);
        assert_eq!(relevance_group(8.4, best_score), 3);
        assert_eq!(relevance_group(7.0, best_score), 3);
        assert_eq!(relevance_group(3.0, best_score), 1);
    }
}