use std::{path::PathBuf, sync::Arc};

use anyhow::{bail, Context, Result};
use chrono::Utc;
use dashmap::DashMap;
use serde::{Deserialize, Serialize};
use tokio::io::AsyncWriteExt as _;
use tracing::info;
use uuid::Uuid;

use super::types::{P95Stats, RunnerLease, RunnerSpec, SpawnRequest, SpawnResponse};

pub type BuildMetrics = Arc<DashMap<String, (f64, f64, f64, u64)>>;

/// Environment-driven configuration for the Nu runner script backend.
///
/// All fields map to env vars; `RUNNER_SCRIPT` is the only required one —
/// its absence causes `try_new()` to return `None` (pool disabled).
#[derive(Clone)]
struct RunnerConfig {
    /// Absolute path to the `runner.nu` script (e.g. `catalog/providers/hetzner/runner.nu`).
    script: PathBuf,
    /// Nu binary name or path. Defaults to `nu`.
    nu_bin: String,
    /// Provider location passed as `location` in the spawn JSON.
    location: String,
    /// SSH key name registered with the provider.
    ssh_key_ref: String,
    /// Default image ref — label selector or direct ID. Overridden per-request by `SpawnRequest.image`.
    image_ref: String,
    /// Optional private network name.
    network_ref: Option<String>,
    /// Optional firewall name.
    firewall_ref: Option<String>,
}

impl RunnerConfig {
    fn from_env() -> Option<Self> {
        let script = std::env::var("RUNNER_SCRIPT").ok().map(PathBuf::from)?;
        Some(Self {
            script,
            nu_bin: std::env::var("NU_BIN").unwrap_or_else(|_| "nu".into()),
            location: std::env::var("RUNNER_LOCATION").unwrap_or_else(|_| "fsn1".into()),
            ssh_key_ref: std::env::var("RUNNER_SSH_KEY_REF")
                .unwrap_or_else(|_| "orchestrator-buildkit".into()),
            image_ref: std::env::var("RUNNER_IMAGE_REF")
                .unwrap_or_else(|_| "app=buildkit-runner-golden".into()),
            network_ref: std::env::var("RUNNER_NETWORK_REF").ok(),
            firewall_ref: std::env::var("RUNNER_FIREWALL_REF").ok(),
        })
    }
}

#[derive(Serialize)]
struct SpawnInput<'a> {
    cpu: u32,
    memory_gb: u32,
    disk_gb: u32,
    time_budget_min: u32,
    image_ref: &'a str,
    location: &'a str,
    ssh_key_ref: &'a str,
    #[serde(skip_serializing_if = "Option::is_none")]
    network_ref: Option<&'a str>,
    #[serde(skip_serializing_if = "Option::is_none")]
    firewall_ref: Option<&'a str>,
}

#[derive(Deserialize)]
struct SpawnOutput {
    handle: String,
    ssh_host: String,
    ssh_port: u16,
    expires_at: Option<String>,
}

#[derive(Clone)]
pub struct VmPool {
    config: RunnerConfig,
    leases: Arc<DashMap<Uuid, RunnerLease>>,
    build_metrics: BuildMetrics,
}

impl VmPool {
    pub fn new(build_metrics: BuildMetrics) -> Result<Self> {
        let config = RunnerConfig::from_env()
            .context("RUNNER_SCRIPT environment variable is required for vm_pool")?;
        Ok(Self {
            config,
            leases: Arc::new(DashMap::new()),
            build_metrics,
        })
    }

    /// Returns `None` when `RUNNER_SCRIPT` is not set — pool disabled gracefully.
    pub fn try_new(build_metrics: BuildMetrics) -> Option<Self> {
        Self::new(build_metrics).ok()
    }

