//! End-to-end smoke: a tiny agent runs against fakes and completes.

use klieo_core::agent::{Agent, AgentContext};
use klieo_core::ids::ThreadId;
use klieo_core::llm::{ToolCall, ToolDef};
use klieo_core::memory::Episode;
use klieo_core::runtime::{run_steps, RunOptions};
use klieo_core::test_utils::{fake_context, FakeLlmClient, FakeLlmStep, FakeToolInvoker};
use std::sync::Arc;

#[derive(Clone)]
struct EchoAgent;

#[async_trait::async_trait]
impl Agent for EchoAgent {
    type Input = String;
    type Output = String;
    type Error = klieo_core::Error;

    fn name(&self) -> &str {
        "echo"
    }
    fn system_prompt(&self) -> &str {
        "you are a tiny test agent"
    }
    fn tools(&self) -> &[ToolDef] {
        &[]
    }

    async fn run(
        &self,
        ctx: AgentContext,
        input: Self::Input,
    ) -> Result<Self::Output, Self::Error> {
        let thread = ThreadId::new("test-thread");
        ctx.short_term
            .append(
                thread.clone(),
                klieo_core::Message {
                    role: klieo_core::Role::User,
                    content: input,
                    tool_calls: vec![],
                    tool_call_id: None,
                },
            )
            .await?;
        let final_text =
            run_steps(&ctx, self.system_prompt(), thread, RunOptions::default()).await?;
        Ok(final_text)
    }
}

fn ctx_with(
    llm: Arc<dyn klieo_core::LlmClient>,
    tools: Arc<dyn klieo_core::ToolInvoker>,
) -> AgentContext {
    let mut ctx = fake_context("echo");
    ctx.llm = llm;
    ctx.tools = tools;
    ctx
}

#[tokio::test]
async fn agent_completes_after_text_response() {
    let llm = Arc::new(
        FakeLlmClient::new("fake-llm").with_steps(vec![FakeLlmStep::Text("hello back".into())]),
    );
    let tools = Arc::new(FakeToolInvoker::new());
    let ctx = ctx_with(llm, tools);
    let agent = EchoAgent;
    let out = agent.run(ctx, "hello".into()).await.unwrap();
    assert_eq!(out, "hello back");
}

#[tokio::test]
async fn agent_dispatches_tool_call_then_completes() {
    let llm = Arc::new(FakeLlmClient::new("fake-llm").with_steps(vec![
        FakeLlmStep::ToolCalls(vec![ToolCall::new(
            "t1",
            "echo",
            serde_json::json!({"x": 1}),
        )]),
        FakeLlmStep::Text("done".into()),
    ]));
    let tools = Arc::new(
        FakeToolInvoker::new().with_tool("echo", "echoes args", |args| {
            Ok(serde_json::json!({"echoed": args}))
        }),
    );
    let ctx = ctx_with(llm, tools);
    let episodic = ctx.episodic.clone();
    let run_id = ctx.run_id;
    let agent = EchoAgent;
    let out = agent.run(ctx, "hi".into()).await.unwrap();
    assert_eq!(out, "done");

    let eps = episodic.replay(run_id).await.unwrap();
    let tool_eps = eps
        .iter()
        .filter(|e| matches!(e, Episode::ToolCall { .. }))
        .count();
    let llm_eps = eps
        .iter()
        .filter(|e| matches!(e, Episode::LlmCall { .. }))
        .count();
    assert_eq!(tool_eps, 1, "one tool call recorded");
    assert_eq!(llm_eps, 2, "two LLM calls recorded");
}

#[tokio::test]
async fn agent_aborts_on_cancellation() {
    let llm = Arc::new(
        FakeLlmClient::new("fake-llm")
            .with_steps(vec![FakeLlmStep::Text("would not reach".into())]),
    );
    let tools = Arc::new(FakeToolInvoker::new());
    let ctx = ctx_with(llm, tools);
    let cancel = ctx.cancel.clone();
    cancel.cancel();
    let agent = EchoAgent;
    let err = agent.run(ctx, "hi".into()).await.err().unwrap();
    assert!(matches!(err, klieo_core::Error::Cancelled));
}

