rigg-client 1.6.1

//! Microsoft Foundry REST API client
//!
//! Manages Foundry agents via the project-scoped `/agents` API
//! (v1 stable data plane at `https://{account}.services.ai.azure.com/api/projects/{project}`).

use std::time::Duration;

use reqwest::{Client, Method, StatusCode};
use serde_json::{Map, Value};
use tracing::{debug, instrument, warn};

use rigg_core::config::FoundryServiceConfig;

use crate::auth::{AuthProvider, get_auth_provider_for};
use crate::error::ClientError;

/// Maximum number of retry attempts for retryable errors
const MAX_RETRIES: u32 = 3;

/// Initial backoff delay in seconds
const INITIAL_BACKOFF_SECS: u64 = 1;

/// Microsoft Foundry API client
pub struct FoundryClient {
    http: Client,
    auth: Box<dyn AuthProvider>,
    base_url: String,
    project: String,
    api_version: String,
    /// Optional `Foundry-Features` header values gating preview features
    /// (e.g. `HostedAgents=V1Preview`).
    features: Vec<String>,
}

impl FoundryClient {
    /// Create a client from a workspace foundry connection (v1 data plane).
    pub fn from_connection(
        conn: &rigg_core::workspace::FoundryConnection,
    ) -> Result<Self, ClientError> {
        let auth = get_auth_provider_for(rigg_core::ServiceDomain::Foundry)?;
        let http = Client::builder().timeout(Duration::from_secs(30)).build()?;
        Ok(Self {
            http,
            auth,
            base_url: conn.url(),
            project: conn.project.clone(),
            api_version: conn
                .api_version
                .clone()
                .unwrap_or_else(|| rigg_core::registry::FOUNDRY_API_VERSION.to_string()),
            features: Vec::new(),
        })
    }

    /// Create a new Foundry client from service configuration (legacy).
    pub fn new(config: &FoundryServiceConfig) -> Result<Self, ClientError> {
        let auth = get_auth_provider_for(rigg_core::ServiceDomain::Foundry)?;
        let http = Client::builder().timeout(Duration::from_secs(30)).build()?;

        Ok(Self {
            http,
            auth,
            base_url: config.service_url(),
            project: config.project.clone(),
            api_version: rigg_core::registry::FOUNDRY_API_VERSION.to_string(),
            features: Vec::new(),
        })
    }

    /// Create with a custom auth provider (for testing)
    pub fn with_auth(
        base_url: String,
        project: String,
        api_version: String,
        auth: Box<dyn AuthProvider>,
    ) -> Result<Self, ClientError> {
        let http = Client::builder().timeout(Duration::from_secs(30)).build()?;

        Ok(Self {
            http,
            auth,
            base_url,
            project,
            api_version,
            features: Vec::new(),
        })
    }

    /// Enable preview feature gates sent via the `Foundry-Features` header.
    pub fn with_features(mut self, features: &[&str]) -> Self {
        self.features = features.iter().map(|s| s.to_string()).collect();
        self
    }

    /// Build URL for the agents collection
    fn agents_url(&self) -> String {
        format!(
            "{}/api/projects/{}/agents?api-version={}",
            self.base_url, self.project, self.api_version
        )
    }

    /// Build URL for a specific agent
    fn agent_url(&self, id: &str) -> String {
        format!(
            "{}/api/projects/{}/agents/{}?api-version={}",
            self.base_url,
            self.project,
            urlencoding::encode(id),
            self.api_version
        )
    }

    /// Build URL for creating/updating agent versions
    fn agent_versions_url(&self, name: &str) -> String {
        format!(
            "{}/api/projects/{}/agents/{}/versions?api-version={}",
            self.base_url,
            self.project,
            urlencoding::encode(name),
            self.api_version
        )
    }

