//! Pool sandbox tests against a real `monty` worker. Skipped unless
//! `MONTY_WORKER` points at one (`cargo install monty-runtime --locked`).

use std::path::PathBuf;
use std::sync::Arc;
use std::time::{Duration, Instant};

use async_trait::async_trait;
use asyncai::code_mode::{CodeModeLimits, CodeModeOptions, SandboxKind};
use asyncai::{CodeModeToolExecutor, ToolCallOutput, ToolDefinition, ToolError, ToolExecutor};
use serde_json::{json, Value};

struct Pods;

#[async_trait]
impl ToolExecutor for Pods {
    async fn list_tools(&self) -> Result<Vec<ToolDefinition>, ToolError> {
        Ok(vec![
            ToolDefinition::new(
                "list_pods",
                None::<String>,
                json!({"type": "object", "properties": {"namespace": {"type": "string"}}}),
            ),
            ToolDefinition::new("slow", None::<String>, json!({"type": "object"})),
        ])
    }

    async fn call_tool(&self, name: &str, _arguments: Value) -> Result<ToolCallOutput, ToolError> {
        match name {
            "list_pods" => {
                tokio::time::sleep(Duration::from_millis(150)).await;
                Ok(ToolCallOutput::text(
                    json!({"items": [
                        {"metadata": {"name": "a"}, "status": {"phase": "Running"}},
                        {"metadata": {"name": "b"}, "status": {"phase": "Failed"}},
                    ]})
                    .to_string(),
                ))
            }
            _ => {
                tokio::time::sleep(Duration::from_secs(30)).await;
                Ok(ToolCallOutput::text("late"))
            }
        }
    }
}

fn worker() -> Option<PathBuf> {
    let path = PathBuf::from(std::env::var_os("MONTY_WORKER")?);
    path.is_file().then_some(path)
}

async fn pool(limits: CodeModeLimits) -> Option<CodeModeToolExecutor> {
    let Some(worker) = worker() else {
        eprintln!(
            "{}",
            json!({"level": "info", "event": "test.skipped", "reason": "MONTY_WORKER not set"})
        );
        return None;
    };
    let executor = CodeModeToolExecutor::with_options(
        Arc::new(Pods),
        CodeModeOptions {
            limits,
            sandbox: SandboxKind::Pool,
            monty_worker: Some(worker),
            ..CodeModeOptions::default()
        },
    )
    .await
    .unwrap();
    assert_eq!(executor.sandbox_name(), "pool");
    Some(executor)
}

async fn run(executor: &dyn ToolExecutor, code: &str) -> ToolCallOutput {
    executor
        .call_tool("run_code", json!({"code": code}))
        .await
        .unwrap()
}

#[tokio::test(flavor = "multi_thread")]
async fn pool_gathers_filters_and_keeps_state() {
    let Some(executor) = pool(CodeModeLimits::default()).await else {
        return;
    };
    let scope = executor.run_scope().unwrap();
    let started = Instant::now();
    let output = run(
        scope.as_ref(),
        "import asyncio\n\
         a, b = await asyncio.gather(list_pods(namespace='x'), list_pods(namespace='y'))\n\
         phases = count_by(a, 'status.phase')\n\
         {'names': names(a) + names(b), 'phases': phases}",
    )
    .await;
    assert!(!output.is_error, "{}", output.content);
    assert_eq!(
        serde_json::from_str::<Value>(&output.content).unwrap(),
        json!({"names": ["a", "b", "a", "b"], "phases": {"Running": 1, "Failed": 1}})
    );
    assert!(started.elapsed() < Duration::from_millis(1500));
    assert_eq!(run(scope.as_ref(), "phases['Failed']").await.content, "1");
}

#[tokio::test(flavor = "multi_thread")]
async fn pool_worker_memory_cap_resets_state_and_agent_survives() {
    let Some(executor) = pool(CodeModeLimits {
        max_memory: 64 * 1024 * 1024,
        ..CodeModeLimits::default()
    })
    .await
    else {
        return;
    };
    let scope = executor.run_scope().unwrap();
    run(scope.as_ref(), "keep = 1").await;
    let blown = run(scope.as_ref(), "x = 'a' * (512 * 1024 * 1024)\nlen(x)").await;
    assert!(blown.is_error, "{}", blown.content);
    assert!(blown.content.contains("MemoryError"), "{}", blown.content);
    let after = run(scope.as_ref(), "1 + 1").await;
    assert!(!after.is_error, "{}", after.content);
    assert_eq!(after.content, "2");
}

#[tokio::test(flavor = "multi_thread")]
async fn pool_tool_timeout_and_runaway_loops_are_bounded() {
    let Some(executor) = pool(CodeModeLimits {
        tool_timeout: Duration::from_millis(200),
        max_duration: Duration::from_millis(300),
        ..CodeModeLimits::default()
    })
    .await
    else {
        return;
    };
    let scope = executor.run_scope().unwrap();
    let timed_out = run(scope.as_ref(), "await slow()").await;
    assert!(
        timed_out.content.contains("timed out after 200ms"),
        "{}",
        timed_out.content
    );

    let started = Instant::now();
    let spin = run(scope.as_ref(), "while True:\n    pass").await;
    assert!(spin.is_error, "{}", spin.content);
    assert!(started.elapsed() < Duration::from_secs(15));
    assert_eq!(run(scope.as_ref(), "2 * 21").await.content, "42");
}

#[tokio::test]
async fn pool_scopes_report_variables_and_reject_helper_shadowing() {
    let Some(executor) = pool(CodeModeLimits::default()).await else {
        return;
    };
    let scope = executor.thread_scope("pool-thread").unwrap();
    let run = |code: &'static str| {
        let scope = scope.clone();
        async move {
            scope
                .call_tool("run_code", json!({"code": code}))
                .await
                .unwrap()
        }
    };
    let fetched = run(
        "pods = await list_pods(namespace='a')\nphases = count_by(pods, 'status', 'phase')\nphases",
    )
    .await;
    assert!(!fetched.is_error, "{}", fetched.content);
    assert_eq!(fetched.content, r#"{"Failed":1,"Running":1}"#);
    let note = scope.context_note().unwrap();
    assert!(
        note.contains("pods: dict[1]") && note.contains("phases: dict[2]"),
        "{note}"
    );
    assert!(run("dig = 1").await.is_error);
    assert!(run("to_json({'a': (1, 2)})")
        .await
        .content
        .contains("[1, 2]"));
}