anakin-sdk 0.1.0

use std::time::Duration;

use serde::de::DeserializeOwned;
use serde_json::{json, Value};

use crate::error::{Error, Result};
use crate::sessions::Sessions;
use crate::types::{
    AgenticSearchResult, CrawlResult, Document, MapResult, SearchResult, WireResult,
};

/// SDK version baked into the User-Agent header.
pub const VERSION: &str = "0.1.0";

const DEFAULT_BASE_URL: &str = "https://api.anakin.io/v1";

/// Asynchronous client for the Anakin web-scraping API.
///
/// ```no_run
/// use anakin::Client;
///
/// # async fn run() -> anakin::Result<()> {
/// let client = Client::builder().api_key("ak-...").build()?;
/// let doc = client.scrape("https://example.com").await?;
/// println!("{}", doc.markdown.unwrap_or_default());
/// # Ok(())
/// # }
/// ```
///
/// The client is `Clone` and `Send + Sync` — share a single instance across
/// tasks rather than building a new one per request.
#[derive(Clone)]
pub struct Client {
    http: reqwest::Client,
    api_key: String,
    base_url: String,
    max_retries: u32,
    poll_interval: Duration,
    poll_max_interval: Duration,
    poll_timeout: Duration,
}

impl std::fmt::Debug for Client {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Client")
            .field("base_url", &self.base_url)
            .field("max_retries", &self.max_retries)
            .field("poll_interval", &self.poll_interval)
            .field("poll_max_interval", &self.poll_max_interval)
            .field("poll_timeout", &self.poll_timeout)
            .finish_non_exhaustive()
    }
}

impl Client {
    /// Start building a client. See [`ClientBuilder`].
    pub fn builder() -> ClientBuilder {
        ClientBuilder::default()
    }

