rigg 2.0.1

//! Shared remote access for sync commands: one façade over the three clients
//! (Search data plane, Foundry v1 data plane, ARM control plane), created
//! lazily per project + environment.

use anyhow::{Context, Result, bail};
use serde_json::Value;
use tokio::sync::OnceCell;

use colored::Colorize;
use rigg_client::arm_resources::ArmResourceClient;
use rigg_client::client::AzureSearchClient;
use rigg_client::error::ClientError;
use rigg_client::foundry::FoundryClient;
use rigg_core::registry::{self, Domain};
use rigg_core::resources::{ResourceKind, ResourceRef};
use rigg_core::workspace::{FoundryConnection, Project, ResolvedEnv, SearchConnection};

pub struct Remote {
    search_conn: Option<SearchConnection>,
    foundry_conn: Option<FoundryConnection>,
    search: OnceCell<AzureSearchClient>,
    foundry: OnceCell<FoundryClient>,
    arm: OnceCell<ArmResourceClient>,
}

impl Remote {
    /// Build for a project in an environment. Connections are optional; using
    /// a kind whose connection is missing yields a clear error.
    pub fn for_project(env: &ResolvedEnv, _project: &Project) -> Remote {
        Remote::for_env(env)
    }

    /// Build for an environment alone — the project has never mattered here
    /// (targets are environment-level), and callers without one (the auth
    /// engine, which reports across every project) should not have to invent
    /// a `Project` to reach the connections.
    pub fn for_env(env: &ResolvedEnv) -> Remote {
        Remote {
            search_conn: env.search().cloned(),
            foundry_conn: env.foundry().cloned(),
            search: OnceCell::new(),
            foundry: OnceCell::new(),
            arm: OnceCell::new(),
        }
    }

    /// Print the actual Azure targets (service names + resolved base URLs)
    /// so the user can verify where cloud operations go before anything
    /// happens — the URLs shown are exactly what the clients request against.
    pub fn print_targets(&self) {
        for line in self.target_lines() {
            println!("{line}");
        }
    }

    /// The same lines [`Self::print_targets`] prints, one per target, for
    /// callers that route output through `say!` instead of printing straight
    /// to stdout (so the line still reaches stderr in json mode rather than
    /// vanishing).
    pub fn target_lines(&self) -> Vec<String> {
        let mut lines = Vec::new();
        if let Some(s) = &self.search_conn {
            lines.push(format!("  Search:  {} → {}", s.service.bold(), s.url()));
        }
        if let Some(f) = &self.foundry_conn {
            lines.push(format!(
                "  Foundry: {}/{} → {}",
                f.account.bold(),
                f.project,
                f.url()
            ));
        }
        lines
    }

    pub fn has_search(&self) -> bool {
        self.search_conn.is_some()
    }

    pub fn has_foundry(&self) -> bool {
        self.foundry_conn.is_some()
    }

    /// Which kinds are reachable with the configured connections?
    pub fn supported_kinds(&self) -> Vec<ResourceKind> {
        ResourceKind::all()
            .iter()
            .copied()
            .filter(|k| match registry::meta(*k).domain {
                Domain::Search => self.has_search(),
                Domain::FoundryData | Domain::FoundryArm => self.has_foundry(),
            })
            .collect()
    }

    async fn search(&self) -> Result<&AzureSearchClient> {
        let conn = self
            .search_conn
            .as_ref()
            .context("no search connection configured for this project/environment")?;
        self.search
            .get_or_try_init(|| async { Ok(AzureSearchClient::from_connection(conn)?) })
            .await
    }

    async fn foundry(&self) -> Result<&FoundryClient> {
        let conn = self
            .foundry_conn
            .as_ref()
            .context("no foundry connection configured for this project/environment")?;
        self.foundry
            .get_or_try_init(|| async { Ok(FoundryClient::from_connection(conn)?) })
            .await
    }

    async fn arm(&self) -> Result<&ArmResourceClient> {
        let conn = self
            .foundry_conn
            .as_ref()
            .context("no foundry connection configured for this project/environment")?;
        self.arm
            .get_or_try_init(|| async {
                ArmResourceClient::for_account(&conn.account, &conn.project)
                    .await
                    .map_err(anyhow::Error::from)
            })
            .await
    }

    /// GET one resource; Ok(None) when it does not exist remotely.
    pub async fn get(&self, r: &ResourceRef) -> Result<Option<Value>> {
        match registry::meta(r.kind).domain {
            Domain::Search => match self.search().await?.get(r.kind, &r.name).await {
                Ok(v) => Ok(Some(v)),
                Err(ClientError::NotFound { .. }) => Ok(None),
                Err(e) => Err(e.into()),
            },
            Domain::FoundryData => match self.foundry().await?.get_agent(&r.name).await {
                Ok(v) => Ok(Some(v)),
                Err(ClientError::NotFound { .. }) => Ok(None),
                Err(e) => Err(e.into()),
            },
            Domain::FoundryArm => Ok(self.arm().await?.get(r.kind, &r.name).await?),
        }
    }

