//! Runtime wrapper around `call_server` that handles the
//! `InputRequiredResult` retry loop.
//!
//! When the server returns `InputRequiredResult` (per SEP-2322), the
//! client must re-issue the same method with the user's answers
//! plus the echoed `requestState`. This module:
//!
//! 1. Calls `dispatch::call_server` once.
//! 2. If `recognize_input_required` matches the response, picks the
//!    answer set from the supplied `InputResponseFixture` (or, in
//!    the no-fixture case, fails fast so a non-interactive CI run
//!    surfaces the missing answer source).
//! 3. Validates every required `InputRequest` has a matching
//!    `InputResponse`, builds the retry `params` via
//!    `build_retry_params`, and re-dispatches.
//! 4. Stops after a max-rounds cap so a misbehaving server cannot
//!    loop forever.

use std::collections::BTreeMap;
use std::time::Duration;

use serde_json::Value;

use crate::network::ProxyConfig;
use crate::protocol::elicitation::{
    build_retry_params, recognize_input_required, InputRequest, InputResponse, RequestState,
};
use crate::suite::RestInputResponderSpec;

/// Default cap on the number of elicitation rounds the runner will
/// follow before giving up. Five matches the SEP-2322 example flow
/// and is enough for the deepest documented follow-up chain.
pub const DEFAULT_MAX_ELICITATION_ROUNDS: u8 = 5;

/// Source of canned answers for non-interactive (CI) runs.
///
/// v1 ships a static map: each `InputRequest::id` resolves to a
/// pre-supplied `InputResponse`. A future ticket can layer an
/// interactive variant (read from stdin, validate type) on the same
/// trait without changing call sites.
#[derive(Debug, Default, Clone)]
pub struct InputResponseFixture {
    responses: std::collections::HashMap<String, Value>,
}

impl InputResponseFixture {
    /// Build a fixture from an `(id, value)` map.
    pub fn from_map<I, K>(entries: I) -> Self
    where
        I: IntoIterator<Item = (K, Value)>,
        K: Into<String>,
    {
        Self {
            responses: entries.into_iter().map(|(k, v)| (k.into(), v)).collect(),
        }
    }

    /// Look up the answer for one `InputRequest` id.
    pub fn lookup(&self, id: &str) -> Option<&Value> {
        self.responses.get(id)
    }

    /// True when the fixture carries no entries; the runner uses
    /// this to take the "no answers available" branch immediately
    /// rather than try to satisfy a real elicitation against an
    /// empty map.
    pub fn is_empty(&self) -> bool {
        self.responses.is_empty()
    }
}

/// Reasons the runner could not satisfy an `InputRequiredResult`.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ElicitationError {
    /// A required `InputRequest` had no matching entry in the
    /// supplied [`InputResponseFixture`]. The runner cannot dispatch
    /// the retry without that answer.
    MissingResponse {
        /// The `InputRequest::id` the runner needed an answer for.
        request_id: String,
    },
    /// The server kept asking for more input after the configured
    /// max-rounds cap fired. Surfaces as a test failure so a buggy
    /// server cannot wedge the runner.
    MaxRoundsExceeded {
        /// The cap the runner used (typically
        /// `DEFAULT_MAX_ELICITATION_ROUNDS`).
        max_rounds: u8,
    },
    /// The fixture supplied no answers and the server elicited.
    /// Distinct from `MissingResponse` so the test failure message
    /// can point at "non-interactive mode without fixture" rather
    /// than at a specific missing id.
    NoFixture,
    /// An answer's value did not match the `kind` the server's
    /// [`InputRequest`] declared (`string` / `number` / `boolean`).
    /// Catches an author typo (a number where a string is required)
    /// and a misbehaving REST responder alike.
    WrongKind {
        /// The `InputRequest::id` whose answer was wrong-typed.
        request_id: String,
        /// The `kind` the request declared.
        expected: String,
    },
    /// The REST responder could not be reached or returned an
    /// unusable reply (non-2xx, invalid JSON, missing/duplicate/unknown
    /// ids). Carries a redacted human-readable reason; never the
    /// response body or any header value.
    Rest(String),
}

