//! Amazon Bedrock provider real endpoint smoke test, focused on the thinking wire format.
//!
//! Verifies the per-model `ThinkingWireFormat` plumbing end-to-end against a live Bedrock
//! endpoint: a model configured `Adaptive` must accept a thinking request (Opus 4.6+ reject
//! the legacy `enabled` + `budget_tokens` shape with a 400), and the same model under
//! `Legacy` reproduces that 400 — proving the format actually changes the wire body.
//!
//! Usage:
//!
//! ```bash
//! # AWS creds from the environment (profile / env vars), as configured in .bashrc.
//! BEDROCK_MODEL='arn:aws:bedrock:...:application-inference-profile/...' \
//!   cargo run -p defect-llm --example bedrock_smoke -- [scenario]
//! ```
//!
//! `[scenario]` ∈ `text | thinking-adaptive | thinking-legacy | all`, defaults to `all`.
//!
//! Optional env:
//! - `BEDROCK_MODEL`: the model id / inference-profile ARN (default below)
//! - `AWS_REGION` / `AWS_PROFILE`: standard AWS resolution
//! - `RUST_LOG=defect_llm=debug,toac=debug` enables protocol-level debug logging

mod common;

use std::sync::Arc;

use agent_client_protocol_schema::StopReason as AcpStopReason;
use defect_agent::llm::{LlmProvider, SamplingParams};
use defect_llm::protocol::anthropic_messages::ThinkingWireFormat;
use defect_llm::provider::bedrock::{BedrockConfig, BedrockModel, BedrockProvider};

use common::{
    EXIT_FAIL, EXIT_OK, build_session, env_string, init_tracing, print_fail, print_pass,
    print_skip, run_turn_and_print, sampling_with_thinking, scenario_from_args,
};

/// Default to the opus-4-8 discount inference profile from `.defect/config.local.toml`.
const DEFAULT_MODEL: &str =
    "arn:aws:bedrock:us-east-1:258049739530:application-inference-profile/7k9nkdq0in6f";
const THINKING_BUDGET_TOKENS: u32 = 2048;

const TEXT_PROMPT: &str = "Say hello in one short sentence.";
const THINKING_PROMPT: &str = "Think step by step: a farmer has 17 sheep and all but 9 die. How many are left? \
     Show your reasoning briefly, then give the final number.";

#[tokio::main]
async fn main() {
    init_tracing();

    let model = env_string("BEDROCK_MODEL").unwrap_or_else(|| DEFAULT_MODEL.to_string());
    let scenario = scenario_from_args();
    println!("=== bedrock smoke: scenario={scenario} model={model} ===");

    let mut failed = 0u32;
    let mut ran = 0u32;

    for label in scenarios_for(&scenario) {
        ran += 1;
        let outcome = run_scenario(label, &model).await;
        match outcome {
            ScenarioOutcome::Pass => print_pass(label),
            ScenarioOutcome::Skip(reason) => print_skip(label, &reason),
            ScenarioOutcome::Fail(err) => {
                failed += 1;
                print_fail(label, &err);
            }
        }
    }

    println!("\n=== bedrock smoke done: ran={ran} failed={failed} ===");
    std::process::exit(if failed > 0 { EXIT_FAIL } else { EXIT_OK });
}

fn scenarios_for(name: &str) -> Vec<&'static str> {
    match name {
        "text" => vec!["text"],
        "thinking-adaptive" => vec!["thinking-adaptive"],
        "thinking-legacy" => vec!["thinking-legacy"],
        _ => vec!["text", "thinking-adaptive", "thinking-legacy"],
    }
}

enum ScenarioOutcome {
    Pass,
    Skip(String),
    Fail(String),
}

/// Builds a Bedrock provider whose single configured model uses the given thinking format.
async fn provider_for(
    model: &str,
    format: ThinkingWireFormat,
) -> Result<Arc<dyn LlmProvider>, String> {
    let cfg = BedrockConfig {
        default_model: Some(model.to_string()),
        models: vec![BedrockModel {
            id: model.to_string(),
            context_window: None,
            max_output_tokens: None,
            thinking_format: Some(format),
        }],
        ..Default::default()
    };
    BedrockProvider::new(cfg)
        .await
        .map(|p| Arc::new(p) as Arc<dyn LlmProvider>)
        .map_err(|e| format!("provider init failed: {e}"))
}

async fn run_scenario(label: &str, model: &str) -> ScenarioOutcome {
    println!("\n--- running: {label} ---");
    let res = match label {
        "text" => scenario_text(model).await,
        "thinking-adaptive" => scenario_thinking(model, ThinkingWireFormat::Adaptive).await,
        "thinking-legacy" => scenario_thinking_legacy(model).await,
        other => Err(format!("unknown scenario {other}")),
    };
    match res {
        Ok(None) => ScenarioOutcome::Pass,
        Ok(Some(reason)) => ScenarioOutcome::Skip(reason),
        Err(e) => ScenarioOutcome::Fail(e),
    }
}

async fn scenario_text(model: &str) -> Result<Option<String>, String> {
    let provider = provider_for(model, ThinkingWireFormat::Adaptive).await?;
    let session = build_session(provider, model, SamplingParams::default()).await;
    let (stop, text, _hits) = run_turn_and_print(session, TEXT_PROMPT)
        .await
        .map_err(|e| e.to_string())?;
    if !matches!(stop, AcpStopReason::EndTurn) {
        return Err(format!("unexpected stop reason: {stop:?}"));
    }
    if text.trim().is_empty() {
        return Err("empty assistant text".to_string());
    }
    Ok(None)
}

/// The adaptive path must succeed on Opus 4.8 with thinking requested.
async fn scenario_thinking(
    model: &str,
    format: ThinkingWireFormat,
) -> Result<Option<String>, String> {
    let provider = provider_for(model, format).await?;
    let sampling = sampling_with_thinking(Some(THINKING_BUDGET_TOKENS));
    let session = build_session(provider, model, sampling).await;
    let (stop, text, hits) = run_turn_and_print(session, THINKING_PROMPT)
        .await
        .map_err(|e| e.to_string())?;
    if !matches!(stop, AcpStopReason::EndTurn) {
        return Err(format!("unexpected stop reason: {stop:?}"));
    }
    if text.trim().is_empty() {
        return Err("empty assistant text".to_string());
    }
    if hits.thought_text.trim().is_empty() {
        return Ok(Some(format!(
            "no thinking text emitted on model {model} (adaptive accepted, but no summary returned)"
        )));
    }
    Ok(None)
}

/// The legacy path is expected to FAIL on Opus 4.8 (it rejects `enabled` + `budget_tokens`).
/// A turn error here is the *expected* outcome and counts as PASS; an unexpectedly clean
/// turn is the surprising result and counts as FAIL.
async fn scenario_thinking_legacy(model: &str) -> Result<Option<String>, String> {
    let provider = provider_for(model, ThinkingWireFormat::Legacy).await?;
    let sampling = sampling_with_thinking(Some(THINKING_BUDGET_TOKENS));
    let session = build_session(provider, model, sampling).await;
    match run_turn_and_print(session, THINKING_PROMPT).await {
        Err(e) => {
            println!("legacy format rejected as expected: {e}");
            Ok(None)
        }
        Ok((stop, _text, _hits)) => Ok(Some(format!(
            "legacy format unexpectedly accepted on {model} (stop={stop:?}); \
             this model may still support the enabled+budget shape"
        ))),
    }
}