    /// Invoke an agent with a single prompt via the Foundry responses API
    /// (`POST {project}/openai/v1/responses`); returns the raw OpenAI-shaped
    /// response object.
    pub async fn agent_respond(&self, agent: &str, input: &str) -> Result<Value, ClientError> {
        // NOTE: the /openai/v1/ route rejects an api-version query parameter
        // (the version is in the path) — verified live.
        let url = format!(
            "{}/api/projects/{}/openai/v1/responses",
            self.base_url, self.project
        );
        let body = serde_json::json!({
            "agent_reference": { "name": agent, "type": "agent_reference" },
            "input": input,
            "stream": false
        });
        self.request_with_retry(Method::POST, &url, Some(&body))
            .await?
            .ok_or_else(|| ClientError::Api {
                status: 0,
                message: "agent response had no body".to_string(),
            })
    }

    /// Execute an HTTP request
    async fn request(
        &self,
        method: Method,
        url: &str,
        body: Option<&Value>,
    ) -> Result<Option<Value>, ClientError> {
        let token = self.auth.get_token()?;

        let mut request = self
            .http
            .request(method.clone(), url)
            .header("Authorization", format!("Bearer {}", token))
            .header("Content-Type", "application/json");

        if !self.features.is_empty() {
            request = request.header("Foundry-Features", self.features.join(","));
        }

        if let Some(json) = body {
            request = request.json(json);
        }

        debug!("Request: {} {}", method, url);
        let response = request.send().await?;
        let status = response.status();

        if status == StatusCode::NO_CONTENT {
            return Ok(None);
        }

        let body = response.text().await?;

        if status.is_success() {
            if body.is_empty() {
                Ok(None)
            } else {
                let value: Value = serde_json::from_str(&body)?;
                Ok(Some(value))
            }
        } else {
            match status {
                StatusCode::NOT_FOUND => Err(ClientError::NotFound {
                    kind: "agent".to_string(),
                    name: url.to_string(),
                }),
                StatusCode::TOO_MANY_REQUESTS => {
                    let retry_after = 60;
                    Err(ClientError::RateLimited { retry_after })
                }
                StatusCode::SERVICE_UNAVAILABLE => Err(ClientError::ServiceUnavailable(body)),
                _ => Err(ClientError::from_response_with_url(
                    status.as_u16(),
                    &body,
                    Some(url),
                )),
            }
        }
    }

    /// Execute an HTTP request with retry logic
    async fn request_with_retry(
        &self,
        method: Method,
        url: &str,
        body: Option<&Value>,
    ) -> Result<Option<Value>, ClientError> {
        let mut attempt = 0u32;
        loop {
            match self.request(method.clone(), url, body).await {
                Ok(value) => return Ok(value),
                Err(err) if err.is_retryable() && attempt < MAX_RETRIES => {
                    let delay = match &err {
                        ClientError::RateLimited { retry_after } => {
                            Duration::from_secs(*retry_after)
                        }
                        _ => Duration::from_secs(INITIAL_BACKOFF_SECS * 2u64.pow(attempt)),
                    };
                    warn!(
                        "Request {} {} failed (attempt {}/{}): {}. Retrying in {:?}",
                        method,
                        url,
                        attempt + 1,
                        MAX_RETRIES + 1,
                        err,
                        delay,
                    );
                    tokio::time::sleep(delay).await;
                    attempt += 1;
                }
                Err(err) => return Err(err),
            }
        }
    }

    /// List all agents in the project
    #[instrument(skip(self))]
    pub async fn list_agents(&self) -> Result<Vec<Value>, ClientError> {
        let url = self.agents_url();
        let response = self.request_with_retry(Method::GET, &url, None).await?;

        match response {
            Some(value) => {
                let items = value
                    .get("data")
                    .and_then(|v| v.as_array())
                    .cloned()
                    .unwrap_or_default();
                // Flatten versioned response into flat agent objects
                Ok(items.iter().map(flatten_agent_response).collect())
            }
            None => Ok(Vec::new()),
        }
    }

    /// Get a specific agent by ID
    #[instrument(skip(self))]
    pub async fn get_agent(&self, id: &str) -> Result<Value, ClientError> {
        let url = self.agent_url(id);
        let response = self.request_with_retry(Method::GET, &url, None).await?;

        let raw = response.ok_or_else(|| ClientError::NotFound {
            kind: "Agent".to_string(),
            name: id.to_string(),
        })?;
        Ok(flatten_agent_response(&raw))
    }