impl std::fmt::Display for ElicitationError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::MissingResponse { request_id } => {
                write!(f, "no fixture entry for inputRequest `{request_id}`")
            }
            Self::MaxRoundsExceeded { max_rounds } => write!(
                f,
                "server kept eliciting after {max_rounds} rounds; aborting"
            ),
            Self::NoFixture => write!(
                f,
                "server returned InputRequiredResult but no inputResponses fixture was supplied"
            ),
            Self::WrongKind {
                request_id,
                expected,
            } => write!(
                f,
                "answer for inputRequest `{request_id}` is not a {expected}"
            ),
            Self::Rest(reason) => write!(f, "REST elicitation responder: {reason}"),
        }
    }
}

/// Check that one answer's JSON type matches the `kind` its
/// [`InputRequest`] declared. Applied to both static and REST answers so
/// a wrong-typed answer fails the same way regardless of source. Unknown
/// `kind` values (future-reserved) are not constrained.
fn answer_matches_kind(kind: &str, value: &Value) -> bool {
    match kind {
        "string" => value.is_string(),
        "number" => value.is_number(),
        "boolean" => value.is_boolean(),
        _ => true,
    }
}

/// Validate every answer against its request's `kind`. The answers are
/// matched to requests by id; an answer with no matching request is
/// ignored here (the REST path rejects unknown ids separately).
pub fn validate_answer_kinds(
    requests: &[InputRequest],
    answers: &[InputResponse],
) -> Result<(), ElicitationError> {
    for answer in answers {
        if let Some(request) = requests.iter().find(|r| r.id == answer.id) {
            if !answer_matches_kind(&request.kind, &answer.value) {
                return Err(ElicitationError::WrongKind {
                    request_id: request.id.clone(),
                    expected: request.kind.clone(),
                });
            }
        }
    }
    Ok(())
}

/// Decide whether the runner has every answer it needs.
///
/// Walks the server's `input_requests`, looks each one up on the
/// fixture, and returns the typed [`InputResponse`] list when
/// every required request resolves. Returns a typed error when a
/// required id is missing.
///
/// Non-required (`required: false`) requests are skipped when the
/// fixture has no entry; their `InputResponse` is simply omitted.
pub fn collect_responses(
    requests: &[InputRequest],
    fixture: &InputResponseFixture,
) -> Result<Vec<InputResponse>, ElicitationError> {
    if fixture.is_empty() {
        return Err(ElicitationError::NoFixture);
    }
    let mut out = Vec::with_capacity(requests.len());
    for request in requests {
        match fixture.lookup(&request.id) {
            Some(value) => out.push(InputResponse {
                id: request.id.clone(),
                value: value.clone(),
            }),
            None if !request.required => continue,
            None => {
                return Err(ElicitationError::MissingResponse {
                    request_id: request.id.clone(),
                });
            }
        }
    }
    Ok(out)
}

/// Pure helper the executor uses on each elicitation round.
/// Returns the retry params on success or an [`ElicitationError`]
/// when the fixture cannot satisfy the request. Kept pure so the
/// retry-loop test can drive every branch without a live transport.
pub fn build_retry_envelope(
    requests: &[InputRequest],
    state: &RequestState,
    original_params: &Value,
    fixture: &InputResponseFixture,
) -> Result<Value, ElicitationError> {
    let answers = collect_responses(requests, fixture)?;
    Ok(build_retry_params(original_params, state, &answers))
}