    /// Browser-session management endpoints.
    pub fn sessions(&self) -> Sessions<'_> {
        Sessions::new(self)
    }

    // ── public endpoints ─────────────────────────────────────────────────

    /// Scrape a single URL. Polls internally until the job reaches a terminal status.
    pub async fn scrape(&self, url: &str) -> Result<Document> {
        self.scrape_with(url, None).await
    }

    /// Scrape with extra options (passed as a JSON object).
    pub async fn scrape_with(&self, url: &str, opts: Option<Value>) -> Result<Document> {
        let body = build_body(json!({ "url": url }), opts);
        let submit: Value = self.send_json(reqwest::Method::POST, "/url-scraper", Some(body)).await?;
        let job_id = require_string(&submit, "job_id")?;
        let poll = self.poll_job(&format!("/url-scraper/{job_id}")).await?;
        decode_field(&poll, "result")
    }

    /// Discover links for a domain. Polls until the job completes.
    pub async fn map(&self, url: &str) -> Result<MapResult> {
        self.map_with(url, None).await
    }

    pub async fn map_with(&self, url: &str, opts: Option<Value>) -> Result<MapResult> {
        let body = build_body(json!({ "url": url }), opts);
        let submit: Value = self.send_json(reqwest::Method::POST, "/map", Some(body)).await?;
        let job_id = require_string(&submit, "job_id")?;
        let poll = self.poll_job(&format!("/map/{job_id}")).await?;
        decode_field(&poll, "result")
    }

    /// Crawl a site starting from a URL. Polls until the job completes.
    pub async fn crawl(&self, url: &str) -> Result<CrawlResult> {
        self.crawl_with(url, None).await
    }

    pub async fn crawl_with(&self, url: &str, opts: Option<Value>) -> Result<CrawlResult> {
        let body = build_body(json!({ "url": url }), opts);
        let submit: Value = self.send_json(reqwest::Method::POST, "/crawl", Some(body)).await?;
        let job_id = require_string(&submit, "job_id")?;
        let poll = self.poll_job(&format!("/crawl/{job_id}")).await?;
        decode_field(&poll, "result")
    }

    /// Synchronous web search.
    pub async fn search(&self, query: &str) -> Result<SearchResult> {
        self.search_with(query, None).await
    }

    pub async fn search_with(&self, query: &str, opts: Option<Value>) -> Result<SearchResult> {
        let body = build_body(json!({ "prompt": query }), opts);
        let v: Value = self.send_json(reqwest::Method::POST, "/search", Some(body)).await?;
        serde_json::from_value(v).map_err(|e| Error::Other(format!("decode response: {e}")))
    }

    /// AI-synthesised agentic search. Polls until the job completes.
    pub async fn agentic_search(&self, prompt: &str) -> Result<AgenticSearchResult> {
        self.agentic_search_with(prompt, None).await
    }

    pub async fn agentic_search_with(&self, prompt: &str, opts: Option<Value>) -> Result<AgenticSearchResult> {
        let body = build_body(json!({ "prompt": prompt }), opts);
        let submit: Value = self.send_json(reqwest::Method::POST, "/agentic-search", Some(body)).await?;
        let job_id = require_string(&submit, "job_id")?;
        let poll = self.poll_job(&format!("/agentic-search/{job_id}")).await?;
        decode_field(&poll, "result")
    }

    /// Execute a Wire (Holocron) action by ID. Polls until the job completes.
    pub async fn wire(&self, action_id: &str, params: Option<Value>) -> Result<WireResult> {
        let mut body = json!({ "action_id": action_id });
        if let Some(p) = params {
            body["params"] = p;
        }
        let submit: Value = self.send_json(reqwest::Method::POST, "/holocron/task", Some(body)).await?;
        let job_id = require_string(&submit, "job_id")?;
        let poll = self.poll_job(&format!("/holocron/task/{job_id}")).await?;
        decode_field(&poll, "result")
    }

    // ── HTTP plumbing (visible to sub-modules) ───────────────────────────

    pub(crate) async fn send_json(
        &self,
        method: reqwest::Method,
        path: &str,
        body: Option<Value>,
    ) -> Result<Value> {
        let url = format!("{}{}", self.base_url, path);
        let mut last_resp: Option<reqwest::Response> = None;

        for attempt in 0..=self.max_retries {
            if attempt > 0 {
                let delay = backoff(attempt, last_resp.as_ref());
                tokio::time::sleep(delay).await;
            }

            let mut req = self
                .http
                .request(method.clone(), &url)
                .header("X-API-Key", &self.api_key)
                .header("Accept", "application/json")
                .header("User-Agent", format!("anakin-rust/{VERSION}"));
            if let Some(ref b) = body {
                req = req.json(b);
            }

            let resp = match req.send().await {
                Ok(r) => r,
                Err(e) => {
                    if attempt == self.max_retries {
                        return Err(Error::Network {
                            message: format!(
                                "http request after {} retries: {}",
                                self.max_retries, e
                            ),
                            source: Some(Box::new(e)),
                        });
                    }
                    continue;
                }
            };

            let status = resp.status().as_u16();
            if should_retry(status) && attempt < self.max_retries {
                last_resp = Some(resp);
                continue;
            }
            return self.handle_response(resp).await;
        }

        // Unreachable: the loop body always returns or continues, and the
        // final iteration returns via handle_response or the err branch.
        Err(Error::Other("retry loop exited unexpectedly".into()))
    }

    async fn handle_response(&self, resp: reqwest::Response) -> Result<Value> {
        let status = resp.status();
        let retry_after_header = resp
            .headers()
            .get("Retry-After")
            .and_then(|v| v.to_str().ok())
            .map(|s| s.to_string());

        if status.is_success() {
            let bytes = resp.bytes().await.map_err(|e| Error::Network {
                message: format!("read response: {e}"),
                source: Some(Box::new(e)),
            })?;
            if bytes.is_empty() {
                return Ok(Value::Object(serde_json::Map::new()));
            }
            return serde_json::from_slice(&bytes)
                .map_err(|e| Error::Other(format!("decode response: {e}")));
        }

        let body_bytes = resp.bytes().await.unwrap_or_default();
        Err(map_error(status.as_u16(), retry_after_header.as_deref(), &body_bytes))
    }

    async fn poll_job(&self, path: &str) -> Result<Value> {
        let deadline = std::time::Instant::now() + self.poll_timeout;
        let mut delay = self.poll_interval;

        loop {
            let v = self.send_json(reqwest::Method::GET, path, None).await?;
            let status = v.get("status").and_then(|s| s.as_str()).unwrap_or("");
            let error = v.get("error").and_then(|s| s.as_str()).unwrap_or("");
            let job_id = v.get("job_id").and_then(|s| s.as_str()).map(|s| s.to_string());

            if status == "completed" || status == "succeeded" {
                return Ok(v);
            }
            if status == "failed" {
                return Err(Error::JobFailed {
                    job_id,
                    reason: error.to_string(),
                });
            }

            if std::time::Instant::now() > deadline {
                return Err(Error::JobTimeout {
                    job_id,
                    elapsed: self.poll_timeout,
                });
            }

            tokio::time::sleep(delay).await;
            let next_ms = (delay.as_millis() as f64 * 1.5) as u64;
            let capped = std::cmp::min(next_ms, self.poll_max_interval.as_millis() as u64);
            delay = Duration::from_millis(capped);
        }
    }
}

// ── helpers ──────────────────────────────────────────────────────────────

fn build_body(base: Value, extra: Option<Value>) -> Value {
    let mut obj = base.as_object().cloned().unwrap_or_default();
    if let Some(Value::Object(map)) = extra {
        for (k, v) in map {
            obj.insert(k, v);
        }
    }
    Value::Object(obj)
}

fn should_retry(status: u16) -> bool {
    status == 429 || (500..600).contains(&status)
}

fn backoff(attempt: u32, prev: Option<&reqwest::Response>) -> Duration {
    if let Some(r) = prev {
        if let Some(ra) = r
            .headers()
            .get("Retry-After")
            .and_then(|v| v.to_str().ok())
            .and_then(|s| s.parse::<u64>().ok())
        {
            if ra > 0 {
                return Duration::from_secs(ra);
            }
        }
    }
    let ms = (2_u64.saturating_pow(attempt.saturating_sub(1))) * 500;
    Duration::from_millis(std::cmp::min(ms, 30_000))
}

