//! Qwen server defaults and the native application's execution settings.
//! Overrides are model-owned values; never mutate the process environment.
use super::session_cache::{BoundaryPolicy, DecodeOptions, TurnOptions};
use super::{ar, prefill, sampling::SamplerConfig, speculative};
use crate::settings::{EngineSettings, QWEN_DEFAULT_PREFILL_CHUNK};
use std::collections::BTreeMap;

pub(super) type Overrides = BTreeMap<&'static str, &'static str>;

pub(super) fn truthy(value: Option<&str>) -> bool {
    value.is_some_and(|s| {
        matches!(
            s.trim().to_ascii_lowercase().as_str(),
            "1" | "true" | "yes" | "on"
        )
    })
}

fn falsey(value: Option<&str>) -> bool {
    value.is_some_and(|s| {
        matches!(
            s.trim().to_ascii_lowercase().as_str(),
            "0" | "false" | "no" | "off"
        )
    })
}

pub(super) fn sdpa_head_chunk_enabled(get: impl Fn(&'static str) -> Option<String>) -> bool {
    !falsey(get("MTPLX_QWEN4_VERIFY_SDPA_HEAD_CHUNK").as_deref())
}

pub(super) fn sorted_gather_enabled(get: impl Fn(&'static str) -> Option<String>) -> bool {
    !falsey(get("MTPLX_MOE_SORTED_GATHER_KERNEL").as_deref())
        && !truthy(get("MTPLX_FORCE_GPU_FAMILY_FALLBACK").as_deref())
}

fn strict_bool(name: &str, raw: Option<String>) -> Result<bool, String> {
    match raw
        .as_deref()
        .unwrap_or_default()
        .trim()
        .to_ascii_lowercase()
        .as_str()
    {
        "" | "0" | "false" | "no" | "off" | "disable" | "disabled" => Ok(false),
        "1" | "true" | "yes" | "on" | "enable" | "enabled" => Ok(true),
        _ => Err(format!("{name}={raw:?} is not a boolean")),
    }
}

pub(super) fn value(overrides: &Overrides, name: &str) -> Option<String> {
    overrides
        .get(name)
        .map(|s| (*s).to_owned())
        .or_else(|| std::env::var(name).ok())
}

/// _server_runtime_env_overrides for this family, serial AR/native batched MTP.
pub(super) fn server_overrides(mtp: bool, get: impl Fn(&str) -> Option<String>) -> Overrides {
    let mut result = Overrides::new();
    if mtp {
        result.insert("MTPLX_SKIP_VERIFY_SNAPSHOT", "0");
    }
    for key in [
        "MTPLX_FAMILY_CAPTURE_COMMIT",
        "MTPLX_FUSED_HC_V3",
        "MTPLX_FUSED_GDN_INPROJ",
        "MTPLX_FUSED_GATE_UP",
        "MTPLX_FUSED_GDN_CONVNORM",
        "MTPLX_FUSED_GDN_STEP",
        "MTPLX_FUSED_CONVNORM_VERIFY",
        "MTPLX_QSA_GATHER",
    ] {
        if get(key).is_none() {
            result.insert(key, "1");
        }
    }
    let compiled = ["MTPLX_COMPILED_GDN", "MTPLX_QWEN4EXP_COMPILE"];
    if !compiled
        .iter()
        .any(|key| get(key).is_some_and(|s| !truthy(Some(&s))))
    {
        for key in compiled {
            if get(key).is_none() {
                result.insert(key, "1");
            }
        }
    }
    if get("MTPLX_NAX_VERIFY").is_none() {
        result.insert("MTPLX_NAX_VERIFY", "0");
    }
    result
}

pub(super) fn integer(raw: Option<&str>, default: usize, minimum: usize) -> Result<usize, String> {
    let number = raw
        .and_then(|s| s.trim().replace('_', "").parse::<i128>().ok())
        .unwrap_or(default as i128)
        .max(minimum as i128);
    usize::try_from(number).map_err(|_| "MTPLX setting exceeds the supported integer range".into())
}

/// Application chunk size is the native request-local override, not an env
/// write. The bounded/sustained route keeps only the final prefill logits.
pub(super) fn turn_options(
    settings: &EngineSettings,
    get: impl Fn(&str) -> Option<String>,
) -> Result<(TurnOptions<'static>, Option<BoundaryPolicy>), String> {
    if !settings.model.is_qwen38() || settings.context_tokens <= 0 {
        return Err("canonical Qwen requires Qwen settings and a positive context".into());
    }
    let flag = |key| truthy(get(key).as_deref());
    let enabled = |key| !falsey(get(key).as_deref());
    let number = |key, default, minimum| integer(get(key).as_deref(), default, minimum);
    let mtp = settings.speculative.glm_mtp;
    let chunk = if settings.execution.prefill_chunk == 0 {
        QWEN_DEFAULT_PREFILL_CHUNK
    } else {
        settings.execution.prefill_chunk
    };
    if chunk > i32::MAX as u32 {
        return Err("Qwen prefill chunk exceeds int32".into());
    }
    let decode = if mtp {
        DecodeOptions::Mtp(speculative::Options {
            prompt_ids: &[],
            depth: 3,
            max_tokens: usize::MAX,
            seed: 0,
            stops: &[],
            committed_history: true,
            family_capture: strict_bool(
                "MTPLX_FAMILY_CAPTURE_COMMIT",
                get("MTPLX_FAMILY_CAPTURE_COMMIT"),
            )?,
            lazy_history: flag("MTPLX_LAZY_MTP_HISTORY_APPEND"),
            defer_repair_eval: strict_bool(
                "MTPLX_DEFER_REPAIR_EVAL",
                get("MTPLX_DEFER_REPAIR_EVAL"),
            )?,
            capture_final_state: true,
        })
    } else {
        DecodeOptions::Ar(ar::Options {
            max_tokens: usize::MAX,
            seed: 0,
            stop_tokens: &[],
            sync_eval: !flag("MTPLX_ASYNC_AR")
                || get("MTPLX_EVAL_AUDIT").is_some_and(|s| !s.is_empty()),
            return_hidden: false,
            capture_final_state: true,
        })
    };
    let boundaries = if enabled("MTPLX_GDN_BOUNDARY_CAPTURE") {
        Some(BoundaryPolicy {
            tail_interval: number("MTPLX_GDN_BOUNDARY_TAIL_INTERVAL", 256, 0)?,
            max_count: number("MTPLX_GDN_BOUNDARY_MAX", 8, 2)?,
        })
    } else {
        None
    };
    Ok((
        TurnOptions {
            prefill: prefill::Options {
                chunk_size: Some(chunk as usize),
                final_logits_only: true,
                external_cache_only: flag("MTPLX_PREFILL_OMLX_EXTERNAL")
                    || (flag("MTPLX_PREFILL_STOCK_CACHE_ONLY")
                        && flag("MTPLX_ALLOW_UNSAFE_PREFILL_STOCK_CACHE_ONLY")),
                external_emit_logits: enabled("MTPLX_PREFILL_EXTERNAL_EMIT_LOGITS"),
                return_hidden: mtp,
            },
            decode,
            sampler: SamplerConfig::default(),
            draft_sampler: SamplerConfig::default(),
            lazy_kv: enabled("MTPLX_SESSION_LAZY_SNAPSHOT"),
            store_min_suffix: if enabled("MTPLX_SESSION_STORE_ON_PREFILL") {
                Some(number(
                    "MTPLX_SESSION_STORE_ON_PREFILL_MIN_SUFFIX",
                    1024,
                    1,
                )?)
            } else {
                None
            },
            fused_max: number("MTPLX_SMALL_SUFFIX_FUSED_MAX", 512, 0)?,
        },
        boundaries,
    ))
}

#[test]
fn retired_ar_pipeline_flag_cannot_change_server_or_turn_policy() {
    let mut settings = crate::settings::effective_settings(
        crate::model::ModelChoice::Qwen38FlashNext,
        &crate::settings::GenerationPreferences::default(),
        &crate::settings::RuntimePreferences::default(),
        std::path::Path::new("/unused-policy-models"),
    )
    .unwrap()
    .engine;
    settings.speculative.glm_mtp = false;
    for legacy in [None, Some("0"), Some("1")] {
        let lookup = |key: &str| {
            (key == "MTPLX_AR_PIPELINE")
                .then_some(legacy)
                .flatten()
                .map(str::to_owned)
        };
        let overrides = server_overrides(false, lookup);
        assert_eq!(overrides, server_overrides(false, |_| None));
        let (options, _) = turn_options(&settings, |key| {
            assert_ne!(key, "MTPLX_AR_PIPELINE", "retired option was read");
            overrides
                .get(key)
                .map(|s| (*s).to_owned())
                .or_else(|| lookup(key))
        })
        .unwrap();
        let DecodeOptions::Ar(ar) = options.decode else {
            panic!("AR request changed modes");
        };
        assert!(ar.sync_eval && ar.capture_final_state);
        assert!(!ar.return_hidden);
    }
}

#[test]
#[ignore = "requires DS4_QWEN38_POLICY_REFERENCE from pinned MTPLX; CPU-only"]
fn qwen_server_policy_matches_original_functions() {
    let path = std::env::var_os("DS4_QWEN38_POLICY_REFERENCE").unwrap();
    let contents = std::fs::read_to_string(path).unwrap();
    let mut count = 0;
    for case in contents
        .lines()
        .filter_map(|line| serde_json::from_str::<serde_json::Value>(line).ok())
    {
        if case["event"] != "policy_reference" {
            continue;
        }
        let mtp = case["mtp"].as_bool().unwrap();
        let lookup = |key: &str| case["environment"][key].as_str().map(str::to_owned);
        let overrides = server_overrides(mtp, lookup);
        if case["server_only"] == true {
            // This scoped comparison owns AR retirement and the existing
            // defaults. New upstream features are outside #104; retain the
            // complete legacy policy comparison below for its own oracle.
            assert!(
                !case["overrides"]
                    .as_object()
                    .unwrap()
                    .contains_key("MTPLX_AR_PIPELINE")
            );
            assert!(!overrides.contains_key("MTPLX_AR_PIPELINE"));
            for (key, value) in &overrides {
                assert_eq!(case["overrides"][key], *value, "default {key} case {count}");
            }
            count += 1;
            continue;
        }
        assert_eq!(
            serde_json::to_value(&overrides).unwrap(),
            case["overrides"],
            "overrides case {count}"
        );
        let get = |key: &str| {
            overrides
                .get(key)
                .map(|s| (*s).to_owned())
                .or_else(|| lookup(key))
        };
        let mut settings = crate::settings::effective_settings(
            crate::model::ModelChoice::Qwen38FlashNext,
            &crate::settings::GenerationPreferences::default(),
            &crate::settings::RuntimePreferences::default(),
            std::path::Path::new("/unused-policy-models"),
        )
        .unwrap()
        .engine;
        settings.speculative.glm_mtp = mtp;
        let resolved = turn_options(&settings, get);
        if case["error"].is_string() {
            assert!(
                resolved.is_err(),
                "invalid configuration accepted in case {count}"
            );
            count += 1;
            continue;
        }
        let (actual, boundary) = resolved.unwrap();
        let expected = &case["options"];
        assert_eq!(actual.prefill.chunk_size.unwrap(), expected["chunk"]);
        assert_eq!(actual.prefill.final_logits_only, expected["final_only"]);
        assert_eq!(
            actual.prefill.external_cache_only,
            expected["external_cache_only"]
        );
        assert_eq!(
            actual.prefill.external_emit_logits,
            expected["external_emit_logits"]
        );
        assert_eq!(actual.lazy_kv, expected["lazy_kv"]);
        assert_eq!(actual.store_min_suffix.is_some(), expected["store"]);
        if let Some(min) = actual.store_min_suffix {
            assert_eq!(min, expected["store_min"]);
        }
        assert_eq!(actual.fused_max, expected["fused_max"]);
        assert_eq!(boundary.is_some(), expected["boundaries"]);
        if let Some(b) = boundary {
            assert_eq!(b.max_count, expected["boundary_max"]);
            assert_eq!(b.tail_interval, expected["boundary_tail"]);
        }
        match actual.decode {
            DecodeOptions::Ar(o) => {
                assert!(!mtp);
                assert_eq!(o.sync_eval, expected["sync_eval"]);
            }
            DecodeOptions::Mtp(o) => {
                assert!(mtp && o.committed_history && o.capture_final_state);
                assert_eq!(o.depth, 3);
                assert_eq!(o.family_capture, expected["capture"]);
                assert_eq!(o.lazy_history, expected["lazy_history"]);
                assert_eq!(o.defer_repair_eval, expected["defer_repair"]);
            }
        }
        settings.execution.prefill_chunk = 77;
        assert_eq!(
            turn_options(&settings, get).unwrap().0.prefill.chunk_size,
            Some(77)
        );
        count += 1;
    }
    assert_eq!(count, 268);
}

#[test]
fn sorted_moe_switch_and_family_rehearsal() {
    assert!(sorted_gather_enabled(|_| None));
    for off in ["0", " false ", "NO", "off"] {
        assert!(!sorted_gather_enabled(|key| (key
            == "MTPLX_MOE_SORTED_GATHER_KERNEL")
            .then(|| off.to_owned())));
    }
    for force in ["1", " true ", "YES", "on"] {
        assert!(!sorted_gather_enabled(|key| (key
            == "MTPLX_FORCE_GPU_FAMILY_FALLBACK")
            .then(|| force.to_owned())));
    }
    assert!(sorted_gather_enabled(|_| Some("unexpected".into())));
}