/// Walk an elicitation response chain and return either the final
/// non-elicitation result or an [`ElicitationError`].
///
/// `next_response` is a closure the caller drives: each invocation
/// dispatches a `tools/call` (or sibling) with the supplied params
/// and returns the wrapped `{"result":...}` envelope. Keeping it a
/// closure means this module never touches the transport, so the
/// unit tests exercise the loop against a synthetic in-memory
/// server.
pub fn run_elicitation_chain<F>(
    initial_response: Value,
    initial_params: Value,
    fixture: &InputResponseFixture,
    max_rounds: u8,
    mut next_response: F,
) -> Result<Value, ElicitationError>
where
    F: FnMut(Value) -> Value,
{
    let mut response = initial_response;
    let mut params = initial_params;
    for _round in 0..max_rounds {
        let Some(irr) = recognize_input_required(&response) else {
            return Ok(response);
        };
        let retry_params =
            build_retry_envelope(&irr.input_requests, &irr.request_state, &params, fixture)?;
        response = next_response(retry_params.clone());
        params = retry_params;
    }
    if recognize_input_required(&response).is_some() {
        return Err(ElicitationError::MaxRoundsExceeded { max_rounds });
    }
    Ok(response)
}

/// Live answer source for an `InputRequiredResult` elicitation, built
/// from the parsed suite spec at plan-build time (WOR-1383). `Static`
/// answers from a pre-supplied map; `Rest` POSTs each elicitation to an
/// HTTP endpoint.
#[derive(Debug, Clone)]
pub enum Responder {
    /// Deterministic answers from the suite's `input_responses` map.
    Static(InputResponseFixture),
    /// Dynamic answers from a REST endpoint.
    Rest(RestResponder),
}

/// REST elicitation responder: POSTs each elicitation and reads back the
/// `inputResponses`. The `reqwest::Client` is built once at plan time.
#[derive(Debug, Clone)]
pub struct RestResponder {
    url: String,
    headers: BTreeMap<String, String>,
    client: reqwest::Client,
}

/// Context for one elicitation round: the call being retried plus the
/// round index. A responder may use it to answer; the runner also stamps
/// it (redacted) into the trace.
#[derive(Debug, Clone)]
pub struct ElicitCallContext {
    /// Test name (for the trace and the REST payload).
    pub test_name: String,
    /// Server name the test targets.
    pub server: String,
    /// Tool being called.
    pub tool: String,
    /// The call's `arguments` object.
    pub arguments: Value,
    /// The opaque `requestState` echoed on the retry.
    pub request_state: RequestState,
    /// 1-based elicitation round.
    pub round: u8,
}

impl Responder {
    /// Resolve answers for one elicitation round.
    pub async fn answer(
        &self,
        requests: &[InputRequest],
        ctx: &ElicitCallContext,
    ) -> Result<Vec<InputResponse>, ElicitationError> {
        match self {
            Responder::Static(fixture) => {
                let answers = collect_responses(requests, fixture)?;
                validate_answer_kinds(requests, &answers)?;
                Ok(answers)
            }
            Responder::Rest(rest) => rest.answer(requests, ctx).await,
        }
    }

    /// `static` or `rest`, for the trace `responder_kind` field.
    pub fn kind(&self) -> &'static str {
        match self {
            Responder::Static(_) => "static",
            Responder::Rest(_) => "rest",
        }
    }

    /// True when answers come from live IO (REST). Such a test is
    /// treated like one with a `transform` at the cache gate: its result
    /// is not cached, since a dynamic answer could differ next run.
    pub fn is_dynamic(&self) -> bool {
        matches!(self, Responder::Rest(_))
    }
}

impl RestResponder {
    /// Build the live responder from its spec and the run-wide proxy.
    /// Fails fast at plan-build time on a bad url / header / proxy rather
    /// than deferring the error to the first elicitation.
    pub fn from_spec(spec: &RestInputResponderSpec, proxy: &ProxyConfig) -> Result<Self, String> {
        let timeout = Duration::from_millis(spec.timeout_ms.unwrap_or(30_000));
        let builder = proxy
            .apply(reqwest::Client::builder().timeout(timeout))
            .map_err(|e| format!("input_responder proxy: {e}"))?;
        let client = builder
            .build()
            .map_err(|e| format!("input_responder client: {e}"))?;
        Ok(Self {
            url: spec.url.clone(),
            headers: spec.headers.clone(),
            client,
        })
    }

