//! [`Highflame`] — async HTTP client for Highflame guardrails.

use std::{fmt, sync::Arc, time::Duration};

use reqwest::header::{HeaderMap, HeaderValue, ACCEPT, AUTHORIZATION, USER_AGENT};
use serde_json::json;
use tokio::sync::Mutex;

use crate::{
    error::HighflameError,
    stream::{parse_sse_response, SseEvent},
    types_gen::{
        ContentType, DebugPoliciesResponse, DetectRequest, DetectResponse, GuardRequest,
        GuardResponse, ListDetectorsResponse, Mode, TokenResponse, ToolContext,
    },
};

// ── Constants ─────────────────────────────────────────────────────────────────

/// SDK version — kept in sync with Cargo.toml via `make version-sync`.
pub const VERSION: &str = env!("CARGO_PKG_VERSION");

const DEFAULT_TIMEOUT: Duration = Duration::from_secs(30);
const DEFAULT_MAX_RETRIES: u32 = 2;
/// Refresh the token this many seconds before it actually expires.
const TOKEN_REFRESH_BUFFER: Duration = Duration::from_secs(60);
/// HTTP status codes that trigger a retry.
const RETRY_STATUS_CODES: &[u16] = &[429, 500, 502, 503, 504];

const SAAS_BASE_URL: &str = "https://api.highflame.ai";
const SAAS_TOKEN_URL: &str = "https://control.highflame.ai/v1/auth/oauth2/token";

// ── Token cache ───────────────────────────────────────────────────────────────

#[derive(Debug, Clone)]
struct CachedToken {
    access_token: String,
    expires_at: tokio::time::Instant,
    account_id: String,
    project_id: String,
    #[allow(dead_code)]
    gateway_id: String,
}

// ── Options ───────────────────────────────────────────────────────────────────

/// Configuration for [`Highflame`].
///
/// ```no_run
/// use highflame::HighflameOptions;
/// use std::time::Duration;
///
/// let opts = HighflameOptions::new("hf_sk_my_key")
///     .base_url("https://shield.self-hosted.example.com")
///     .timeout(Duration::from_secs(10))
///     .max_retries(1);
/// ```
#[derive(Debug, Clone)]
pub struct HighflameOptions {
    pub(crate) api_key: String,
    pub(crate) base_url: Option<String>,
    pub(crate) token_url: Option<String>,
    pub(crate) timeout: Option<Duration>,
    pub(crate) max_retries: Option<u32>,
    pub(crate) account_id: Option<String>,
    pub(crate) project_id: Option<String>,
    pub(crate) default_headers: Option<std::collections::HashMap<String, String>>,
}

impl HighflameOptions {
    /// Create options from just an API key; all other fields take defaults.
    pub fn new(api_key: impl Into<String>) -> Self {
        Self {
            api_key: api_key.into(),
            base_url: None,
            token_url: None,
            timeout: None,
            max_retries: None,
            account_id: None,
            project_id: None,
            default_headers: None,
        }
    }

    /// Override the guard service base URL (default: Highflame SaaS endpoint).
    pub fn base_url(mut self, url: impl Into<String>) -> Self {
        self.base_url = Some(url.into());
        self
    }

    /// Override the token exchange URL (default: Highflame SaaS endpoint).
    /// Used for `hf_sk_*` (service keys), `hfa_*` (agent keys), and `zid_sk_*` (ZeroID keys).
    pub fn token_url(mut self, url: impl Into<String>) -> Self {
        self.token_url = Some(url.into());
        self
    }

    /// Request timeout (default: 30 s).
    pub fn timeout(mut self, t: Duration) -> Self {
        self.timeout = Some(t);
        self
    }

    /// Number of retries on 429 / 5xx (default: 2).
    pub fn max_retries(mut self, n: u32) -> Self {
        self.max_retries = Some(n);
        self
    }

    /// Override `X-Account-ID` header; takes precedence over the token exchange
    /// value.
    pub fn account_id(mut self, id: impl Into<String>) -> Self {
        self.account_id = Some(id.into());
        self
    }