    /// Create a new agent (creates first version)
    ///
    /// Takes a flat agent definition and wraps it in the API format
    /// before posting to `/agents/{name}/versions`.
    #[instrument(skip(self, definition))]
    pub async fn create_agent(&self, definition: &Value) -> Result<Value, ClientError> {
        let name = definition
            .get("name")
            .and_then(|n| n.as_str())
            .ok_or_else(|| ClientError::Api {
                status: 400,
                message: "Agent definition missing 'name' field".to_string(),
            })?;
        let payload = wrap_agent_payload(definition);
        let url = self.agent_versions_url(name);
        let response = self
            .request_with_retry(Method::POST, &url, Some(&payload))
            .await?;

        let raw = response.ok_or_else(|| ClientError::Api {
            status: 500,
            message: "No response body from agent creation".to_string(),
        })?;
        Ok(flatten_agent_response(&raw))
    }

    /// Update an existing agent (creates new version)
    ///
    /// Takes a flat agent definition and wraps it in the API format
    /// before posting to `/agents/{name}/versions`.
    #[instrument(skip(self, definition))]
    pub async fn update_agent(&self, id: &str, definition: &Value) -> Result<Value, ClientError> {
        let payload = wrap_agent_payload(definition);
        let url = self.agent_versions_url(id);
        let response = self
            .request_with_retry(Method::POST, &url, Some(&payload))
            .await?;

        let raw = response.ok_or_else(|| ClientError::Api {
            status: 500,
            message: "No response body from agent update".to_string(),
        })?;
        Ok(flatten_agent_response(&raw))
    }

    /// Delete an agent
    #[instrument(skip(self))]
    pub async fn delete_agent(&self, id: &str) -> Result<(), ClientError> {
        let url = self.agent_url(id);
        self.request_with_retry(Method::DELETE, &url, None).await?;
        Ok(())
    }

    /// Get the authentication method being used
    pub fn auth_method(&self) -> &'static str {
        self.auth.method_name()
    }
}

/// Wrap a flat agent definition into the API request format.
///
/// Converts from flat: `{ "name", "model", "instructions", "tools", ... }`
/// To API format:
/// ```json
/// {
///   "metadata": {...},
///   "description": "...",
///   "definition": {
///     "kind": "prompt",
///     "model": "...",
///     "instructions": "...",
///     "tools": [...]
///   }
/// }
/// ```
fn wrap_agent_payload(flat: &Value) -> Value {
    let obj = match flat.as_object() {
        Some(o) => o,
        None => return flat.clone(),
    };

    // Fields that go at the version level (outside definition)
    const VERSION_LEVEL_FIELDS: &[&str] = &["metadata", "description"];

    // Fields that are response-only and should not be sent
    const EXCLUDED_FIELDS: &[&str] = &["id", "name", "version", "created_at", "object"];

    let mut wrapper = Map::new();
    let mut definition = Map::new();

    for (key, value) in obj {
        if EXCLUDED_FIELDS.contains(&key.as_str()) {
            continue;
        } else if VERSION_LEVEL_FIELDS.contains(&key.as_str()) {
            wrapper.insert(key.clone(), value.clone());
        } else {
            definition.insert(key.clone(), value.clone());
        }
    }

    // Ensure kind is set (default to "prompt")
    if !definition.contains_key("kind") {
        definition.insert("kind".to_string(), Value::String("prompt".to_string()));
    }

    wrapper.insert("definition".to_string(), Value::Object(definition));
    Value::Object(wrapper)
}