    async fn answer(
        &self,
        requests: &[InputRequest],
        ctx: &ElicitCallContext,
    ) -> Result<Vec<InputResponse>, ElicitationError> {
        let body = serde_json::json!({
            "test": ctx.test_name,
            "server": ctx.server,
            "tool": ctx.tool,
            "arguments": ctx.arguments,
            "round": ctx.round,
            "requestState": ctx.request_state.as_str(),
            "inputRequests": requests,
        });
        let mut request = self.client.post(&self.url).json(&body);
        for (name, value) in &self.headers {
            request = request.header(name, value);
        }
        let response = request.send().await.map_err(|e| {
            ElicitationError::Rest(format!("POST failed: {}", redact_send_error(&e)))
        })?;
        let status = response.status();
        if !status.is_success() {
            return Err(ElicitationError::Rest(format!(
                "endpoint returned HTTP {}",
                status.as_u16()
            )));
        }
        let parsed: Value = response
            .json()
            .await
            .map_err(|_| ElicitationError::Rest("response body was not valid JSON".to_string()))?;
        let raw = parsed
            .get("inputResponses")
            .and_then(Value::as_array)
            .ok_or_else(|| {
                ElicitationError::Rest("response is missing an `inputResponses` array".to_string())
            })?;
        let answers: Vec<InputResponse> = serde_json::from_value(Value::Array(raw.clone()))
            .map_err(|_| {
                ElicitationError::Rest(
                    "`inputResponses` entries must be `{ id, value }` objects".to_string(),
                )
            })?;
        validate_rest_answers(requests, &answers)?;
        Ok(answers)
    }
}

/// Validate a REST responder's answers: no duplicate ids, no unknown
/// ids, every required id answered, and every answer well-typed.
fn validate_rest_answers(
    requests: &[InputRequest],
    answers: &[InputResponse],
) -> Result<(), ElicitationError> {
    let mut seen = std::collections::HashSet::new();
    for answer in answers {
        if !seen.insert(answer.id.as_str()) {
            return Err(ElicitationError::Rest(format!(
                "duplicate answer for id `{}`",
                answer.id
            )));
        }
        if !requests.iter().any(|r| r.id == answer.id) {
            return Err(ElicitationError::Rest(format!(
                "answer for unknown id `{}`",
                answer.id
            )));
        }
    }
    for request in requests {
        if request.required && !answers.iter().any(|a| a.id == request.id) {
            return Err(ElicitationError::Rest(format!(
                "no answer for required id `{}`",
                request.id
            )));
        }
    }
    validate_answer_kinds(requests, answers)
}