    /// List all resources of a kind. Every returned item carries a "name".
    pub async fn list(&self, kind: ResourceKind) -> Result<Vec<Value>> {
        let items = match registry::meta(kind).domain {
            Domain::Search => self.search().await?.list(kind).await?,
            Domain::FoundryData => self.foundry().await?.list_agents().await?,
            Domain::FoundryArm => self.arm().await?.list(kind).await?,
        };
        Ok(items)
    }

    /// Create or update; returns the server's post-write document
    /// (GETs it back when the API returns 204/no body) for canonicalization.
    ///
    /// A timed-out PUT is ambiguous — the server may have completed it (an
    /// indexer create validates connections and starts its first run before
    /// responding). Resolve the ambiguity instead of failing: GET the
    /// resource and accept the write when the server's document semantically
    /// matches what was sent.
    pub async fn put(&self, r: &ResourceRef, body: &Value) -> Result<Value> {
        match self.put_inner(r, body).await {
            Ok(v) => Ok(v),
            Err(e) if is_timeout(&e) => match self.get(r).await {
                Ok(Some(server_doc))
                    if rigg_core::normalize::semantic_eq(r.kind, body, &server_doc) =>
                {
                    Ok(server_doc)
                }
                _ => Err(e.context(format!(
                    "the request timed out and {r} does not (yet) match what was sent — \
                         re-run the command; it resumes safely"
                ))),
            },
            Err(e) => Err(e),
        }
    }

    async fn put_inner(&self, r: &ResourceRef, body: &Value) -> Result<Value> {
        match registry::meta(r.kind).domain {
            Domain::Search => {
                let client = self.search().await?;
                match client.create_or_update(r.kind, &r.name, body).await? {
                    Some(v) => Ok(v),
                    None => client.get(r.kind, &r.name).await.map_err(Into::into),
                }
            }
            Domain::FoundryData => {
                let client = self.foundry().await?;
                let exists = client.get_agent(&r.name).await.is_ok();
                let result = if exists {
                    client.update_agent(&r.name, body).await?
                } else {
                    client.create_agent(body).await?
                };
                Ok(result)
            }
            Domain::FoundryArm => Ok(self.arm().await?.put(r.kind, &r.name, body).await?),
        }
    }

    // ------------------------------------------------------------------
    // Runtime operations (rigg az)
    // ------------------------------------------------------------------

    pub async fn indexer_run(&self, name: &str) -> Result<()> {
        Ok(self.search().await?.indexer_run(name).await?)
    }

    pub async fn indexer_reset(&self, name: &str) -> Result<()> {
        Ok(self.search().await?.indexer_reset(name).await?)
    }

    pub async fn indexer_status(&self, name: &str) -> Result<Value> {
        Ok(self.search().await?.indexer_status(name).await?)
    }

    pub async fn index_stats(&self, name: &str) -> Result<Value> {
        Ok(self.search().await?.index_stats(name).await?)
    }

    pub async fn search_docs(&self, index: &str, body: &Value) -> Result<Value> {
        Ok(self.search().await?.search_docs(index, body).await?)
    }

    pub async fn kb_retrieve(&self, kb: &str, body: &Value) -> Result<Value> {
        Ok(self.search().await?.kb_retrieve(kb, body).await?)
    }

    pub async fn agent_ask(&self, agent: &str, input: &str) -> Result<Value> {
        Ok(self.foundry().await?.agent_respond(agent, input).await?)
    }

    /// Delete a remote resource (idempotent: missing is not an error).
    pub async fn delete(&self, r: &ResourceRef) -> Result<()> {
        let result = match registry::meta(r.kind).domain {
            Domain::Search => self.search().await?.delete(r.kind, &r.name).await,
            Domain::FoundryData => self.foundry().await?.delete_agent(&r.name).await,
            Domain::FoundryArm => return Ok(self.arm().await?.delete(r.kind, &r.name).await?),
        };
        match result {
            Ok(()) => Ok(()),
            Err(ClientError::NotFound { .. }) => Ok(()),
            Err(e) => Err(e.into()),
        }
    }

    /// Fetch every remote resource of the supported kinds as (ref, doc).
    pub async fn snapshot(&self) -> Result<Vec<(ResourceRef, Value)>> {
        let mut out = Vec::new();
        for kind in self.supported_kinds() {
            let items = self
                .list(kind)
                .await
                .with_context(|| format!("failed to list remote {}", kind.directory_name()))?;
            for item in items {
                let Some(name) = item.get("name").and_then(Value::as_str) else {
                    continue;
                };
                if rigg_core::resources::validate_resource_name(name).is_err() {
                    continue;
                }
                out.push((ResourceRef::new(kind, name.to_string()), item));
            }
        }
        Ok(out)
    }
}

