//! smoke_test — Auth smoke test for the Highflame Rust SDK.
//!
//! Verifies:
//!   1. Direct JWT auth   → GET /v1/shield/detectors, POST /v1/shield/guard, GET /v1/shield/debug/policies
//!   2. Service-key flow  → POST token exchange, JWT used as Bearer,
//!      token cached across calls
//!
//! # Environment variables
//!
//! | Variable            | Description                                              |
//! |---------------------|----------------------------------------------------------|
//! | `HIGHFLAME_API_KEY` | `hf_sk_xxx` service key (primary credential)            |
//! | `HIGHFLAME_JWT`     | Raw JWT for direct-auth tests (optional)                 |
//! | `SHIELD_BASE_URL`   | Override guard service URL (default: Highflame SaaS)    |
//! | `SHIELD_TOKEN_URL`  | Override token exchange URL (default: Highflame SaaS)   |
//!
//! # Usage
//!
//! ```bash
//! # Service-key only (SaaS endpoints):
//! HIGHFLAME_API_KEY=hf_sk_xxx cargo run --example smoke_test
//!
//! # Override endpoints (e.g. dev environment):
//! SHIELD_BASE_URL=https://api-dev.highflame.dev \
//! SHIELD_TOKEN_URL=https://control-dev.highflame.dev/v1/auth/oauth2/token \
//! HIGHFLAME_API_KEY=hf_sk_xxx cargo run --example smoke_test
//! ```
//!
//! Exit codes: `0` — all checks passed, `1` — one or more checks failed.

use std::{env, io::IsTerminal as _, process};

use highflame::{ContentType, GuardRequest, Highflame, HighflameOptions};

// ── Colour helpers ────────────────────────────────────────────────────────────

struct Out {
    tty: bool,
}

impl Out {
    fn new() -> Self {
        Self {
            tty: std::io::stdout().is_terminal(),
        }
    }
    fn c(&self, code: &str, s: &str) -> String {
        if self.tty {
            format!("{code}{s}\x1b[0m")
        } else {
            s.to_string()
        }
    }
    fn pass(&self) -> String {
        self.c("\x1b[32m", "✓")
    }
    fn fail(&self) -> String {
        self.c("\x1b[31m", "✗")
    }
    fn info(&self) -> String {
        self.c("\x1b[34m", "·")
    }
    fn warn(&self) -> String {
        self.c("\x1b[33m", "!")
    }
    fn skip(&self) -> String {
        self.c("\x1b[33m", "–")
    }
}

// ── Check runner ──────────────────────────────────────────────────────────────

struct Results {
    checks: Vec<(String, bool)>,
    out: Out,
}

impl Results {
    fn new() -> Self {
        Self {
            checks: Vec::new(),
            out: Out::new(),
        }
    }

    fn record(&mut self, label: &str, ok: bool) {
        if ok {
            println!("  {}  {label}", self.out.pass());
        } else {
            println!("  {}  {label}", self.out.fail());
        }
        self.checks.push((label.to_string(), ok));
    }

    fn summary(&self) -> bool {
        let passed = self.checks.iter().filter(|(_, ok)| *ok).count();
        let total = self.checks.len();
        println!("\n{}", "─".repeat(57));
        if passed == total {
            println!("{}  All {total} checks passed", self.out.pass());
            true
        } else {
            println!("{}  {passed}/{total} checks passed", self.out.fail());
            for (label, ok) in &self.checks {
                let mark = if *ok {
                    self.out.pass()
                } else {
                    self.out.fail()
                };
                println!("       {mark}  {label}");
            }
            false
        }
    }
}

// ── Main ──────────────────────────────────────────────────────────────────────

