mx 0.1.227

use anyhow::{Context, Result};
use serde::Deserialize;
use surrealdb::sql::Thing;

use crate::types::Relationship;

use super::{ExistsRow, SurrealConnection, SurrealDatabase};

// =========================================================================
// RELATIONSHIP OPERATIONS
// =========================================================================

impl SurrealDatabase {
    /// Add a relationship between knowledge entries
    pub fn add_relationship(&self, from: &str, to: &str, rel_type: &str) -> Result<()> {
        Self::runtime().block_on(self.add_relationship_async(from, to, rel_type))
    }

    async fn add_relationship_async(&self, from: &str, to: &str, rel_type: &str) -> Result<()> {
        let from_id = from.strip_prefix("kn-").unwrap_or(from);
        let to_id = to.strip_prefix("kn-").unwrap_or(to);

        let from_thing = Thing::from(("knowledge", from_id));
        let to_thing = Thing::from(("knowledge", to_id));
        let rel_type_thing = Thing::from(("relationship_type", rel_type));

        with_db!(self, db, {
            db.query("RELATE $from->relates_to->$to SET relationship_type = $rel_type, created_at = time::now()")
                .bind(("from", from_thing))
                .bind(("to", to_thing))
                .bind(("rel_type", rel_type_thing))
                .await
                .context("Failed to create relationship")
        })?;

        Ok(())
    }

    /// List all relationships for a knowledge entry
    pub fn list_relationships(&self, entry_id: &str) -> Result<Vec<Relationship>> {
        Self::runtime().block_on(self.list_relationships_async(entry_id))
    }

    async fn list_relationships_async(&self, entry_id: &str) -> Result<Vec<Relationship>> {
        let id_part = entry_id.strip_prefix("kn-").unwrap_or(entry_id);
        let entry_thing = Thing::from(("knowledge", id_part));

        // Use meta::id() to extract plain string IDs from Thing record links.
        // Direct deserialization of Thing fields via serde_json::Value fails
        // because surrealdb::sql::Thing serializes as an untagged enum tuple
        // that serde_json cannot round-trip. meta::id() returns a plain string.
        #[derive(Debug, Deserialize)]
        struct RelRow {
            id: String,
            from_entry_id: String,
            to_entry_id: String,
            relationship_type: String,
            #[serde(default)]
            created_at: Option<String>,
        }

        let mut response = with_db!(self, db, {
            db.query(
                "SELECT meta::id(id) AS id,
                        meta::id(in) AS from_entry_id,
                        meta::id(out) AS to_entry_id,
                        meta::id(relationship_type) AS relationship_type,
                        <string>created_at AS created_at
                 FROM relates_to
                 WHERE in = $entry OR out = $entry
                 ORDER BY created_at DESC",
            )
            .bind(("entry", entry_thing))
            .await
            .context("Failed to query relationships")
        })?;

        let results: Vec<RelRow> = response.take(0)?;
        let relationships = results
            .into_iter()
            .map(|row| Relationship {
                id: row.id,
                from_entry_id: format!("kn-{}", row.from_entry_id),
                to_entry_id: format!("kn-{}", row.to_entry_id),
                relationship_type: row.relationship_type,
                created_at: row.created_at.unwrap_or_else(|| "unknown".to_string()),
            })
            .collect();

        Ok(relationships)
    }

    /// Delete a relationship by from/to/type triple
    pub fn delete_relationship(&self, from: &str, to: &str, rel_type: &str) -> Result<bool> {
        Self::runtime().block_on(self.delete_relationship_async(from, to, rel_type))
    }

    /// Delete a relationship edge by its record ID (e.g. "abc123" or the raw SurrealDB ID).
    pub fn delete_relationship_by_id(&self, id: &str) -> Result<bool> {
        Self::runtime().block_on(self.delete_relationship_by_id_async(id))
    }

    async fn delete_relationship_by_id_async(&self, id: &str) -> Result<bool> {
        // SurrealDB's RETURN BEFORE yields Thing-typed fields that serde_json cannot
        // round-trip. Instead: SELECT with meta::id() to check existence, then DELETE
        // without a RETURN clause (which is safe to take as Vec<Value> empty).
        //
        // NOTE: There is a TOCTOU window between the SELECT and DELETE — the edge
        // could be deleted by another caller between the two queries.  This is
        // acceptable for a single-user CLI tool where concurrent mutation is rare.
        let mut check = with_db!(self, db, {
            db.query("SELECT meta::id(id) AS id FROM relates_to WHERE meta::id(id) = $id LIMIT 1")
                .bind(("id", id.to_string()))
                .await
                .context("Failed to check relationship existence")
        })?;

        let exists: Vec<ExistsRow> = check.take(0)?;
        if exists.is_empty() {
            return Ok(false);
        }

        with_db!(self, db, {
            db.query("DELETE relates_to WHERE meta::id(id) = $id")
                .bind(("id", id.to_string()))
                .await
                .context("Failed to delete relationship by id")
        })?;

        Ok(true)
    }