    pub async fn spawn_runner(&self, req: SpawnRequest) -> Result<SpawnResponse> {
        let cfg = &self.config;
        let image_ref = if req.image.is_empty() { &cfg.image_ref } else { &req.image };

        let input = SpawnInput {
            cpu: req.cpu,
            memory_gb: req.memory_gb,
            disk_gb: req.disk_gb,
            time_budget_min: req.time_budget_min,
            image_ref,
            location: &cfg.location,
            ssh_key_ref: &cfg.ssh_key_ref,
            network_ref: cfg.network_ref.as_deref(),
            firewall_ref: cfg.firewall_ref.as_deref(),
        };

        let json_in = serde_json::to_string(&input).context("serialize SpawnInput")?;
        let stdout = dispatch(&cfg.nu_bin, &cfg.script, "spawn", &json_in).await?;
        let out: SpawnOutput = serde_json::from_str(stdout.trim())
            .context("deserialize SpawnOutput from runner.nu")?;

        let lease_id = Uuid::new_v4();
        let expires_at = Utc::now() + chrono::Duration::minutes(req.time_budget_min as i64);

        let lease = RunnerLease {
            lease_id,
            runner_handle: out.handle.clone(),
            ssh_host: out.ssh_host.clone(),
            ssh_port: out.ssh_port,
            spec: RunnerSpec {
                cpu: req.cpu,
                memory_gb: req.memory_gb,
                disk_gb: req.disk_gb,
                time_budget_min: req.time_budget_min,
            },
            workspace: req.workspace,
            created_at: Utc::now(),
            expires_at,
        };

        self.leases.insert(lease_id, lease);
        info!(lease_id = %lease_id, ssh_host = %out.ssh_host, "runner lease issued");

        Ok(SpawnResponse {
            lease_id,
            ssh_host: out.ssh_host,
            ssh_port: out.ssh_port,
            expires_at: expires_at.to_rfc3339(),
        })
    }

    pub async fn destroy_runner(&self, lease_id: Uuid) -> Result<()> {
        let (_, lease) = self
            .leases
            .remove(&lease_id)
            .ok_or_else(|| anyhow::anyhow!("lease {lease_id} not found"))?;

        let json_in = serde_json::json!({ "handle": lease.runner_handle }).to_string();
        dispatch(&self.config.nu_bin, &self.config.script, "destroy", &json_in).await?;

        info!(lease_id = %lease_id, "runner VM destroyed");
        Ok(())
    }

    pub fn p95_stats(&self, workspace: &str) -> Option<P95Stats> {
        p95_from_metrics(&self.build_metrics, workspace)
    }

    pub fn record_build_metrics(
        &self,
        workspace: &str,
        cpu_secs: f64,
        mem_peak_mb: f64,
        duration_secs: f64,
    ) {
        accumulate_metrics(&self.build_metrics, workspace, cpu_secs, mem_peak_mb, duration_secs);
    }
}

pub fn accumulate_metrics(
    store: &BuildMetrics,
    workspace: &str,
    cpu_secs: f64,
    mem_peak_mb: f64,
    duration_secs: f64,
) {
    let mut entry = store.entry(workspace.to_string()).or_insert((0.0, 0.0, 0.0, 0));
    entry.0 += cpu_secs;
    entry.1 += mem_peak_mb;
    entry.2 += duration_secs;
    entry.3 += 1;
}

pub fn p95_from_metrics(store: &BuildMetrics, workspace: &str) -> Option<P95Stats> {
    store.get(workspace).map(|acc| {
        let (sum_cpu, sum_mem, sum_dur, count) = *acc;
        if count == 0 {
            return P95Stats { cpu_p95: 0.0, memory_mb_p95: 0.0, duration_secs_p95: 0.0, sample_count: 0 };
        }
        P95Stats {
            cpu_p95: sum_cpu / count as f64,
            memory_mb_p95: sum_mem / count as f64,
            duration_secs_p95: sum_dur / count as f64,
            sample_count: count,
        }
    })
}

async fn dispatch(nu_bin: &str, script: &PathBuf, command: &str, json_input: &str) -> Result<String> {
    let mut child = tokio::process::Command::new(nu_bin)
        .arg(script)
        .arg(command)
        .stdin(std::process::Stdio::piped())
        .stdout(std::process::Stdio::piped())
        .stderr(std::process::Stdio::piped())
        .spawn()
        .with_context(|| format!("failed to spawn `nu {} {command}`", script.display()))?;

    if let Some(mut stdin) = child.stdin.take() {
        stdin
            .write_all(json_input.as_bytes())
            .await
            .with_context(|| format!("failed to write to runner `{command}` stdin"))?;
    }

    let out = child
        .wait_with_output()
        .await
        .with_context(|| format!("runner `{command}` process error"))?;

    if !out.status.success() {
        let stderr = String::from_utf8_lossy(&out.stderr);
        bail!(
            "runner `{command}` exit {}: {}",
            out.status.code().unwrap_or(-1),
            stderr.lines().last().unwrap_or("(no stderr)")
        );
    }

    String::from_utf8(out.stdout).context("runner stdout not valid UTF-8")
}