    /// Override `X-Project-ID` header; takes precedence over the token exchange
    /// value.
    pub fn project_id(mut self, id: impl Into<String>) -> Self {
        self.project_id = Some(id.into());
        self
    }

    /// Extra headers merged into every request. Useful for custom routing
    /// or identity headers (e.g., `X-Custom-Routing: my-service`).
    ///
    /// Custom headers cannot override SDK-managed headers (`Authorization`,
    /// `Accept`, `User-Agent`, `X-Account-ID`, `X-Project-ID`).
    pub fn default_headers(mut self, headers: std::collections::HashMap<String, String>) -> Self {
        self.default_headers = Some(headers);
        self
    }
}

// ── Inner shared state ────────────────────────────────────────────────────────

struct Inner {
    base_url: String,
    api_key: String,
    token_url: Option<String>,
    timeout: Duration,
    max_retries: u32,
    override_account_id: Option<String>,
    override_project_id: Option<String>,
    default_headers: HeaderMap,
    static_headers: HeaderMap,
    http: reqwest::Client,
    token_cache: Mutex<Option<CachedToken>>,
}

// ── Resource types ────────────────────────────────────────────────────────────

/// Guard resource — evaluate content against guard policies.
pub struct GuardResource<'a> {
    client: &'a Highflame,
}

impl<'a> GuardResource<'a> {
    /// Evaluate content against guard policies.
    ///
    /// `POST /v1/shield/guard` — full detect + Cedar evaluate.
    pub async fn evaluate(&self, request: &GuardRequest) -> Result<GuardResponse, HighflameError> {
        let body = serde_json::to_value(request)?;
        let resp = self
            .client
            .send_with_retry(reqwest::Method::POST, "/v1/shield/guard", Some(&body))
            .await?;
        if !resp.status().is_success() {
            return Err(parse_api_error(resp).await);
        }
        Ok(resp.json().await?)
    }

    /// Shorthand: evaluate a user prompt.
    pub async fn evaluate_prompt(
        &self,
        content: &str,
        mode: Option<Mode>,
        session_id: Option<&str>,
    ) -> Result<GuardResponse, HighflameError> {
        self.evaluate(&GuardRequest {
            content: content.to_string(),
            content_type: ContentType::Prompt,
            action: "process_prompt".to_string(),
            mode,
            session_id: session_id.map(str::to_string),
            ..Default::default()
        })
        .await
    }

    /// Shorthand: evaluate a tool call.
    pub async fn evaluate_tool_call(
        &self,
        tool_name: &str,
        arguments: Option<std::collections::HashMap<String, serde_json::Value>>,
        mode: Option<Mode>,
        session_id: Option<&str>,
    ) -> Result<GuardResponse, HighflameError> {
        self.evaluate(&GuardRequest {
            content: format!("Tool call: {tool_name}"),
            content_type: ContentType::ToolCall,
            action: "call_tool".to_string(),
            mode,
            session_id: session_id.map(str::to_string),
            tool: Some(ToolContext {
                name: tool_name.to_string(),
                is_builtin: false,
                arguments,
                ..Default::default()
            }),
            ..Default::default()
        })
        .await
    }

    /// Evaluate content with SSE streaming.
    ///
    /// `POST /v1/shield/guard/stream` — returns a `Stream` of [`SseEvent`]s.
    pub async fn stream(
        &self,
        request: &GuardRequest,
    ) -> Result<impl futures_util::Stream<Item = Result<SseEvent, HighflameError>>, HighflameError>
    {
        let url = format!("{}/v1/shield/guard/stream", self.client.inner.base_url);
        let mut headers = self.client.get_auth_headers().await?;
        headers.insert(ACCEPT, HeaderValue::from_static("text/event-stream"));

        let resp = self
            .client
            .inner
            .http
            .post(&url)
            .headers(headers)
            .json(request)
            .timeout(self.client.inner.timeout)
            .send()
            .await
            .map_err(|e| HighflameError::ApiConnection(e.to_string()))?;

        if !resp.status().is_success() {
            return Err(parse_api_error(resp).await);
        }

        Ok(parse_sse_response(resp))
    }
}

/// Detect resource — run detectors without policy evaluation.
pub struct DetectResource<'a> {
    client: &'a Highflame,
}