    async fn delete_relationship_async(
        &self,
        from: &str,
        to: &str,
        rel_type: &str,
    ) -> Result<bool> {
        let from_id = from.strip_prefix("kn-").unwrap_or(from);
        let to_id = to.strip_prefix("kn-").unwrap_or(to);

        let from_thing = Thing::from(("knowledge", from_id));
        let to_thing = Thing::from(("knowledge", to_id));
        let rel_type_thing = Thing::from(("relationship_type", rel_type));

        // SurrealDB's RETURN BEFORE yields Thing-typed fields that serde_json cannot
        // round-trip. Check existence with meta::id() SELECT first, then DELETE without
        // a RETURN clause to avoid the deserialization error.
        let mut check = with_db!(self, db, {
            db.query(
                "SELECT meta::id(id) AS id FROM relates_to
                 WHERE in = $from AND out = $to AND relationship_type = $rel_type
                 LIMIT 1",
            )
            .bind(("from", from_thing.clone()))
            .bind(("to", to_thing.clone()))
            .bind(("rel_type", rel_type_thing.clone()))
            .await
            .context("Failed to check relationship existence")
        })?;

        let exists: Vec<ExistsRow> = check.take(0)?;
        if exists.is_empty() {
            return Ok(false);
        }

        with_db!(self, db, {
            db.query(
                "DELETE relates_to
                 WHERE in = $from AND out = $to AND relationship_type = $rel_type",
            )
            .bind(("from", from_thing))
            .bind(("to", to_thing))
            .bind(("rel_type", rel_type_thing))
            .await
            .context("Failed to delete relationship")
        })?;

        Ok(true)
    }

    /// Get facts extracted from a specific session
    pub fn get_facts_for_session(&self, session_id: &str) -> Result<Vec<String>> {
        Self::runtime().block_on(self.get_facts_for_session_async(session_id))
    }

    async fn get_facts_for_session_async(&self, session_id: &str) -> Result<Vec<String>> {
        let session_id_normalized = session_id.strip_prefix("kn-").unwrap_or(session_id);
        let session_thing = Thing::from(("knowledge", session_id_normalized));

        let mut response = with_db!(self, db, {
            db.query(
                "SELECT VALUE meta::id(in) FROM relates_to
                 WHERE out = $session_id AND relationship_type = relationship_type:extracted_from",
            )
            .bind(("session_id", session_thing))
            .await
            .context("Failed to query facts for session")
        })?;

        let fact_ids: Vec<String> = response.take(0).unwrap_or_default();
        let facts_with_prefix: Vec<String> = fact_ids
            .into_iter()
            .map(|id| format!("kn-{}", id))
            .collect();

        Ok(facts_with_prefix)
    }

    /// Get the session a fact was extracted from
    pub fn get_session_for_fact(&self, fact_id: &str) -> Result<Option<String>> {
        Self::runtime().block_on(self.get_session_for_fact_async(fact_id))
    }

    async fn get_session_for_fact_async(&self, fact_id: &str) -> Result<Option<String>> {
        let fact_id_normalized = fact_id.strip_prefix("kn-").unwrap_or(fact_id);
        let fact_thing = Thing::from(("knowledge", fact_id_normalized));

        let mut response = with_db!(self, db, {
            db.query(
                "SELECT VALUE meta::id(out) FROM relates_to
                 WHERE in = $fact AND relationship_type = relationship_type:extracted_from",
            )
            .bind(("fact", fact_thing))
            .await
            .context("Failed to query session for fact")
        })?;

        let session_ids: Vec<String> = response.take(0).unwrap_or_default();

        Ok(session_ids.first().map(|id| format!("kn-{}", id)))
    }

    // =========================================================================
    // TAG OPERATIONS (not exposed in public API, handled via knowledge entry)
    // =========================================================================

    /// Get tags for an entry
    pub fn get_tags_for_entry(&self, entry_id: &str) -> Result<Vec<String>> {
        Self::runtime().block_on(self.get_tags_for_entry_async(entry_id))
    }

    pub(super) async fn get_tags_for_entry_async(&self, entry_id: &str) -> Result<Vec<String>> {
        let id_part = entry_id.strip_prefix("kn-").unwrap_or(entry_id);
        let entry_thing = Thing::from(("knowledge", id_part));

        let mut tags_response = with_db!(self, db, {
            db.query("SELECT VALUE out.name FROM tagged_with WHERE in = $knowledge")
                .bind(("knowledge", entry_thing))
                .await
                .context("Failed to query tags")
        })?;

        let tags: Vec<String> = tags_response.take(0).unwrap_or_default();
        Ok(tags)
    }

    /// Set tags for an entry - handled automatically by upsert_knowledge
    pub fn set_tags_for_entry(&self, _entry_id: &str, _tags: &[String]) -> Result<()> {
        // Tags are managed via upsert_knowledge, this is a no-op for compatibility
        Ok(())
    }

