vibe-tests 0.0.1

//! EngineTests builder — constructs test engine for MCP servers.

use std::collections::HashMap;
use std::path::PathBuf;
use std::time::{Duration, Instant};
use tempfile::tempdir;
use tracing::Level;

use crate::EngineTests;
use crate::base::alias::{OnLog, OnRun, OnStart, OnStop};
use crate::base::error::TestError;
use crate::base::error::TestsResult;
use crate::base::tee_writer::TeeWriter;
use crate::engine::engine_env::EngineEnv;
use crate::engine::engine_events::EngineEvents;
use crate::engine::engine_state::EngineState;
use crate::env::env_log::EnvLog;
use crate::env::env_run::EnvRun;
use crate::env::env_start::EnvStart;
use crate::env::env_stop::EnvStop;

/// Builder for VibeTests.
pub struct EngineBuilder {
    /// Optional docker-compose file for test infrastructure.
    compose_file: Option<PathBuf>,
    /// MCP server URL for health check and tool calls.
    mcp_host: String,
    /// Ollama API host URL.
    ollama_host: String,
    /// Models to test queries against.
    ollama_models: Vec<String>,
    /// Unload other models before test to free GPU memory.
    ollama_exclusive: bool,
    /// Startup callback.
    on_start: Option<OnStart>,
    /// Per-test callback.
    on_run: Option<OnRun>,
    /// Shutdown callback.
    on_stop: Option<OnStop>,
    /// Log event handler.
    on_log: Option<OnLog>,
    /// Log level for internal tracing.
    log_level: Level,
    /// Timeout for all HTTP requests.
    timeout: Duration,
}

impl EngineBuilder {
    /// Creates a new builder with defaults.
    pub fn new() -> Self {
        Self {
            compose_file: None,
            mcp_host: String::new(),
            ollama_host: String::new(),
            ollama_models: Vec::new(),
            ollama_exclusive: true,
            on_start: None,
            on_run: None,
            on_stop: None,
            on_log: None,
            log_level: Level::ERROR,
            timeout: Duration::from_secs(60),
        }
    }

    /// Sets the docker-compose file path for the test environment (optional).
    pub fn compose_file(mut self, path: impl Into<PathBuf>) -> Self {
        self.compose_file = Some(path.into());
        self
    }

    /// Sets the MCP server URL (health check + tool calls).
    pub fn mcp_host(mut self, host: impl Into<String>) -> Self {
        self.mcp_host = host.into();
        self
    }

    /// Sets the Ollama API host URL.
    pub fn ollama_host(mut self, host: impl Into<String>) -> Self {
        self.ollama_host = host.into();
        self
    }

    /// Sets the Ollama models to test against.
    pub fn ollama_models(mut self, models: &[&str]) -> Self {
        self.ollama_models = models.iter().map(|s| s.to_string()).collect();
        self
    }

    /// Unload other models before each test to ensure full GPU memory for current model.
    pub fn ollama_exclusive(mut self, yes: bool) -> Self {
        self.ollama_exclusive = yes;
        self
    }

    /// Sets the startup callback. Receives EnvStart with isolated home directory and tee writer.
    /// Returns optional key-value data passed to on_stop.
    /// Compose (if specified) is already up when this runs.
    pub fn on_start<F, Fut>(mut self, f: F) -> Self
    where
        F: FnOnce(EnvStart) -> Fut + Send + 'static,
        Fut: Future<Output = TestsResult<Option<HashMap<String, String>>>> + Send + 'static,
    {
        self.on_start = Some(Box::new(move |env| Box::pin(f(env))));
        self
    }
    /// Sets the per-test callback. Receives EnvRun with home and current model.
    /// Runs before each test() call.
    pub fn on_run<F, Fut>(mut self, f: F) -> Self
    where
        F: Fn(EnvRun) -> Fut + Send + Sync + 'static,
        Fut: Future<Output = TestsResult<()>> + Send + 'static,
    {
        self.on_run = Some(Box::new(move |env| Box::pin(f(env))));
        self
    }

    /// Sets the shutdown callback. Receives EnvStop with home and total duration.
    /// Runs before compose down and home cleanup.
    pub fn on_stop<F>(mut self, f: F) -> Self
    where
        F: FnOnce(EnvStop) + Send + 'static,
    {
        self.on_stop = Some(Box::new(f));
        self
    }

    /// Sets the log event handler.
    pub fn on_log<F>(mut self, f: F) -> Self
    where
        F: Fn(EnvLog) + Send + Sync + 'static,
    {
        self.on_log = Some(Box::new(f));
        self
    }

    /// Sets the log level for internal tracing.
    pub fn log_level(mut self, level: Level) -> Self {
        self.log_level = level;
        self
    }

    /// Sets the timeout for all operations.
    pub fn timeout(mut self, timeout: Duration) -> Self {
        self.timeout = timeout;
        self
    }

    /// Validates and returns the test engine.
    pub fn build(self) -> TestsResult<EngineTests> {
        // Validate required parameters
        if self.ollama_host.is_empty() {
            return Err(TestError::Setup("ollama_host must not be empty".into()));
        }
        if self.ollama_models.is_empty() {
            return Err(TestError::Setup("ollama_models must not be empty".into()));
        }
        if self.mcp_host.is_empty() {
            return Err(TestError::Setup("mcp_host must not be empty".into()));
        }

        // Create isolated home directory for this test run
        let home =
            tempdir().map_err(|e| TestError::Setup(format!("Failed to create temp dir: {}", e)))?;

        // Create TeeWriter: writes child output to log file and forwards to on_log callback
        let log_path = home.path().join("vibe-tests.log");
        let tee = TeeWriter::new(log_path, self.log_level, self.on_log)
            .map_err(|e| TestError::Setup(format!("Failed to create TeeWriter: {}", e)))?;

        // Assemble the test engine
        Ok(EngineTests {
            env: EngineEnv {
                compose_file: self.compose_file,
                mcp_host: self.mcp_host,
                ollama_host: self.ollama_host,
                ollama_models: self.ollama_models,
                ollama_exclusive: self.ollama_exclusive,
                log_level: self.log_level,
                timeout: self.timeout,
            },
            events: EngineEvents {
                on_start: self.on_start,
                on_run: self.on_run,
                on_stop: self.on_stop,
            },
            state: EngineState {
                initialized: false,
                home,
                tee,
                compose: None,
                runner: None,
                start_data: None,
                start_time: Instant::now(),
            },
        })
    }
}