use crate::config::*;
use crate::pre_serve::*;
use std::collections::{HashMap, HashSet};

fn tcls() -> Cluster {
    Cluster {
        name: "test-cluster".to_string(),
        master_ib_ip: Some("192.168.1.1".to_string()),
        nodes: vec![
            Node {
                name: "node0".to_string(),
                ssh_ip: "10.0.0.1".to_string(),
                ib_ip: Some("192.168.1.1".to_string()),
                rank: 0,
                host_key: None,
                entity: None,
            },
            Node {
                name: "node1".to_string(),
                ssh_ip: "10.0.0.2".to_string(),
                ib_ip: Some("192.168.1.2".to_string()),
                rank: 1,
                host_key: None,
                entity: None,
            },
        ],
        extra_entities: vec![],
    }
}

fn tcfg() -> Config {
    Config {
        ssh_user: "u".to_string(),
        ssh_key_path: "/k".into(),
        clusters: vec![tcls()],
        nccl: HashMap::new(),
        hf_cache_dir: None,
        init_port: 25000,
        before_serve: HashMap::new(),
        exec_aliases: HashMap::new(),
        home_assistant: None,
        global_mods_dir: None,
        drop_caches: true,
        mount_aliases: HashMap::new(),
        recipe_sources: Vec::new(),
        docker_image_sources: Vec::new(),
        cluster_sources: Vec::new(),
    }
}

fn tnccl() -> NcclConfig {
    NcclConfig {
        ib_hca: "rocep1s0f0".to_string(),
        socket_ifname: "enp1s0f0np0".to_string(),
        ib_gid_index: Some(3),
        ib_disable: 0,
        gloo_socket_ifname: None,
        ucx_net_devices: None,
        nccl_cross_nic: None,
        nccl_proto: None,
        nccl_p2p_level: None,
    }
}

#[test]
fn test_recipe_deserialize_minimal() {
    let r: Recipe = toml::from_str("\nname = \"M\"\nmodel = \"t/m\"\n").unwrap();
    assert_eq!(r.name, "M");
    assert_eq!(r.engine, "vllm");
    assert_eq!(r.port, 30000);
}

#[test]
fn test_recipe_deserialize_full() {
    let r: Recipe = toml::from_str("engine = \"sglang\"\nname = \"Full\"\ndocker_image = \"o/i:t\"\nmodel = \"o/m\"\nport = 30000\nrequired_nodes = 4\npipeline_parallel_size = 2\ndistributed_backend = \"ray\"\ncluster = \"c\"\nargs = [\"--a\", \"v\"]\n[env]\nK = \"V\"\n").unwrap();
    assert_eq!(r.engine, "sglang");
    assert_eq!(r.required_nodes, Some(4));
    assert_eq!(r.distributed_backend, Some(DistributedBackend::Ray));
}

#[test]
fn test_recipe_deserialize_no_required_nodes() {
    let r: Recipe = toml::from_str("name = \"F\"\nmodel = \"t/m\"\n").unwrap();
    assert_eq!(r.required_nodes, None);
    assert_eq!(r.pipeline_parallel_size, None);
}

#[test]
fn test_composed_source_exclude_deserialize() {
    let toml = r#"
name = "C"
model = "t/m"
[composed_model]
mount = "/root/model"
[[composed_model.sources]]
repo = "base"
[[composed_model.sources]]
repo = "overlay"
glob = "L*.safetensors"
exclude = ["L5.safetensors", "L8.safetensors"]
[[composed_model.files]]
name = "cfg.json"
content = "{}"
"#;
    let r: Recipe = toml::from_str(toml).unwrap();
    let composed = r.composed_model.unwrap();
    assert_eq!(composed.sources.len(), 2);
    let overlay = &composed.sources[1];
    assert_eq!(overlay.repo, "overlay");
    assert_eq!(overlay.glob.as_deref(), Some("L*.safetensors"));
    let exclude = overlay
        .exclude
        .as_ref()
        .expect("exclude should deserialize");
    assert_eq!(exclude.len(), 2);
    assert!(exclude.contains(&"L5.safetensors".to_string()));
    assert!(exclude.contains(&"L8.safetensors".to_string()));
}

