use crate::config::Config;
use crate::dto::{EmbedRequest, EmbedResponse};
use crate::errors::OrchestratorError;
use std::process::Stdio;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::process::{Child, Command};
use tokio::sync::Mutex;
use tracing::{info, warn};

pub struct RustBackend {
    cfg: Arc<Config>,
    client: reqwest::Client,
    spawn_mutex: Mutex<()>,
}

impl RustBackend {
    pub fn new(cfg: Arc<Config>) -> Self {
        let client = reqwest::Client::builder()
            .timeout(Duration::from_secs(1800))
            .build()
            .expect("reqwest client");
        Self {
            cfg,
            client,
            spawn_mutex: Mutex::new(()),
        }
    }

    /// Spawn a fresh Rust process sized for this bucket, run the request, tear down.
    pub async fn run_one(
        &self,
        max_len: usize,
        req: &EmbedRequest,
    ) -> Result<EmbedResponse, OrchestratorError> {
        let _g = self.spawn_mutex.lock().await;
        let base_url = format!("http://127.0.0.1:{}", self.cfg.rust_port);

        info!(max_len, port = self.cfg.rust_port, "spawning rust backend");
        let mut child: Child = Command::new(&self.cfg.rust_bin)
            .env("EMBED_PORT", self.cfg.rust_port.to_string())
            .env("EMBED_MAX_LEN", max_len.to_string())
            .env("EMBED_MAX_BATCH", "1")
            .env("MODEL_PATH", &self.cfg.rust_model_path)
            .stdout(Stdio::inherit())
            .stderr(Stdio::inherit())
            .kill_on_drop(true)
            .spawn()
            .map_err(|e| OrchestratorError::Rust(format!("spawn failed: {e}")))?;

        // Poll /healthz until ready or timeout.
        let deadline =
            Instant::now() + Duration::from_millis(self.cfg.rust_start_timeout_ms);
        loop {
            if Instant::now() > deadline {
                let _ = child.kill().await;
                return Err(OrchestratorError::Rust(
                    "rust backend did not become healthy in time".into(),
                ));
            }
            match self.client.get(format!("{base_url}/healthz")).send().await {
                Ok(r) if r.status().is_success() => break,
                _ => tokio::time::sleep(Duration::from_millis(200)).await,
            }
        }

        // Run the embedding request.
        let result = async {
            let resp = self
                .client
                .post(format!("{base_url}/embed"))
                .json(req)
                .send()
                .await
                .map_err(|e| OrchestratorError::Rust(e.to_string()))?;
            if !resp.status().is_success() {
                let status = resp.status();
                let body = resp.text().await.unwrap_or_default();
                return Err(OrchestratorError::Rust(format!("{status}: {body}")));
            }
            resp.json::<EmbedResponse>()
                .await
                .map_err(|e| OrchestratorError::Rust(e.to_string()))
        }
        .await;

        // Tear down regardless of success.
        if let Err(e) = child.kill().await {
            warn!(error=?e, "failed to kill rust backend");
        }
        let _ = child.wait().await;
        result
    }
}