impl<'a> DetectResource<'a> {
    /// Run detectors without policy evaluation.
    ///
    /// `POST /v1/shield/detect` — detection only, no Cedar evaluate.
    pub async fn run(&self, request: &DetectRequest) -> Result<DetectResponse, HighflameError> {
        let body = serde_json::to_value(request)?;
        let resp = self
            .client
            .send_with_retry(reqwest::Method::POST, "/v1/shield/detect", Some(&body))
            .await?;
        if !resp.status().is_success() {
            return Err(parse_api_error(resp).await);
        }
        Ok(resp.json().await?)
    }
}

/// Detectors resource — list available detectors.
pub struct DetectorsResource<'a> {
    client: &'a Highflame,
}

impl<'a> DetectorsResource<'a> {
    /// List available detectors.
    ///
    /// `GET /v1/shield/detectors` — returns detector names, tiers, and health.
    pub async fn list(&self) -> Result<ListDetectorsResponse, HighflameError> {
        let resp = self
            .client
            .send_with_retry(reqwest::Method::GET, "/v1/shield/detectors", None)
            .await?;
        if !resp.status().is_success() {
            return Err(parse_api_error(resp).await);
        }
        Ok(resp.json().await?)
    }
}

/// Debug resource — inspect loaded policies and other debug information.
pub struct DebugResource<'a> {
    client: &'a Highflame,
}

impl<'a> DebugResource<'a> {
    /// Inspect loaded policies.
    ///
    /// `GET /v1/shield/debug/policies` — returns loaded policies with metadata.
    pub async fn policies(
        &self,
        product: Option<&str>,
    ) -> Result<DebugPoliciesResponse, HighflameError> {
        let product = product.unwrap_or("guardrails");
        let path = format!("/v1/shield/debug/policies?product={product}");
        let resp = self
            .client
            .send_with_retry(reqwest::Method::GET, &path, None)
            .await?;
        if !resp.status().is_success() {
            return Err(parse_api_error(resp).await);
        }
        Ok(resp.json().await?)
    }
}

// ── Client ────────────────────────────────────────────────────────────────────

/// Async HTTP client for Highflame guardrails.
///
/// Cheap to clone — all clones share the underlying connection pool and token
/// cache via [`Arc`].
///
/// # Resource-based API
///
/// ```no_run
/// use highflame::{Highflame, HighflameOptions};
///
/// # #[tokio::main] async fn main() -> Result<(), Box<dyn std::error::Error>> {
/// let client = Highflame::new(HighflameOptions::new("hf_sk_my_key"));
///
/// let resp = client.guard().evaluate_prompt("Hello", None, None).await?;
/// println!("{:?}", resp.decision);
///
/// let policies = client.debug().policies(None).await?;
/// println!("policies: {:?}", policies);
/// # Ok(()) }
/// ```
#[derive(Clone)]
pub struct Highflame {
    inner: Arc<Inner>,
}