#[test]
fn test_config_deserialize() {
    let c: Config = toml::from_str("ssh_user = \"a\"\nssh_key_path = \"/p\"\n[[clusters]]\nname = \"p\"\nmaster_ib_ip = \"1\"\nnodes = [{name=\"n1\",ssh_ip=\"10.0.0.10\",ib_ip=\"1\",rank=0},{name=\"n2\",ssh_ip=\"10.0.0.11\",ib_ip=\"2\",rank=1}]\n[nccl.\"p\"]\nib_hca = \"m\"\nsocket_ifname = \"e\"\nib_gid_index = 3\nib_disable = 0\n").unwrap();
    assert_eq!(c.clusters.len(), 1);
    assert_eq!(c.clusters[0].nodes.len(), 2);
}

#[test]
fn test_config_deserialize_with_host_key() {
    let c: Config = toml::from_str("ssh_user = \"a\"\nssh_key_path = \"/k\"\n[[clusters]]\nname = \"s\"\nmaster_ib_ip = \"1\"\nnodes = [\n    { name = \"n\", ssh_ip = \"2\", ib_ip = \"1\", rank = 0, host_key = \"SHA256:abc\" },\n]\n").unwrap();
    assert_eq!(c.clusters[0].nodes[0].host_key, Some("SHA256:abc".into()));
}

#[test]
fn test_select_cluster_by_name() {
    let mut cfg = tcfg();
    cfg.nccl.insert("test-cluster".into(), tnccl());
    let (c, n) = select_cluster(&cfg, "test-cluster", None, &HashSet::new()).unwrap();
    assert_eq!(c.name, "test-cluster");
    assert_eq!(n.unwrap().ib_hca, "rocep1s0f0");
}

#[test]
fn test_select_cluster_auto() {
    let mut cfg = tcfg();
    cfg.nccl.insert("test-cluster".into(), tnccl());
    let (c, _n) = select_cluster(&cfg, "auto", None, &HashSet::new()).unwrap();
    assert_eq!(c.name, "test-cluster");
}

#[test]
fn test_select_cluster_auto_no_nccl_allows_single_node() {
    // With required_nodes = None, NCCL config is not required — the cluster
    // is selected even without NCCL config (the node count is unknown at
    // selection time, so we can't require it).
    let cfg = tcfg();
    let (c, n) = select_cluster(&cfg, "auto", None, &HashSet::new()).unwrap();
    assert_eq!(c.name, "test-cluster");
    assert!(n.is_none());
}

#[test]
fn test_select_cluster_multi_node_requires_nccl() {
    // Multi-node (required_nodes > 1) without NCCL config → error.
    assert!(select_cluster(&tcfg(), "auto", Some(2), &HashSet::new()).is_err());
}

#[test]
fn test_select_cluster_not_found() {
    let mut cfg = tcfg();
    cfg.clusters = vec![];
    assert!(select_cluster(&cfg, "x", None, &HashSet::new()).is_err());
}

#[test]
fn test_select_cluster_by_name_found() {
    let mut nccl = HashMap::new();
    nccl.insert("test-cluster".into(), tnccl());
    let mut cfg = tcfg();
    cfg.nccl = nccl;
    let c = select_cluster_by_name(&cfg, "test-cluster").unwrap();
    assert_eq!(c.name, "test-cluster");
}

#[test]
fn test_select_cluster_by_name_not_found() {
    assert!(select_cluster_by_name(&tcfg(), "x").is_err());
}

#[test]
fn test_select_cluster_by_prefix() {
    let mut cfg = tcfg();
    cfg.nccl.insert("test-cluster".into(), tnccl());
    let c = select_cluster_by_name(&cfg, "test-c").unwrap();
    assert_eq!(c.name, "test-cluster");
    assert!(select_cluster_by_name(&cfg, "no-such").is_err());
}

#[test]
fn test_default_engine_is_vllm() {
    let r: Recipe = toml::from_str("name = \"x\"\nmodel = \"m\"\n").unwrap();
    assert_eq!(r.engine, "vllm");
}

#[test]
fn test_default_distributed_backend_is_none() {
    let r: Recipe = toml::from_str("name = \"x\"\nmodel = \"m\"\n").unwrap();
    assert_eq!(r.distributed_backend, None);
}

#[test]
fn test_deploy_defaults_to_docker() {
    let r: Recipe = toml::from_str("name = \"x\"\nmodel = \"m\"\n").unwrap();
    assert_eq!(r.deploy, DeployTarget::Docker);
}