#[tokio::test]
async fn agent_errors_on_max_steps() {
    let llm = Arc::new(FakeLlmClient::new("fake-llm").with_steps(vec![
        // tool-call loop — runtime will exceed max_steps before a Text step
        FakeLlmStep::ToolCalls(vec![ToolCall::new("t1", "echo", serde_json::json!({}))]),
        FakeLlmStep::ToolCalls(vec![ToolCall::new("t2", "echo", serde_json::json!({}))]),
        FakeLlmStep::ToolCalls(vec![ToolCall::new("t3", "echo", serde_json::json!({}))]),
    ]));
    let tools = Arc::new(FakeToolInvoker::new().with_tool("echo", "", Ok));
    let ctx = ctx_with(llm, tools);

    struct LowAgent;
    #[async_trait::async_trait]
    impl Agent for LowAgent {
        type Input = ();
        type Output = String;
        type Error = klieo_core::Error;
        fn name(&self) -> &str {
            "low"
        }
        fn system_prompt(&self) -> &str {
            ""
        }
        fn tools(&self) -> &[ToolDef] {
            &[]
        }
        async fn run(&self, ctx: AgentContext, _: ()) -> Result<Self::Output, Self::Error> {
            run_steps(
                &ctx,
                "",
                ThreadId::new("t"),
                RunOptions::default().with_max_steps(2),
            )
            .await
        }
    }
    let err = LowAgent.run(ctx, ()).await.err().unwrap();
    assert!(matches!(
        err,
        klieo_core::Error::MaxStepsExceeded { steps: 2 }
    ));
}

#[tokio::test]
async fn run_loop_emits_run_attributed_when_tenant_label_is_set() {
    let llm =
        Arc::new(FakeLlmClient::new("fake-llm").with_steps(vec![FakeLlmStep::Text("ok".into())]));
    let tools = Arc::new(FakeToolInvoker::new());
    let ctx = ctx_with(llm, tools).with_tenant_label("hashed-tenant-001".into());
    let episodic = ctx.episodic.clone();
    let run_id = ctx.run_id;
    EchoAgent.run(ctx, "hi".into()).await.unwrap();

    let eps = episodic.replay(run_id).await.unwrap();
    let attributed: Vec<&Episode> = eps
        .iter()
        .filter(|e| matches!(e, Episode::RunAttributed { .. }))
        .collect();
    assert_eq!(
        attributed.len(),
        1,
        "exactly one RunAttributed at run entry; got {eps:?}"
    );
    match attributed[0] {
        Episode::RunAttributed { tenant_label } => {
            assert_eq!(tenant_label, "hashed-tenant-001");
        }
        other => panic!("unexpected episode: {other:?}"),
    }
}

#[tokio::test]
async fn run_loop_emits_no_run_attributed_when_tenant_label_absent() {
    let llm =
        Arc::new(FakeLlmClient::new("fake-llm").with_steps(vec![FakeLlmStep::Text("ok".into())]));
    let tools = Arc::new(FakeToolInvoker::new());
    // No `with_tenant_label` call — the default path.
    let ctx = ctx_with(llm, tools);
    let episodic = ctx.episodic.clone();
    let run_id = ctx.run_id;
    EchoAgent.run(ctx, "hi".into()).await.unwrap();

    let eps = episodic.replay(run_id).await.unwrap();
    let count = eps
        .iter()
        .filter(|e| matches!(e, Episode::RunAttributed { .. }))
        .count();
    assert_eq!(count, 0, "no attribution episode without a label");
}

#[tokio::test]
async fn run_loop_emits_run_origin_verbatim_when_parent_anchor_is_set() {
    let llm =
        Arc::new(FakeLlmClient::new("fake-llm").with_steps(vec![FakeLlmStep::Text("ok".into())]));
    let tools = Arc::new(FakeToolInvoker::new());
    let anchor = "sha256:deadbeefcafe";
    let ctx = ctx_with(llm, tools).with_parent_anchor(anchor.into());
    let episodic = ctx.episodic.clone();
    let run_id = ctx.run_id;
    EchoAgent.run(ctx, "hi".into()).await.unwrap();

    let eps = episodic.replay(run_id).await.unwrap();
    let origins: Vec<&Episode> = eps
        .iter()
        .filter(|e| matches!(e, Episode::RunOrigin { .. }))
        .collect();
    assert_eq!(
        origins.len(),
        1,
        "exactly one RunOrigin at run entry; got {eps:?}"
    );
    match origins[0] {
        Episode::RunOrigin { parent_anchor } => {
            assert_eq!(
                parent_anchor, anchor,
                "anchor recorded verbatim, not hashed or rewritten"
            );
        }
        other => panic!("unexpected episode: {other:?}"),
    }
}

#[tokio::test]
async fn run_loop_emits_no_run_origin_when_parent_anchor_absent() {
    let llm =
        Arc::new(FakeLlmClient::new("fake-llm").with_steps(vec![FakeLlmStep::Text("ok".into())]));
    let tools = Arc::new(FakeToolInvoker::new());
    let ctx = ctx_with(llm, tools);
    let episodic = ctx.episodic.clone();
    let run_id = ctx.run_id;
    EchoAgent.run(ctx, "hi".into()).await.unwrap();

    let eps = episodic.replay(run_id).await.unwrap();
    let count = eps
        .iter()
        .filter(|e| matches!(e, Episode::RunOrigin { .. }))
        .count();
    assert_eq!(count, 0, "no origin episode without an anchor");
}