impl Highflame {
    /// Create a new [`Highflame`] client.
    pub fn new(options: HighflameOptions) -> Self {
        let base_url = options
            .base_url
            .unwrap_or_else(|| SAAS_BASE_URL.to_string());
        let base_url = base_url.trim_end_matches('/').to_string();

        let needs_exchange = options.api_key.starts_with("hf_sk")
            || options.api_key.starts_with("hfa_")
            || options.api_key.starts_with("zid_sk");

        let token_url = if needs_exchange {
            Some(
                options
                    .token_url
                    .unwrap_or_else(|| SAAS_TOKEN_URL.to_string()),
            )
        } else {
            None
        };

        let timeout = options.timeout.unwrap_or(DEFAULT_TIMEOUT);
        let max_retries = options.max_retries.unwrap_or(DEFAULT_MAX_RETRIES);

        // Build default_headers from user-provided map.
        let mut custom_headers = HeaderMap::new();
        if let Some(ref map) = options.default_headers {
            for (k, v) in map {
                let name = reqwest::header::HeaderName::from_bytes(k.as_bytes())
                    .expect("default_headers key contains invalid header name");
                let val = HeaderValue::from_str(v)
                    .expect("default_headers value contains invalid header characters");
                custom_headers.insert(name, val);
            }
        }

        // SDK-managed headers are set *after* custom headers so they cannot be
        // overridden by user-provided default_headers.
        let mut static_headers = custom_headers.clone();
        static_headers.insert(
            AUTHORIZATION,
            HeaderValue::from_str(&format!("Bearer {}", options.api_key))
                .expect("api_key contains invalid header characters"),
        );
        static_headers.insert(ACCEPT, HeaderValue::from_static("application/json"));
        static_headers.insert(
            USER_AGENT,
            HeaderValue::from_str(&format!("highflame-rust/{VERSION}"))
                .expect("invalid user-agent"),
        );
        if let Some(ref id) = options.account_id {
            static_headers.insert(
                reqwest::header::HeaderName::from_static("x-account-id"),
                HeaderValue::from_str(id).expect("account_id contains invalid header characters"),
            );
        }
        if let Some(ref id) = options.project_id {
            static_headers.insert(
                reqwest::header::HeaderName::from_static("x-project-id"),
                HeaderValue::from_str(id).expect("project_id contains invalid header characters"),
            );
        }

        let http = reqwest::Client::builder()
            .build()
            .expect("failed to build reqwest client");

        Self {
            inner: Arc::new(Inner {
                base_url,
                api_key: options.api_key,
                token_url,
                timeout,
                max_retries,
                override_account_id: options.account_id,
                override_project_id: options.project_id,
                default_headers: custom_headers,
                static_headers,
                http,
                token_cache: Mutex::new(None),
            }),
        }
    }

    // ── Resource accessors ────────────────────────────────────────────────────

