use serde::Deserialize;
use std::collections::HashMap;

pub use crate::pre_serve::{
    BeforeServeTable, DistributedBackend, ModsTable, OnFailure, PreServeCommand, RecipeArgs,
};

#[derive(Clone, Debug, Deserialize)]
pub struct Node {
    pub name: String,
    pub ssh_ip: String,
    pub ib_ip: String,
    pub rank: u32,
    #[serde(default)]
    pub host_key: Option<String>,
    #[serde(default)]
    pub entity: Option<String>,
}

#[derive(Clone, Debug, Deserialize)]
pub struct NcclConfig {
    pub ib_hca: String,
    pub socket_ifname: String,
    pub ib_gid_index: u32,
    pub ib_disable: u32,
    #[serde(default)]
    pub gloo_socket_ifname: Option<String>,
    #[serde(default)]
    pub ucx_net_devices: Option<String>,
    #[serde(default)]
    pub nccl_cross_nic: Option<u32>,
    #[serde(default)]
    pub nccl_proto: Option<String>,
    #[serde(default)]
    pub nccl_p2p_level: Option<String>,
}

#[derive(Clone, Debug, Deserialize)]
pub struct Cluster {
    pub name: String,
    pub master_ib_ip: String,
    pub nodes: Vec<Node>,
    #[serde(default)]
    pub extra_entities: Vec<String>,
}

#[derive(Clone, Deserialize)]
pub struct Config {
    pub ssh_user: String,
    pub ssh_key_path: std::path::PathBuf,
    pub clusters: Vec<Cluster>,
    #[serde(default)]
    pub nccl: HashMap<String, NcclConfig>,
    #[serde(default)]
    pub hf_cache_dir: Option<String>,
    #[serde(default = "default_init_port")]
    pub init_port: u16,
    #[serde(default)]
    pub before_serve: HashMap<String, BeforeServeTable>,
    #[serde(default)]
    pub exec_aliases: HashMap<String, String>,
    #[serde(default)]
    pub home_assistant: Option<HomeAssistantConfig>,
    #[serde(default)]
    pub global_mods_dir: Option<std::path::PathBuf>,
}

#[derive(Clone, Deserialize)]
pub struct HomeAssistantConfig {
    pub url: String,
    pub token: String,
    #[serde(default = "default_shutdown_delay")]
    pub shutdown_delay_seconds: u64,
}

fn default_shutdown_delay() -> u64 {
    30
}
fn default_init_port() -> u16 {
    25000
}
fn default_port() -> u16 {
    30000
}
fn default_engine() -> String {
    "vllm".to_string()
}
fn default_shm_size() -> String {
    String::new()
}
fn default_tensor_parallel_size() -> u32 {
    1
}
fn default_sync_image() -> bool {
    true
}

#[derive(Clone, Deserialize)]
pub struct Recipe {
    #[serde(default = "default_engine")]
    pub engine: String,
    pub name: String,
    #[serde(default)]
    pub docker_image: String,
    #[serde(default)]
    pub model: String,
    #[serde(default = "default_port")]
    pub port: u16,
    #[serde(default)]
    pub required_nodes: Option<u32>,
    #[serde(default)]
    pub pipeline_parallel_size: Option<u32>,
    #[serde(default = "default_tensor_parallel_size")]
    pub tensor_parallel_size: u32,
    #[serde(default)]
    pub args: RecipeArgs,
    #[serde(default)]
    pub env: HashMap<String, String>,
    #[serde(default)]
    pub cluster: Option<String>,
    #[serde(default)]
    pub entrypoint: Vec<String>,
    pub distributed_backend: Option<DistributedBackend>,
    #[serde(default = "default_shm_size")]
    pub shm_size: String,
    #[serde(default)]
    pub before_serve: Vec<PreServeCommand>,
    #[serde(default)]
    pub mods: ModsTable,
    #[serde(default = "default_sync_image")]
    pub sync_image: bool,
}

impl Config {
    pub fn resolved_hf_cache_dir(&self) -> std::path::PathBuf {
        if let Some(ref dir) = self.hf_cache_dir {
            return std::path::PathBuf::from(dir);
        }
        let home = if self.ssh_user == "root" {
            "/root"
        } else {
            return std::path::PathBuf::from(format!("/home/{}/.cache/huggingface", self.ssh_user));
        };
        std::path::PathBuf::from(format!("{}/.cache/huggingface", home))
    }
}