/// Categorize a reqwest send error without leaking the URL, headers, or
/// any token into diagnostics.
fn redact_send_error(error: &reqwest::Error) -> &'static str {
    if error.is_timeout() {
        "request timed out"
    } else if error.is_connect() {
        "could not connect to the endpoint"
    } else {
        "transport error"
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use serde_json::json;

    fn req(id: &str, kind: &str, required: bool) -> InputRequest {
        InputRequest {
            id: id.into(),
            kind: kind.into(),
            description: None,
            required,
        }
    }

    #[test]
    fn collect_responses_returns_matched_answers_in_request_order() {
        let requests = vec![req("name", "string", true), req("age", "number", true)];
        let fixture = InputResponseFixture::from_map([
            ("age".to_string(), json!(42)),
            ("name".to_string(), json!("ada")),
        ]);
        let answers = collect_responses(&requests, &fixture).unwrap();
        assert_eq!(answers.len(), 2);
        assert_eq!(answers[0].id, "name");
        assert_eq!(answers[0].value, json!("ada"));
        assert_eq!(answers[1].id, "age");
        assert_eq!(answers[1].value, json!(42));
    }

    #[test]
    fn collect_responses_skips_optional_requests_that_have_no_fixture_entry() {
        let requests = vec![
            req("name", "string", true),
            req("nickname", "string", false),
        ];
        let fixture = InputResponseFixture::from_map([("name".to_string(), json!("ada"))]);
        let answers = collect_responses(&requests, &fixture).unwrap();
        assert_eq!(answers.len(), 1);
        assert_eq!(answers[0].id, "name");
    }

    #[test]
    fn collect_responses_errors_when_a_required_request_has_no_answer() {
        let requests = vec![req("name", "string", true)];
        let fixture = InputResponseFixture::from_map([("other".to_string(), json!(""))]);
        let err = collect_responses(&requests, &fixture).unwrap_err();
        assert!(
            matches!(err, ElicitationError::MissingResponse { request_id } if request_id == "name")
        );
    }

    #[test]
    fn collect_responses_errors_when_fixture_is_empty() {
        let requests = vec![req("name", "string", true)];
        let fixture = InputResponseFixture::default();
        assert!(matches!(
            collect_responses(&requests, &fixture).unwrap_err(),
            ElicitationError::NoFixture
        ));
    }

    #[test]
    fn run_elicitation_chain_returns_final_result_after_one_round_trip() {
        let initial_response = json!({
            "result": {
                "inputRequests": [{"id": "name", "kind": "string", "required": true}],
                "requestState": "tok-1"
            }
        });
        let initial_params = json!({"name": "tools/call", "arguments": {}});
        let fixture = InputResponseFixture::from_map([("name".to_string(), json!("ada"))]);
        let final_envelope = json!({"result": {"greeting": "hi ada"}});
        let mut calls = Vec::new();
        let out = run_elicitation_chain(
            initial_response,
            initial_params,
            &fixture,
            DEFAULT_MAX_ELICITATION_ROUNDS,
            |params| {
                calls.push(params.clone());
                final_envelope.clone()
            },
        )
        .expect("elicitation resolves");
        assert_eq!(out, final_envelope);
        assert_eq!(calls.len(), 1);
        // Verify the retry params carried the answer + echoed state.
        let retry = &calls[0];
        assert_eq!(retry["inputResponses"][0]["id"], json!("name"));
        assert_eq!(retry["inputResponses"][0]["value"], json!("ada"));
        assert_eq!(retry["requestState"], json!("tok-1"));
    }

    #[test]
    fn run_elicitation_chain_caps_at_max_rounds() {
        // Server always returns InputRequired; cap is 2.
        let initial = json!({
            "result": {
                "inputRequests": [{"id": "x", "kind": "string", "required": true}],
                "requestState": "tok"
            }
        });
        let fixture = InputResponseFixture::from_map([("x".to_string(), json!(""))]);
        let err =
            run_elicitation_chain(initial.clone(), json!({}), &fixture, 2, |_| initial.clone())
                .unwrap_err();
        assert!(matches!(
            err,
            ElicitationError::MaxRoundsExceeded { max_rounds: 2 }
        ));
    }

    #[test]
    fn run_elicitation_chain_passes_through_non_elicitation_response() {
        let plain = json!({"result": {"ok": true}});
        let fixture = InputResponseFixture::default();
        let out = run_elicitation_chain(plain.clone(), json!({}), &fixture, 5, |_| {
            panic!("closure should not run")
        })
        .expect("plain response passes through");
        assert_eq!(out, plain);
    }

    fn elicit_ctx() -> ElicitCallContext {
        ElicitCallContext {
            test_name: "t".into(),
            server: "s".into(),
            tool: "book_flight".into(),
            arguments: json!({}),
            request_state: RequestState("tok".into()),
            round: 1,
        }
    }

    #[test]
    fn validate_answer_kinds_rejects_a_wrong_typed_answer() {
        let requests = vec![req("dest", "string", true)];
        let answers = vec![InputResponse {
            id: "dest".into(),
            value: json!(42),
        }];
        let err = validate_answer_kinds(&requests, &answers).unwrap_err();
        assert!(
            matches!(err, ElicitationError::WrongKind { ref request_id, ref expected } if request_id == "dest" && expected == "string")
        );
    }

    #[tokio::test]
    async fn static_responder_rejects_a_wrong_typed_static_answer() {
        let requests = vec![req("dest", "string", true)];
        let responder = Responder::Static(InputResponseFixture::from_map([(
            "dest".to_string(),
            json!(42),
        )]));
        let err = responder
            .answer(&requests, &elicit_ctx())
            .await
            .unwrap_err();
        assert!(matches!(err, ElicitationError::WrongKind { .. }));
        assert_eq!(responder.kind(), "static");
        assert!(!responder.is_dynamic());
    }

    #[tokio::test]
    async fn rest_responder_resolves_answers_from_the_endpoint() {
        use wiremock::matchers::{method, path};
        use wiremock::{Mock, MockServer, ResponseTemplate};

        let server = MockServer::start().await;
        Mock::given(method("POST"))
            .and(path("/answer"))
            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
                "inputResponses": [{ "id": "dest", "value": "Denver" }]
            })))
            .mount(&server)
            .await;
        let spec = RestInputResponderSpec {
            url: format!("{}/answer", server.uri()),
            headers: Default::default(),
            timeout_ms: Some(3000),
        };
        let responder =
            Responder::Rest(RestResponder::from_spec(&spec, &ProxyConfig::default()).unwrap());
        assert_eq!(responder.kind(), "rest");
        assert!(responder.is_dynamic());
        let answers = responder
            .answer(&[req("dest", "string", true)], &elicit_ctx())
            .await
            .expect("responder answers");
        assert_eq!(answers.len(), 1);
        assert_eq!(answers[0].id, "dest");
        assert_eq!(answers[0].value, json!("Denver"));
    }

    #[tokio::test]
    async fn rest_responder_fails_on_an_http_error() {
        use wiremock::matchers::{method, path};
        use wiremock::{Mock, MockServer, ResponseTemplate};

        let server = MockServer::start().await;
        Mock::given(method("POST"))
            .and(path("/answer"))
            .respond_with(ResponseTemplate::new(500))
            .mount(&server)
            .await;
        let spec = RestInputResponderSpec {
            url: format!("{}/answer", server.uri()),
            headers: Default::default(),
            timeout_ms: Some(3000),
        };
        let responder =
            Responder::Rest(RestResponder::from_spec(&spec, &ProxyConfig::default()).unwrap());
        let err = responder
            .answer(&[req("dest", "string", true)], &elicit_ctx())
            .await
            .unwrap_err();
        assert!(matches!(err, ElicitationError::Rest(_)), "{err:?}");
    }

    #[tokio::test]
    async fn rest_responder_rejects_an_unknown_answer_id() {
        use wiremock::matchers::{method, path};
        use wiremock::{Mock, MockServer, ResponseTemplate};

        let server = MockServer::start().await;
        Mock::given(method("POST"))
            .and(path("/answer"))
            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
                "inputResponses": [{ "id": "other", "value": "x" }]
            })))
            .mount(&server)
            .await;
        let spec = RestInputResponderSpec {
            url: format!("{}/answer", server.uri()),
            headers: Default::default(),
            timeout_ms: Some(3000),
        };
        let responder =
            Responder::Rest(RestResponder::from_spec(&spec, &ProxyConfig::default()).unwrap());
        let err = responder
            .answer(&[req("dest", "string", true)], &elicit_ctx())
            .await
            .unwrap_err();
        match err {
            ElicitationError::Rest(reason) => assert!(reason.contains("unknown id"), "{reason}"),
            other => panic!("expected a Rest error, got {other:?}"),
        }
    }
}