fn parse_retry_after(header: Option<&str>) -> Duration {
    match header {
        Some(s) => match s.trim().parse::<u64>() {
            Ok(n) => Duration::from_secs(n),
            Err(_) => Duration::ZERO,
        },
        None => Duration::ZERO,
    }
}

fn map_error(status: u16, retry_after_header: Option<&str>, body: &[u8]) -> Error {
    let parsed: Value = serde_json::from_slice(body).unwrap_or(Value::Null);
    let message = parsed
        .get("error")
        .and_then(|v| v.as_str())
        .unwrap_or("")
        .to_string();
    let code = parsed
        .get("code")
        .and_then(|v| v.as_str())
        .map(|s| s.to_string());
    let balance = parsed.get("balance").and_then(|v| v.as_i64()).unwrap_or(0);
    let required = parsed.get("required").and_then(|v| v.as_i64()).unwrap_or(0);
    let message = if message.is_empty() {
        format!("HTTP {status}")
    } else {
        message
    };

    match status {
        400 => Error::InvalidRequest { message, status, code },
        401 => Error::Authentication { message, status, code },
        402 => Error::InsufficientCredits { message, status, code, balance, required },
        429 => Error::RateLimit {
            message,
            status,
            code,
            retry_after: parse_retry_after(retry_after_header),
        },
        s if s >= 500 => Error::Server { message, status, code },
        _ => Error::Other(format!("HTTP {status}: {message}")),
    }
}

fn require_string(v: &Value, field: &str) -> Result<String> {
    v.get(field)
        .and_then(|x| x.as_str())
        .filter(|s| !s.is_empty())
        .map(|s| s.to_string())
        .ok_or_else(|| Error::Other(format!("API response missing required field: {field}")))
}

fn decode_field<T: DeserializeOwned + Default>(parent: &Value, field: &str) -> Result<T> {
    match parent.get(field) {
        Some(v) if !v.is_null() => serde_json::from_value(v.clone())
            .map_err(|e| Error::Other(format!("decode response: {e}"))),
        _ => Ok(T::default()),
    }
}

// ── builder ──────────────────────────────────────────────────────────────

/// Builder for [`Client`]. All knobs have sensible defaults; the only thing
/// you must provide is an API key (or set `ANAKIN_API_KEY`).
pub struct ClientBuilder {
    api_key: Option<String>,
    base_url: Option<String>,
    timeout: Duration,
    max_retries: u32,
    poll_interval: Duration,
    poll_max_interval: Duration,
    poll_timeout: Duration,
    http: Option<reqwest::Client>,
}

impl Default for ClientBuilder {
    fn default() -> Self {
        Self {
            api_key: None,
            base_url: None,
            timeout: Duration::from_secs(60),
            max_retries: 4,
            poll_interval: Duration::from_secs(1),
            poll_max_interval: Duration::from_secs(10),
            poll_timeout: Duration::from_secs(300),
            http: None,
        }
    }
}

impl ClientBuilder {
    pub fn api_key(mut self, key: impl Into<String>) -> Self {
        self.api_key = Some(key.into());
        self
    }

    pub fn base_url(mut self, url: impl Into<String>) -> Self {
        self.base_url = Some(url.into());
        self
    }

    pub fn timeout(mut self, t: Duration) -> Self {
        self.timeout = t;
        self
    }

    pub fn max_retries(mut self, n: u32) -> Self {
        self.max_retries = n;
        self
    }

    pub fn poll_interval(mut self, d: Duration) -> Self {
        self.poll_interval = d;
        self
    }

    pub fn poll_max_interval(mut self, d: Duration) -> Self {
        self.poll_max_interval = d;
        self
    }

    pub fn poll_timeout(mut self, d: Duration) -> Self {
        self.poll_timeout = d;
        self
    }

    /// Inject a pre-built `reqwest::Client` (e.g. with custom proxies).
    pub fn http_client(mut self, c: reqwest::Client) -> Self {
        self.http = Some(c);
        self
    }

    pub fn build(self) -> Result<Client> {
        let api_key = self
            .api_key
            .or_else(|| std::env::var("ANAKIN_API_KEY").ok())
            .filter(|s| !s.is_empty())
            .ok_or_else(|| {
                Error::Other(
                    "no API key — call .api_key(...) on the builder or set ANAKIN_API_KEY".into(),
                )
            })?;

        let http = match self.http {
            Some(c) => c,
            None => reqwest::Client::builder()
                .timeout(self.timeout)
                .user_agent(format!("anakin-rust/{VERSION}"))
                .build()
                .map_err(|e| Error::Other(format!("build http client: {e}")))?,
        };

        Ok(Client {
            http,
            api_key,
            base_url: self
                .base_url
                .unwrap_or_else(|| DEFAULT_BASE_URL.to_string())
                .trim_end_matches('/')
                .to_string(),
            max_retries: self.max_retries,
            poll_interval: self.poll_interval,
            poll_max_interval: self.poll_max_interval,
            poll_timeout: self.poll_timeout,
        })
    }
}