    /// Access guard operations (evaluate, stream).
    pub fn guard(&self) -> GuardResource<'_> {
        GuardResource { client: self }
    }

    /// Access detect operations (run detectors without policy evaluation).
    pub fn detect(&self) -> DetectResource<'_> {
        DetectResource { client: self }
    }

    /// Access detector listing operations.
    pub fn detectors(&self) -> DetectorsResource<'_> {
        DetectorsResource { client: self }
    }

    /// Access debug operations (inspect policies).
    pub fn debug(&self) -> DebugResource<'_> {
        DebugResource { client: self }
    }

    // ── Public accessors ──────────────────────────────────────────────────────

    /// Account ID from the last token exchange, or the constructor override.
    pub async fn account_id(&self) -> String {
        if let Some(ref id) = self.inner.override_account_id {
            return id.clone();
        }
        self.inner
            .token_cache
            .lock()
            .await
            .as_ref()
            .map(|t| t.account_id.clone())
            .unwrap_or_default()
    }

    /// Project ID from the last token exchange, or the constructor override.
    pub async fn project_id(&self) -> String {
        if let Some(ref id) = self.inner.override_project_id {
            return id.clone();
        }
        self.inner
            .token_cache
            .lock()
            .await
            .as_ref()
            .map(|t| t.project_id.clone())
            .unwrap_or_default()
    }

    // ── Token exchange ────────────────────────────────────────────────────────

    fn build_token_headers(&self, token: &CachedToken) -> Result<HeaderMap, HighflameError> {
        // Start from custom default_headers, then overlay SDK-managed headers.
        let mut headers = self.inner.default_headers.clone();
        headers.insert(
            AUTHORIZATION,
            HeaderValue::from_str(&format!("Bearer {}", token.access_token))
                .map_err(|e| HighflameError::ApiConnection(e.to_string()))?,
        );
        headers.insert(ACCEPT, HeaderValue::from_static("application/json"));
        headers.insert(
            USER_AGENT,
            HeaderValue::from_str(&format!("highflame-rust/{VERSION}"))
                .map_err(|e| HighflameError::ApiConnection(e.to_string()))?,
        );

        let account_id = self
            .inner
            .override_account_id
            .as_deref()
            .or_else(|| (!token.account_id.is_empty()).then_some(token.account_id.as_str()));
        let project_id = self
            .inner
            .override_project_id
            .as_deref()
            .or_else(|| (!token.project_id.is_empty()).then_some(token.project_id.as_str()));

        if let Some(id) = account_id {
            headers.insert(
                reqwest::header::HeaderName::from_static("x-account-id"),
                HeaderValue::from_str(id)
                    .map_err(|e| HighflameError::ApiConnection(e.to_string()))?,
            );
        }
        if let Some(id) = project_id {
            headers.insert(
                reqwest::header::HeaderName::from_static("x-project-id"),
                HeaderValue::from_str(id)
                    .map_err(|e| HighflameError::ApiConnection(e.to_string()))?,
            );
        }
        Ok(headers)
    }

    async fn exchange_token(&self) -> Result<CachedToken, HighflameError> {
        let url = self
            .inner
            .token_url
            .as_ref()
            .expect("exchange_token called without a token_url");

        log::debug!("Token exchange: POST {}", url);

        let resp = self
            .inner
            .http
            .post(url)
            .json(&json!({ "grant_type": "api_key", "api_key": self.inner.api_key }))
            .timeout(self.inner.timeout)
            .send()
            .await
            .map_err(|e| {
                log::warn!("Token exchange failed: {}", e);
                HighflameError::ApiConnection(e.to_string())
            })?;

        if !resp.status().is_success() {
            log::warn!("Token exchange error: {}", resp.status());
            return Err(parse_api_error(resp).await);
        }

        let tok: TokenResponse = resp.json().await?;
        log::debug!("Token exchange: success (expires_in={}s)", tok.expires_in);
        let ttl = Duration::from_secs(tok.expires_in as u64).saturating_sub(TOKEN_REFRESH_BUFFER);
        Ok(CachedToken {
            access_token: tok.access_token,
            expires_at: tokio::time::Instant::now() + ttl,
            account_id: tok.account_id.unwrap_or_default(),
            project_id: tok.project_id.unwrap_or_default(),
            gateway_id: tok.gateway_id.unwrap_or_default(),
        })
    }

    /// Return auth headers, exchanging or refreshing the JWT as needed.
    ///
    /// When the `telemetry` feature is enabled and a global `TextMapPropagator`
    /// is registered, W3C traceparent/tracestate headers are automatically
    /// injected for distributed tracing (SDK → Shield).
    pub async fn get_auth_headers(&self) -> Result<HeaderMap, HighflameError> {
        let mut headers = if self.inner.token_url.is_none() {
            self.inner.static_headers.clone()
        } else {
            let mut cache = self.inner.token_cache.lock().await;

            let valid = cache
                .as_ref()
                .filter(|t| tokio::time::Instant::now() < t.expires_at);

            if let Some(token) = valid {
                self.build_token_headers(token)?
            } else {
                let fresh = self.exchange_token().await?;
                let h = self.build_token_headers(&fresh)?;
                *cache = Some(fresh);
                h
            }
        };

        inject_trace_context(&mut headers);
        Ok(headers)
    }

    // ── Request helpers (used by resource types) ──────────────────────────────

    pub(crate) async fn send_with_retry(
        &self,
        method: reqwest::Method,
        path: &str,
        json_body: Option<&serde_json::Value>,
    ) -> Result<reqwest::Response, HighflameError> {
        let url = format!("{}{}", self.inner.base_url, path);
        let mut last_err =
            HighflameError::ApiConnection("request failed after all retries".to_string());
        let mut idempotency_key: Option<String> = None;

        for attempt in 0..=self.inner.max_retries {
            if attempt > 0 {
                let base_ms = 1_000 * 2u64.pow(attempt - 1);
                // ±25% jitter to prevent thundering herd
                let jitter_range = base_ms / 4;
                let jitter = if jitter_range > 0 {
                    let seed = std::time::SystemTime::now()
                        .duration_since(std::time::UNIX_EPOCH)
                        .unwrap_or_default()
                        .subsec_nanos() as u64;
                    (seed % (jitter_range * 2)).saturating_sub(jitter_range)
                } else {
                    0
                };
                let delay_ms = (base_ms as i64 + jitter as i64).max(0) as u64;
                log::debug!(
                    "Retrying {} {} (attempt {}/{}) after {}ms",
                    method,
                    path,
                    attempt + 1,
                    self.inner.max_retries + 1,
                    delay_ms,
                );
                tokio::time::sleep(Duration::from_millis(delay_ms)).await;
                if idempotency_key.is_none() {
                    idempotency_key = Some(uuid_v4());
                }
            }

            log::debug!("{} {}", method, path);
            if let Some(body) = json_body {
                let body_str = body.to_string();
                if body_str.len() <= 1024 {
                    log::debug!("Request body: {}", body_str);
                } else {
                    log::debug!(
                        "Request body: {}... ({} chars)",
                        &body_str[..1024],
                        body_str.len()
                    );
                }
            }

            let mut headers = self.get_auth_headers().await?;
            if let Some(ref key) = idempotency_key {
                headers.insert(
                    reqwest::header::HeaderName::from_static("x-idempotency-key"),
                    HeaderValue::from_str(key)
                        .map_err(|e| HighflameError::ApiConnection(e.to_string()))?,
                );
            }

            let mut req = self
                .inner
                .http
                .request(method.clone(), &url)
                .headers(headers)
                .timeout(self.inner.timeout);

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

            let resp = match req.send().await {
                Ok(r) => r,
                Err(e) if e.is_timeout() => {
                    log::debug!("Request timed out: {} {}", method, path);
                    last_err = HighflameError::ApiConnection(format!("request timed out: {e}"));
                    continue;
                }
                Err(e) => return Err(HighflameError::ApiConnection(e.to_string())),
            };

            log::debug!("{} {} → {}", method, path, resp.status().as_u16());

            if !RETRY_STATUS_CODES.contains(&resp.status().as_u16()) {
                return Ok(resp);
            }

            last_err = parse_api_error(resp).await;
        }

        Err(last_err)
    }
}