#[tokio::main]
async fn main() {
    let out = Out::new();

    let base_url =
        env::var("SHIELD_BASE_URL").unwrap_or_else(|_| "https://api-dev.highflame.dev".to_string());
    let token_url = env::var("SHIELD_TOKEN_URL")
        .unwrap_or_else(|_| "https://control-dev.highflame.dev/v1/auth/oauth2/token".to_string());
    let api_key = env::var("HIGHFLAME_API_KEY").unwrap_or_default();
    let jwt = env::var("HIGHFLAME_JWT").unwrap_or_default();

    println!("\n{}  Base URL  : {base_url}", out.info());
    println!("{}  Token URL : {token_url}", out.info());
    println!();

    if api_key.is_empty() && jwt.is_empty() {
        println!(
            "{}  Neither HIGHFLAME_API_KEY nor HIGHFLAME_JWT is set — cannot run any checks.",
            out.fail()
        );
        process::exit(1);
    }

    let test_req = GuardRequest {
        content: "What is the capital of France?".to_string(),
        content_type: ContentType::Prompt,
        action: "process_prompt".to_string(),
        ..Default::default()
    };

    let mut results = Results::new();

    // ── 1. Direct JWT flow ────────────────────────────────────────────────────

    let key_needs_exchange =
        jwt.starts_with("hf_sk") || jwt.starts_with("hfa_") || jwt.starts_with("zid_sk");

    if jwt.is_empty() {
        println!(
            "{}  Direct JWT tests skipped (HIGHFLAME_JWT not set)",
            out.skip()
        );
    } else if key_needs_exchange {
        println!(
            "{}  Direct JWT tests skipped (HIGHFLAME_JWT looks like a service key)",
            out.skip()
        );
    } else {
        println!("── 1. Direct JWT flow {}", "─".repeat(34));

        let direct = Highflame::new(
            HighflameOptions::new(&jwt)
                .base_url(&base_url)
                .max_retries(0),
        );

        // GET /v1/shield/detectors
        match direct.detectors().list().await {
            Ok(r) => {
                println!("         {} detector(s) registered", r.count);
                results.record("GET  /v1/shield/detectors", true);
            }
            Err(e) => {
                println!("         {e}");
                results.record("GET  /v1/shield/detectors", false);
            }
        }

        // POST /v1/shield/guard
        match direct.guard().evaluate(&test_req).await {
            Ok(resp) => {
                println!(
                    "         decision={:?}  latency_ms={:?}",
                    resp.decision, resp.latency_ms
                );
                results.record("POST /v1/shield/guard  (raw JWT)", true);
            }
            Err(e) => {
                println!("         {e}");
                results.record("POST /v1/shield/guard  (raw JWT)", false);
            }
        }

        // GET /v1/shield/debug/policies
        match direct.debug().policies(None).await {
            Ok(r) => {
                let count = r.policies.len();
                println!("         {} policy/ies loaded", count);
                results.record("GET  /v1/shield/debug/policies", true);
            }
            Err(e) => {
                println!("         {e}");
                results.record("GET  /v1/shield/debug/policies", false);
            }
        }
    }

    // ── 2. Service-key flow ───────────────────────────────────────────────────

    if api_key.is_empty() {
        println!(
            "\n{}  Service-key tests skipped (HIGHFLAME_API_KEY not set)",
            out.skip()
        );
    } else {
        println!("\n── 2. Service-key flow {}", "─".repeat(35));

        let sk = Highflame::new(
            HighflameOptions::new(&api_key)
                .base_url(&base_url)
                .token_url(&token_url)
                .max_retries(0),
        );

        // POST token exchange: get_auth_headers() triggers exchange on first call.
        match sk.get_auth_headers().await {
            Ok(headers) => {
                let auth = headers
                    .get("authorization")
                    .and_then(|v| v.to_str().ok())
                    .unwrap_or("");

                let jwt_ok = auth.starts_with("Bearer ")
                    && &auth["Bearer ".len()..] != api_key.as_str()
                    && auth["Bearer ".len()..].starts_with("eyJ");

                let acct = sk.account_id().await;
                let proj = sk.project_id().await;
                println!("         account_id={acct:?}  project_id={proj:?}");
                if auth.len() > "Bearer ".len() + 40 {
                    println!(
                        "         JWT prefix : {}…",
                        &auth["Bearer ".len().."Bearer ".len() + 40]
                    );
                }
                if !jwt_ok {
                    println!("         exchange did not produce a valid JWT");
                }
                results.record("POST token exchange", jwt_ok);
            }
            Err(e) => {
                println!("         {e}");
                results.record("POST token exchange", false);
            }
        }

        // POST /v1/shield/guard with exchanged JWT.
        match sk.guard().evaluate(&test_req).await {
            Ok(resp) => {
                println!(
                    "         decision={:?}  latency_ms={:?}",
                    resp.decision, resp.latency_ms
                );
                results.record("POST /v1/shield/guard  (JWT auth)", true);
            }
            Err(e) => {
                println!("         {e}");
                results.record("POST /v1/shield/guard  (JWT auth)", false);
            }
        }

        // Token cached across calls: two consecutive get_auth_headers() must return same token.
        let (h1, h2) = (sk.get_auth_headers().await, sk.get_auth_headers().await);
        match (h1, h2) {
            (Ok(a), Ok(b)) => {
                let tok1 = a
                    .get("authorization")
                    .and_then(|v| v.to_str().ok())
                    .unwrap_or("");
                let tok2 = b
                    .get("authorization")
                    .and_then(|v| v.to_str().ok())
                    .unwrap_or("");
                let same = tok1 == tok2;
                if same {
                    println!("         Same token reused across both calls");
                } else {
                    println!("         Token was replaced — unexpected re-exchange");
                }
                results.record("Token cached across calls", same);
            }
            (Err(e), _) | (_, Err(e)) => {
                println!("         {e}");
                results.record("Token cached across calls", false);
            }
        }
    }

    // ── Summary ───────────────────────────────────────────────────────────────

    if results.checks.is_empty() {
        println!(
            "{}  No checks ran — verify your environment variables.",
            out.warn()
        );
        process::exit(1);
    }

    if !results.summary() {
        process::exit(1);
    }
}