#[test]
fn test_deploy_explicit_tmux() {
    let r: Recipe = toml::from_str("name = \"x\"\nmodel = \"m\"\ndeploy = \"tmux\"\n").unwrap();
    assert_eq!(r.deploy, DeployTarget::Tmux);
}

#[test]
fn test_deploy_explicit_docker() {
    let r: Recipe = toml::from_str("name = \"x\"\nmodel = \"m\"\ndeploy = \"docker\"\n").unwrap();
    assert_eq!(r.deploy, DeployTarget::Docker);
}

#[test]
fn test_venv_field() {
    let r: Recipe = toml::from_str(
        "name = \"x\"\nmodel = \"m\"\ndeploy = \"tmux\"\nvenv = \"~/venvs/vllm025\"\n",
    )
    .unwrap();
    assert_eq!(r.deploy, DeployTarget::Tmux);
    assert_eq!(r.venv.as_deref(), Some("~/venvs/vllm025"));
}

#[test]
fn test_venv_defaults_to_none() {
    let r: Recipe = toml::from_str("name = \"x\"\nmodel = \"m\"\n").unwrap();
    assert_eq!(r.venv, None);
}

#[test]
fn test_venv_setup_list_deserializes() {
    let r: Recipe = toml::from_str(
        r#"name = "x"
model = "m"
deploy = "tmux"
venv = "~/venvs/vllm025"
venv_setup = [
    "uv venv ~/venvs/vllm025 -p 3.12",
    "uv pip install vllm==0.25.1 --torch-backend=cu130",
]
"#,
    )
    .unwrap();
    assert_eq!(r.venv.as_deref(), Some("~/venvs/vllm025"));
    let setup = r.venv_setup.to_vec();
    assert_eq!(setup.len(), 2);
    assert_eq!(setup[0], "uv venv ~/venvs/vllm025 -p 3.12");
    assert_eq!(
        setup[1],
        "uv pip install vllm==0.25.1 --torch-backend=cu130"
    );
}

#[test]
fn test_venv_setup_inline_deserializes() {
    let r: Recipe = toml::from_str(
        r#"name = "x"
model = "m"
deploy = "tmux"
venv = "~/venvs/vllm025"
venv_setup = """
    uv venv ~/venvs/vllm025 -p 3.12
    uv pip install vllm==0.25.1 --torch-backend=cu130
"""
"#,
    )
    .unwrap();
    let setup = r.venv_setup.to_vec();
    assert_eq!(setup.len(), 2);
    assert_eq!(setup[0], "uv venv ~/venvs/vllm025 -p 3.12");
    assert_eq!(
        setup[1],
        "uv pip install vllm==0.25.1 --torch-backend=cu130"
    );
}

#[test]
fn test_venv_setup_inline_skips_comments_and_blanks() {
    let r: Recipe = toml::from_str(
        r#"name = "x"
model = "m"
deploy = "tmux"
venv = "~/venvs/vllm025"
venv_setup = """
    # create the venv
    uv venv ~/venvs/vllm025 -p 3.12

    # install deps
    uv pip install vllm==0.25.1
"""
"#,
    )
    .unwrap();
    let setup = r.venv_setup.to_vec();
    assert_eq!(setup.len(), 2);
    assert_eq!(setup[0], "uv venv ~/venvs/vllm025 -p 3.12");
    assert_eq!(setup[1], "uv pip install vllm==0.25.1");
}

#[test]
fn test_venv_setup_defaults_to_empty() {
    let r: Recipe = toml::from_str("name = \"x\"\nmodel = \"m\"\n").unwrap();
    assert!(r.venv_setup.to_vec().is_empty());
}

#[test]
fn default_config_dir_appends_sm121() {
    let Some(d) = default_config_dir() else {
        return;
    };
    assert!(d.ends_with("sm121") && d.parent().is_some());
}