/// Whether an error chain bottoms out in an HTTP timeout.
fn is_timeout(e: &anyhow::Error) -> bool {
    e.chain().any(|cause| {
        cause
            .downcast_ref::<reqwest::Error>()
            .is_some_and(reqwest::Error::is_timeout)
            || cause.to_string().contains("operation timed out")
    })
}

/// Resolve environment-specific values into a push body: `x-rigg-ref`
/// annotations that point at knowledge bases inject the KB's MCP endpoint
/// into a sibling `server_url`/`url` field when empty.
pub fn resolve_cross_service_refs(
    env_search: Option<&SearchConnection>,
    body: &mut Value,
) -> Result<()> {
    let Some(search) = env_search else {
        return Ok(());
    };
    resolve_walk(&search.service, body)
}

/// The knowledge base's MCP endpoint as Foundry expects it (documented form;
/// answers synthesized on the preview api-version).
pub fn kb_mcp_url(search_service: &str, kb: &str) -> String {
    format!(
        "https://{search_service}.search.windows.net/knowledgebases/{kb}/mcp?api-version={}",
        registry::SEARCH_PREVIEW_API_VERSION
    )
}

fn resolve_walk(search_service: &str, value: &mut Value) -> Result<()> {
    match value {
        Value::Object(map) => {
            let kb_ref = map
                .get(registry::X_RIGG_REF)
                .and_then(Value::as_str)
                .and_then(|s| s.split_once('/'))
                .filter(|(dir, _)| *dir == "knowledge-bases")
                .map(|(_, name)| name.to_string());
            if let Some(kb) = kb_ref {
                // The knowledge base exposes an MCP endpoint for agent
                // grounding. The x-rigg-ref annotation is authoritative:
                // the URL is (re)computed for the target environment on
                // every push, so one file set promotes across environments.
                let mcp_url = kb_mcp_url(search_service, &kb);
                let field = ["server_url", "url", "endpoint"]
                    .into_iter()
                    .find(|f| map.contains_key(*f))
                    .unwrap_or("server_url");
                map.insert(field.to_string(), Value::String(mcp_url));
                // Foundry rejects an MCP tool without a `server_label`, and
                // the label is not environment-specific — derive it from the
                // same annotation rather than making every agent file carry
                // a hand-written copy of its knowledge base's name.
                if !map
                    .get("server_label")
                    .and_then(Value::as_str)
                    .is_some_and(|s| !s.trim().is_empty())
                {
                    map.insert(
                        "server_label".to_string(),
                        Value::String(kb.replace(['-', '.'], "_")),
                    );
                }
            }
            for (_, v) in map.iter_mut() {
                resolve_walk(search_service, v)?;
            }
        }
        Value::Array(arr) => {
            for item in arr {
                resolve_walk(search_service, item)?;
            }
        }
        _ => {}
    }
    Ok(())
}

/// Bail with a helpful error when a project has no usable connections.
pub fn ensure_any_connection(remote: &Remote, project: &Project) -> Result<()> {
    if !remote.has_search() && !remote.has_foundry() {
        bail!(
            "project '{}' has no reachable services in this environment (configure `search:` or `foundry:` in rigg.yaml)",
            project.name
        );
    }
    Ok(())
}

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

    #[test]
    fn kb_mcp_url_uses_the_documented_form_and_preview_version() {
        assert_eq!(
            kb_mcp_url("mklabsrch", "regulatory-kb"),
            format!(
                "https://mklabsrch.search.windows.net/knowledgebases/regulatory-kb/mcp?api-version={}",
                rigg_core::registry::SEARCH_PREVIEW_API_VERSION
            )
        );
    }

    /// Foundry rejects an MCP tool with no `server_label`, so a tool that
    /// only carries `x-rigg-ref` gets one derived from the same annotation
    /// (Foundry labels take `[A-Za-z0-9_]`). A label the author wrote is
    /// left alone.
    #[test]
    fn an_mcp_tool_gets_a_server_label_from_its_x_rigg_ref() {
        let search = SearchConnection {
            service: "svc".to_string(),
            ..Default::default()
        };
        let mut body = serde_json::json!({
            "name": "agent",
            "tools": [
                {"type": "mcp", "x-rigg-ref": "knowledge-bases/docs-kb", "server_url": ""},
                {"type": "mcp", "x-rigg-ref": "knowledge-bases/docs-kb", "server_label": "mine"}
            ]
        });
        resolve_cross_service_refs(Some(&search), &mut body).unwrap();
        assert_eq!(body["tools"][0]["server_label"], "docs_kb");
        assert_eq!(body["tools"][1]["server_label"], "mine");
        assert_eq!(body["tools"][0]["server_url"], kb_mcp_url("svc", "docs-kb"));
    }
}