impl fmt::Display for Highflame {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "Highflame(base_url={:?})", self.inner.base_url)
    }
}

impl fmt::Debug for Highflame {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Highflame")
            .field("base_url", &self.inner.base_url)
            .finish()
    }
}

// ── Helpers ───────────────────────────────────────────────────────────────────

/// Parse an RFC 9457 Problem Details body from a non-2xx response.
pub(crate) async fn parse_api_error(resp: reqwest::Response) -> HighflameError {
    let status = resp.status().as_u16();
    let (title, detail) = match resp.json::<serde_json::Value>().await {
        Ok(body) => (
            body["title"].as_str().unwrap_or("Error").to_string(),
            body["detail"].as_str().unwrap_or("").to_string(),
        ),
        Err(_) => ("Error".to_string(), String::new()),
    };

    match status {
        401 => HighflameError::Authentication(detail),
        429 => HighflameError::RateLimit(detail),
        _ => HighflameError::Api {
            status,
            title,
            detail,
        },
    }
}

/// Generate a random UUID v4 for idempotency keys.
fn uuid_v4() -> String {
    uuid::Uuid::new_v4().to_string()
}

// ---------------------------------------------------------------------------
// Optional OpenTelemetry trace context propagation
// ---------------------------------------------------------------------------

/// Inject W3C traceparent/tracestate into outgoing request headers.
///
/// When the `telemetry` feature is enabled and a global propagator is
/// registered, this injects the current span context. Otherwise, it's a no-op.
#[cfg(feature = "telemetry")]
fn inject_trace_context(headers: &mut HeaderMap) {
    struct HeaderInjector<'a>(&'a mut HeaderMap);

    impl opentelemetry::propagation::Injector for HeaderInjector<'_> {
        fn set(&mut self, key: &str, value: String) {
            if let (Ok(name), Ok(val)) = (
                reqwest::header::HeaderName::from_bytes(key.as_bytes()),
                HeaderValue::from_str(&value),
            ) {
                self.0.insert(name, val);
            } else {
                log::warn!("opentelemetry: failed to inject trace context header '{key}'");
            }
        }
    }

    opentelemetry::global::get_text_map_propagator(|p| {
        let ctx = opentelemetry::Context::current();
        p.inject_context(&ctx, &mut HeaderInjector(headers));
    });
}

#[cfg(not(feature = "telemetry"))]
fn inject_trace_context(_headers: &mut HeaderMap) {
    // No-op when telemetry feature is disabled.
}
