use crate::config::Config;
use crate::dto::{EmbedRequest, EmbedResponse};
use crate::errors::OrchestratorError;
use anyhow::{anyhow, Context, Result};
use std::process::Stdio;
use std::time::Duration;
use tokio::process::{Child, Command};
use tracing::{info, warn};

pub struct SwiftBackend {
    _child: Child,
    client: reqwest::Client,
    base_url: String,
}

impl SwiftBackend {
    pub async fn start(cfg: &Config) -> Result<Self> {
        let base_url = format!("http://127.0.0.1:{}", cfg.swift_port);
        info!(bin = ?cfg.swift_bin, port = cfg.swift_port, "spawning swift backend");

        let child = Command::new(&cfg.swift_bin)
            .env("EMBED_PORT", cfg.swift_port.to_string())
            .env("EMBED_PRELOAD", "0.6B")
            .env("MLX_MODEL_DIR_0_6B", &cfg.swift_model_dir)
            .stdout(Stdio::inherit())
            .stderr(Stdio::inherit())
            .kill_on_drop(true)
            .spawn()
            .with_context(|| format!("failed to spawn swift binary {:?}", cfg.swift_bin))?;

        let client = reqwest::Client::builder()
            .timeout(Duration::from_secs(600))
            .build()?;

        // Poll /healthz for up to ~180 s to accommodate MLX warmup of all buckets.
        let deadline = std::time::Instant::now() + Duration::from_secs(180);
        loop {
            if std::time::Instant::now() > deadline {
                return Err(anyhow!("swift backend failed to become healthy within 180s"));
            }
            match client.get(format!("{base_url}/healthz")).send().await {
                Ok(r) if r.status().is_success() => {
                    info!("swift backend ready");
                    break;
                }
                _ => tokio::time::sleep(Duration::from_millis(500)).await,
            }
        }

        Ok(Self {
            _child: child,
            client,
            base_url,
        })
    }

    pub async fn embed(
        &self,
        req: &EmbedRequest,
    ) -> Result<EmbedResponse, OrchestratorError> {
        let resp = self
            .client
            .post(format!("{}/embed", self.base_url))
            .json(req)
            .send()
            .await
            .map_err(|e| OrchestratorError::Swift(e.to_string()))?;
        if !resp.status().is_success() {
            let status = resp.status();
            let body = resp.text().await.unwrap_or_default();
            return Err(OrchestratorError::Swift(format!("{status}: {body}")));
        }
        resp.json::<EmbedResponse>()
            .await
            .map_err(|e| OrchestratorError::Swift(e.to_string()))
    }

    pub async fn healthz(&self) -> bool {
        match self.client.get(format!("{}/healthz", self.base_url)).send().await {
            Ok(r) => r.status().is_success(),
            Err(e) => {
                warn!(error=?e, "swift /healthz failed");
                false
            }
        }
    }
}