#[tokio::test]
async fn test_merge_cluster_sources_adds_new_keeps_main() {
    let dir = std::env::temp_dir().join(format!("sm121-csrc-{}", std::process::id()));
    std::fs::create_dir_all(&dir).unwrap();
    let src_path = dir.join("clusters-extra.toml");
    // "extra" is new; "test-cluster" collides with tcfg's cluster (main wins).
    std::fs::write(
        &src_path,
        r#"
ssh_user = "srcuser"
ssh_key_path = "/src/key"

[[clusters]]
name = "extra"
nodes = [{ name = "e0", ssh_ip = "10.9.0.1", rank = 0 }]

[[clusters]]
name = "test-cluster"
nodes = [{ name = "intruder", ssh_ip = "10.66.0.1", rank = 0 }]

[nccl."extra"]
ib_hca = "mlx5_0"
socket_ifname = "eth0"
ib_disable = 0

[before_serve."extra"]
commands = ["echo hi"]

[home_assistant]
url = "http://ha.lan:8123"
token = "sufficiently-long-test-token"
shutdown_delay_seconds = 30
"#,
    )
    .unwrap();
    let mut cfg = tcfg();
    cfg.cluster_sources = vec![src_path.to_string_lossy().into_owned()];
    cfg.ssh_user = String::new();
    cfg.ssh_key_path = std::path::PathBuf::new();
    merge_cluster_sources(&mut cfg).await.unwrap();
    assert_eq!(cfg.ssh_user, "srcuser", "empty ssh_user filled from source");
    assert_eq!(cfg.ssh_key_path, std::path::PathBuf::from("/src/key"));
    assert_eq!(
        cfg.clusters.len(),
        2,
        "extra added, conflicting name not duplicated"
    );
    let main = cfg
        .clusters
        .iter()
        .find(|c| c.name == "test-cluster")
        .unwrap();
    assert_eq!(
        main.nodes[0].name, "node0",
        "main-config cluster wins on conflict"
    );
    assert!(cfg.clusters.iter().any(|c| c.name == "extra"));
    assert!(cfg.nccl.contains_key("extra"));
    assert!(cfg.before_serve.contains_key("extra"));
    // HA config merges from the source when main has none (power on from
    // other cwds depends on this).
    assert!(
        cfg.home_assistant.is_some(),
        "home_assistant filled from source"
    );
    assert_eq!(
        cfg.home_assistant.as_ref().unwrap().url,
        "http://ha.lan:8123"
    );
    // A main-config ssh_user is never overridden by a source; same rule
    // for home_assistant.
    let mut cfg2 = tcfg();
    cfg2.cluster_sources = vec![src_path.to_string_lossy().into_owned()];
    cfg2.home_assistant = Some(crate::config::HomeAssistantConfig {
        url: "http://main-ha:8123".into(),
        token: "main".into(),
        shutdown_delay_seconds: 5,
    });
    merge_cluster_sources(&mut cfg2).await.unwrap();
    assert_eq!(cfg2.ssh_user, "u", "main ssh_user wins over source");
    assert_eq!(
        cfg2.home_assistant.as_ref().unwrap().url,
        "http://main-ha:8123",
        "main home_assistant wins over source"
    );
}

#[tokio::test]
async fn test_load_recipe_from_recipe_sources() {
    let dir = std::env::temp_dir().join(format!("sm121-rsrc-{}", std::process::id()));
    let rdir = dir.join("remote-recipes").join("remote-recipe");
    std::fs::create_dir_all(&rdir).unwrap();
    std::fs::write(rdir.join("recipe.toml"), "name = \"remote\"\n").unwrap();
    let recipe = load_recipe(
        "remote-recipe",
        &[dir.join("remote-recipes").to_string_lossy().into_owned()],
    )
    .await
    .unwrap();
    assert_eq!(recipe.name, "remote");
}

#[test]
fn effective_master_ip_falls_back_to_rank0_ib_ip() {
    let mut c = tcls();
    c.master_ib_ip = None;
    assert_eq!(c.effective_master_ip(), "192.168.1.1"); // rank-0 ib_ip
}

#[test]
fn effective_master_ip_falls_back_to_ssh_ip_when_ib_missing() {
    let mut c = tcls();
    c.master_ib_ip = None;
    for n in &mut c.nodes {
        n.ib_ip = None;
    }
    assert_eq!(c.effective_master_ip(), "10.0.0.1"); // rank-0 ssh_ip
}

#[test]
fn effective_master_ip_uses_first_node_without_rank0() {
    let mut c = tcls();
    c.master_ib_ip = None;
    for n in &mut c.nodes {
        n.rank = 3; // no rank 0 → first node wins
    }
    assert_eq!(c.effective_master_ip(), "192.168.1.1");
}

#[test]
fn effective_master_ip_prefers_explicit_config() {
    assert_eq!(tcls().effective_master_ip(), "192.168.1.1");
}