/// Flatten a new Foundry agents API response into a flat structure
/// compatible with the agent decomposition pipeline.
///
/// The new Foundry API returns a versioned structure:
/// ```json
/// {
///   "object": "agent",
///   "id": "MyAgent",
///   "name": "MyAgent",
///   "versions": {
///     "latest": {
///       "metadata": {...},
///       "version": "5",
///       "definition": {
///         "kind": "prompt",
///         "model": "gpt-5.2-chat",
///         "instructions": "...",
///         "tools": [...]
///       }
///     }
///   }
/// }
/// ```
///
/// This flattens to: `{ "id", "name", "model", "instructions", "tools", ... }`
fn flatten_agent_response(agent: &Value) -> Value {
    let obj = match agent.as_object() {
        Some(o) => o,
        None => return agent.clone(),
    };

    let mut flat = Map::new();

    // Top-level fields
    if let Some(id) = obj.get("id") {
        flat.insert("id".to_string(), id.clone());
    }
    if let Some(name) = obj.get("name") {
        flat.insert("name".to_string(), name.clone());
    }

    // Version-create responses put the definition at the top level
    // (`{name, version, definition: {...}}`) instead of `versions.latest`.
    let top_level = obj
        .get("definition")
        .is_some()
        .then_some(obj)
        .map(|o| o as &Map<String, Value>);

    // Extract from versions.latest (GET/list shape) or the top level (POST shape)
    if let Some(latest) = obj
        .get("versions")
        .and_then(|v| v.get("latest"))
        .and_then(|l| l.as_object())
        .or(top_level)
    {
        // Version-level fields
        if let Some(metadata) = latest.get("metadata") {
            flat.insert("metadata".to_string(), metadata.clone());
        }
        if let Some(description) = latest.get("description") {
            flat.insert("description".to_string(), description.clone());
        }
        if let Some(version) = latest.get("version") {
            flat.insert("version".to_string(), version.clone());
        }
        if let Some(created_at) = latest.get("created_at") {
            flat.insert("created_at".to_string(), created_at.clone());
        }

        // Definition-level fields (model, instructions, tools, kind, etc.)
        if let Some(definition) = latest.get("definition").and_then(|d| d.as_object()) {
            for (key, value) in definition {
                flat.insert(key.clone(), value.clone());
            }
        }
    }

    // Ensure tools and tool_resources always present (API may omit when empty)
    flat.entry("tools".to_string())
        .or_insert_with(|| Value::Array(Vec::new()));
    flat.entry("tool_resources".to_string())
        .or_insert_with(|| Value::Object(Map::new()));

    Value::Object(flat)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::auth::{AuthError, AuthProvider};
    use serde_json::json;

    struct FakeAuth;
    impl AuthProvider for FakeAuth {
        fn get_token(&self) -> Result<String, AuthError> {
            Ok("fake-token".to_string())
        }
        fn method_name(&self) -> &'static str {
            "Fake"
        }
    }

    fn make_client() -> FoundryClient {
        FoundryClient::with_auth(
            "https://my-ai-svc.services.ai.azure.com".to_string(),
            "my-project".to_string(),
            "v1".to_string(),
            Box::new(FakeAuth),
        )
        .unwrap()
    }

    #[test]
    fn test_agents_url() {
        let client = make_client();
        let url = client.agents_url();
        assert_eq!(
            url,
            "https://my-ai-svc.services.ai.azure.com/api/projects/my-project/agents?api-version=v1"
        );
    }

    #[test]
    fn test_agent_url() {
        let client = make_client();
        let url = client.agent_url("Regulus");
        assert_eq!(
            url,
            "https://my-ai-svc.services.ai.azure.com/api/projects/my-project/agents/Regulus?api-version=v1"
        );
    }

    #[test]
    fn test_agent_versions_url() {
        let client = make_client();
        let url = client.agent_versions_url("KITT");
        assert_eq!(
            url,
            "https://my-ai-svc.services.ai.azure.com/api/projects/my-project/agents/KITT/versions?api-version=v1"
        );
    }

    #[test]
    fn test_auth_method() {
        let client = make_client();
        assert_eq!(client.auth_method(), "Fake");
    }

    #[test]
    fn flatten_handles_version_create_response_shape() {
        // POST /agents/{n}/versions returns the definition at the top level
        let response = json!({
            "object": "agent.version",
            "name": "helper",
            "version": "1",
            "definition": {
                "kind": "prompt",
                "model": "gpt-4.1-mini",
                "instructions": "Be nice.",
                "tools": [{"type": "mcp"}]
            }
        });
        let flat = flatten_agent_response(&response);
        assert_eq!(flat["name"], "helper");
        assert_eq!(flat["model"], "gpt-4.1-mini");
        assert_eq!(flat["instructions"], "Be nice.");
        assert_eq!(flat["tools"][0]["type"], "mcp");
    }

    #[test]
    fn test_wrap_agent_payload() {
        let flat = json!({
            "id": "KITT",
            "name": "KITT",
            "model": "gpt-5.2-chat",
            "kind": "prompt",
            "instructions": "You are KITT.",
            "tools": [{"type": "code_interpreter"}],
            "metadata": {"logo": "kitt.svg"},
            "description": "A smart car",
            "version": "3",
            "created_at": 1234567890
        });

        let wrapped = wrap_agent_payload(&flat);
        let obj = wrapped.as_object().unwrap();

        // Top level: metadata, description, definition
        assert!(obj.contains_key("definition"));
        assert!(obj.contains_key("metadata"));
        assert!(obj.contains_key("description"));

        // Excluded from payload
        assert!(!obj.contains_key("id"));
        assert!(!obj.contains_key("name"));
        assert!(!obj.contains_key("version"));
        assert!(!obj.contains_key("created_at"));

        // Definition should contain model, instructions, tools, kind
        let def = obj.get("definition").unwrap().as_object().unwrap();
        assert_eq!(def.get("model").unwrap(), "gpt-5.2-chat");
        assert_eq!(def.get("kind").unwrap(), "prompt");
        assert_eq!(def.get("instructions").unwrap(), "You are KITT.");
        assert!(def.get("tools").unwrap().as_array().unwrap().len() == 1);

        // Definition should NOT contain excluded or version-level fields
        assert!(!def.contains_key("id"));
        assert!(!def.contains_key("name"));
        assert!(!def.contains_key("metadata"));
    }

    #[test]
    fn test_wrap_agent_payload_adds_default_kind() {
        let flat = json!({
            "name": "simple",
            "model": "gpt-4o",
            "instructions": "Be helpful."
        });

        let wrapped = wrap_agent_payload(&flat);
        let def = wrapped.get("definition").unwrap().as_object().unwrap();
        assert_eq!(def.get("kind").unwrap(), "prompt");
    }

    #[test]
    fn test_flatten_then_wrap_roundtrip() {
        let api_response = json!({
            "object": "agent",
            "id": "KITT",
            "name": "KITT",
            "versions": {
                "latest": {
                    "metadata": {"logo": "kitt.svg"},
                    "version": "3",
                    "description": "Smart car",
                    "created_at": 1234567890,
                    "definition": {
                        "kind": "prompt",
                        "model": "gpt-5.2-chat",
                        "instructions": "You are KITT.",
                        "tools": [{"type": "code_interpreter"}]
                    }
                }
            }
        });

        let flat = flatten_agent_response(&api_response);
        let wrapped = wrap_agent_payload(&flat);

        // The wrapped payload should have a definition with the same content
        let def = wrapped.get("definition").unwrap().as_object().unwrap();
        assert_eq!(def.get("model").unwrap(), "gpt-5.2-chat");
        assert_eq!(def.get("instructions").unwrap(), "You are KITT.");
        assert_eq!(def.get("kind").unwrap(), "prompt");
    }

    #[test]
    fn test_flatten_agent_response_full() {
        let api_response = json!({
            "object": "agent",
            "id": "Regulus",
            "name": "Regulus",
            "versions": {
                "latest": {
                    "metadata": {
                        "logo": "Avatar_Default.svg",
                        "description": "",
                        "modified_at": "1769974547"
                    },
                    "object": "agent.version",
                    "id": "Regulus:5",
                    "name": "Regulus",
                    "version": "5",
                    "description": "",
                    "created_at": 1769974549,
                    "definition": {
                        "kind": "prompt",
                        "model": "gpt-5.2-chat",
                        "instructions": "You are Regulus.",
                        "tools": [
                            {"type": "mcp", "server_label": "kb_test"}
                        ]
                    }
                }
            }
        });

        let flat = flatten_agent_response(&api_response);
        let obj = flat.as_object().unwrap();

        assert_eq!(obj.get("id").unwrap(), "Regulus");
        assert_eq!(obj.get("name").unwrap(), "Regulus");
        assert_eq!(obj.get("model").unwrap(), "gpt-5.2-chat");
        assert_eq!(obj.get("kind").unwrap(), "prompt");
        assert_eq!(obj.get("instructions").unwrap(), "You are Regulus.");
        assert_eq!(obj.get("version").unwrap(), "5");
        assert_eq!(obj.get("description").unwrap(), "");
        assert!(obj.get("metadata").is_some());
        assert!(obj.get("tools").unwrap().as_array().unwrap().len() == 1);

        // Should NOT have the nested versions structure
        assert!(!obj.contains_key("versions"));
        assert!(!obj.contains_key("object"));
    }

    #[test]
    fn test_flatten_agent_response_minimal() {
        let api_response = json!({
            "object": "agent",
            "id": "simple",
            "name": "simple"
        });

        let flat = flatten_agent_response(&api_response);
        let obj = flat.as_object().unwrap();

        assert_eq!(obj.get("id").unwrap(), "simple");
        assert_eq!(obj.get("name").unwrap(), "simple");
        assert!(!obj.contains_key("model"));
        // tools and tool_resources always present with defaults
        assert_eq!(obj.get("tools").unwrap(), &json!([]));
        assert_eq!(obj.get("tool_resources").unwrap(), &json!({}));
    }

    #[test]
    fn test_flatten_agent_response_non_object() {
        let flat = flatten_agent_response(&json!("not an object"));
        assert_eq!(flat, json!("not an object"));
    }

    #[test]
    fn test_wrap_flatten_roundtrip_preserves_tool_permissions() {
        let flat = json!({
            "name": "test-agent",
            "kind": "prompt",
            "model": "gpt-4o",
            "tools": [{
                "type": "mcp",
                "server_label": "kb_test",
                "require_approval": "never",
                "allowed_tools": ["tool_a", "tool_b"]
            }]
        });

        // Wrap for API submission
        let wrapped = wrap_agent_payload(&flat);
        let def_tools = wrapped["definition"]["tools"].as_array().unwrap();
        assert_eq!(def_tools[0]["require_approval"], "never");
        assert_eq!(def_tools[0]["allowed_tools"][0], "tool_a");
        assert_eq!(def_tools[0]["allowed_tools"][1], "tool_b");

        // Simulate API response containing the same tools
        let api_response = json!({
            "object": "agent",
            "id": "test-agent",
            "name": "test-agent",
            "versions": {
                "latest": {
                    "version": "1",
                    "definition": {
                        "kind": "prompt",
                        "model": "gpt-4o",
                        "tools": [{
                            "type": "mcp",
                            "server_label": "kb_test",
                            "require_approval": "never",
                            "allowed_tools": ["tool_a", "tool_b"]
                        }]
                    }
                }
            }
        });

        let flattened = flatten_agent_response(&api_response);
        let tools = flattened["tools"].as_array().unwrap();
        assert_eq!(tools[0]["require_approval"], "never");
        let allowed = tools[0]["allowed_tools"].as_array().unwrap();
        assert_eq!(allowed.len(), 2);
        assert_eq!(allowed[0], "tool_a");
        assert_eq!(allowed[1], "tool_b");
    }

    #[test]
    fn test_flatten_preserves_require_approval_object_form() {
        let api_response = json!({
            "object": "agent",
            "id": "granular-agent",
            "name": "granular-agent",
            "versions": {
                "latest": {
                    "version": "1",
                    "definition": {
                        "kind": "prompt",
                        "model": "gpt-4o",
                        "tools": [{
                            "type": "mcp",
                            "server_label": "kb_test",
                            "require_approval": {
                                "never": {"tool_names": ["safe_tool"]},
                                "always": {"tool_names": ["dangerous_tool"]}
                            }
                        }]
                    }
                }
            }
        });

        let flat = flatten_agent_response(&api_response);
        let ra = &flat["tools"][0]["require_approval"];
        assert_eq!(ra["never"]["tool_names"][0], "safe_tool");
        assert_eq!(ra["always"]["tool_names"][0], "dangerous_tool");

        // Round-trip: flatten → wrap → verify definition still has structured form
        let re_wrapped = wrap_agent_payload(&flat);
        let def_ra = &re_wrapped["definition"]["tools"][0]["require_approval"];
        assert_eq!(def_ra["never"]["tool_names"][0], "safe_tool");
        assert_eq!(def_ra["always"]["tool_names"][0], "dangerous_tool");
    }
}