    /// Batch-fetch tags for many entries in a single round trip.
    ///
    /// Review fix (N+1 hydration in embedding_chunk search, see knowledge.rs
    /// semantic_search): hydrating chunk candidates one at a time via
    /// `get_tags_for_entry_async` per entry means one DB round trip per
    /// unique chunk entry_id, which can run to the thousands under the
    /// archive-tier exclusion case. This collects tags for a whole window of
    /// entry_ids with one query instead.
    pub(super) async fn get_tags_for_entries_async(
        &self,
        entry_ids: &[String],
    ) -> Result<std::collections::HashMap<String, Vec<String>>> {
        let mut out: std::collections::HashMap<String, Vec<String>> =
            std::collections::HashMap::new();
        if entry_ids.is_empty() {
            return Ok(out);
        }

        let things: Vec<Thing> = entry_ids
            .iter()
            .map(|id| {
                let id_part = id.strip_prefix("kn-").unwrap_or(id);
                Thing::from(("knowledge", id_part))
            })
            .collect();

        #[derive(Deserialize)]
        struct TagRow {
            entry_id: String,
            #[serde(default)]
            tags: Vec<String>,
        }

        let mut response = with_db!(self, db, {
            db.query(
                // Both `tagged_with` and `applies_to` carry a UNIQUE composite
                // index on (in, out). A `union` lookup over the `in` prefix of
                // that index (what `in IN $knowledge` plans as) returns nothing,
                // even though the bind is correct and matches the right field --
                // `in = $one` on the same prefix works. Traversing from the
                // record ids directly sidesteps the index question: SurrealDB
                // resolves each entry by its key and walks the graph edge, so
                // the batch never touches that lookup path at all.
                "SELECT meta::id(id) AS entry_id, ->tagged_with->tag.name AS tags \
                 FROM $knowledge",
            )
            .bind(("knowledge", things))
            .await
            .context("Failed to batch-query tags")
        })?;

        let rows: Vec<TagRow> = response
            .take(0)
            .context("Failed to deserialize batch tag rows")?;
        for row in rows {
            out.entry(format!("kn-{}", row.entry_id))
                .or_default()
                .extend(row.tags);
        }

        Ok(out)
    }

    /// Get applicability for an entry
    pub fn get_applicability_for_entry(&self, entry_id: &str) -> Result<Vec<String>> {
        Self::runtime().block_on(self.get_applicability_for_entry_async(entry_id))
    }

    pub(super) async fn get_applicability_for_entry_async(
        &self,
        entry_id: &str,
    ) -> Result<Vec<String>> {
        let id_part = entry_id.strip_prefix("kn-").unwrap_or(entry_id);
        let entry_thing = Thing::from(("knowledge", id_part));

        let mut app_response = with_db!(self, db, {
            db.query("SELECT VALUE meta::id(out) FROM applies_to WHERE in = $knowledge")
                .bind(("knowledge", entry_thing))
                .await
                .context("Failed to query applicability")
        })?;

        let applicability: Vec<String> = app_response
            .take(0)
            .context("Failed to deserialize applicability")?;

        Ok(applicability)
    }

    /// Set applicability for an entry - handled automatically by upsert_knowledge
    pub fn set_applicability_for_entry(&self, _entry_id: &str, _ids: &[String]) -> Result<()> {
        // Applicability is managed via upsert_knowledge, this is a no-op for compatibility
        Ok(())
    }

    /// Batch-fetch applicability for many entries in a single round trip.
    /// See `get_tags_for_entries_async` for why this exists.
    pub(super) async fn get_applicability_for_entries_async(
        &self,
        entry_ids: &[String],
    ) -> Result<std::collections::HashMap<String, Vec<String>>> {
        let mut out: std::collections::HashMap<String, Vec<String>> =
            std::collections::HashMap::new();
        if entry_ids.is_empty() {
            return Ok(out);
        }

        let things: Vec<Thing> = entry_ids
            .iter()
            .map(|id| {
                let id_part = id.strip_prefix("kn-").unwrap_or(id);
                Thing::from(("knowledge", id_part))
            })
            .collect();

        #[derive(Deserialize)]
        struct ApplicabilityRow {
            entry_id: String,
            #[serde(default)]
            applies_ids: Vec<String>,
        }

        let mut response = with_db!(self, db, {
            db.query(
                // Same index shape as the tag query above; see the comment
                // there. `applicability_type` has no `name` field -- its
                // record id *is* its identity -- so the traversal target is
                // `.id`, and `meta::id()` needs to run per array element
                // rather than on the array as a whole (it does not
                // vectorize), hence the explicit `array::map`.
                "SELECT meta::id(id) AS entry_id, \
                        array::map(->applies_to->applicability_type.id, |$v| meta::id($v)) AS applies_ids \
                 FROM $knowledge",
            )
            .bind(("knowledge", things))
            .await
            .context("Failed to batch-query applicability")
        })?;

        let rows: Vec<ApplicabilityRow> = response
            .take(0)
            .context("Failed to deserialize batch applicability rows")?;
        for row in rows {
            out.entry(format!("kn-{}", row.entry_id))
                .or_default()
                .extend(row.applies_ids);
        }

        Ok(out)
    }
}