//! Integration tests for [`EpisodeProjector`] — drives the public
//! `spawn` + `project_run` surface against in-test mocks for
//! `KnowledgeGraph` and `EntityExtractor`.

use async_trait::async_trait;
use chrono::{DateTime, Utc};
use klieo_core::error::MemoryError;
use klieo_core::ids::{FactId, RunId};
use klieo_core::memory::{Episode, EpisodicMemory, Scope};
use klieo_core::test_utils::InMemoryEpisodic;
use klieo_graph_projector::{
    EpisodeProjector, EpisodeProjectorBuilder, HintsOnlyExtractor, ProjectableEpisodic,
};
use klieo_memory_graph::{EntityRef, IndexEntry, KnowledgeGraph, RecallMetrics};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use tokio::sync::broadcast;
use tokio::time::sleep;
use tokio_util::sync::CancellationToken;

#[derive(Clone, Debug)]
struct Indexed {
    fact_id: FactId,
    entity_names: Vec<String>,
}

/// Records every successful `index()` call.
#[derive(Default)]
struct RecordingGraph {
    indexed: Mutex<Vec<Indexed>>,
}

impl RecordingGraph {
    fn snapshot(&self) -> Vec<Indexed> {
        self.indexed.lock().unwrap().clone()
    }
}

#[async_trait]
impl KnowledgeGraph for RecordingGraph {
    async fn index(
        &self,
        _scope: Scope,
        fact_id: &FactId,
        entities: &[EntityRef],
        _text: &str,
        _valid_from: Option<DateTime<Utc>>,
    ) -> Result<(), MemoryError> {
        self.indexed.lock().unwrap().push(Indexed {
            fact_id: fact_id.clone(),
            entity_names: entities.iter().map(|e| e.name.clone()).collect(),
        });
        Ok(())
    }

    async fn neighbors(
        &self,
        _scope: &Scope,
        _entities: &[EntityRef],
    ) -> Result<Vec<FactId>, MemoryError> {
        Ok(Vec::new())
    }
}

/// Always-failing graph that still records how many index attempts
/// have been made so the spawn-keeps-running assertion can probe it.
#[derive(Default)]
struct AlwaysFailGraph {
    attempts: AtomicUsize,
}

/// `EpisodicMemory` whose `replay()` always fails. Used to drive the
/// `ProjectionError::Replay` branch of `project_run`.
struct ReplayFailingEpisodic;

#[async_trait]
impl EpisodicMemory for ReplayFailingEpisodic {
    async fn record(&self, _run: RunId, _event: Episode) -> Result<(), MemoryError> {
        Ok(())
    }
    async fn replay(&self, _run: RunId) -> Result<Vec<Episode>, MemoryError> {
        Err(MemoryError::Store("simulated replay failure".into()))
    }
    async fn list_runs(
        &self,
        _filter: klieo_core::memory::RunFilter,
    ) -> Result<Vec<klieo_core::memory::RunSummary>, MemoryError> {
        Ok(Vec::new())
    }
}

#[async_trait]
impl KnowledgeGraph for AlwaysFailGraph {
    async fn index(
        &self,
        _scope: Scope,
        _fact_id: &FactId,
        _entities: &[EntityRef],
        _text: &str,
        _valid_from: Option<DateTime<Utc>>,
    ) -> Result<(), MemoryError> {
        self.attempts.fetch_add(1, Ordering::SeqCst);
        Err(MemoryError::Store("simulated graph backend failure".into()))
    }

    async fn neighbors(
        &self,
        _scope: &Scope,
        _entities: &[EntityRef],
    ) -> Result<Vec<FactId>, MemoryError> {
        Ok(Vec::new())
    }
}

fn workspace_scope() -> Scope {
    Scope::Workspace("test-ws".into())
}

/// Poll a predicate up to `attempts` × 10ms (≤ 1s total) so async
/// projection assertions don't race the spawn loop.
async fn wait_for(attempts: u32, mut check: impl FnMut() -> bool) -> bool {
    for _ in 0..attempts {
        if check() {
            return true;
        }
        sleep(Duration::from_millis(10)).await;
    }
    check()
}

#[tokio::test]
async fn builder_defaults_indexes_replayed_episodes() {
    let graph = Arc::new(RecordingGraph::default());
    let extractor = Arc::new(HintsOnlyExtractor::default());
    let projector = EpisodeProjector::builder().build(
        Arc::clone(&graph) as Arc<dyn KnowledgeGraph>,
        extractor,
        workspace_scope(),
    );

    let episodic = InMemoryEpisodic::default();
    let run = RunId::new();
    episodic
        .record(
            run,
            Episode::Started {
                agent: "alpha".into(),
            },
        )
        .await
        .unwrap();

    let processed = projector.project_run(&episodic, run).await.unwrap();
    assert_eq!(processed, 1, "builder-produced projector replays episodes");
    assert_eq!(
        graph.snapshot().len(),
        1,
        "entity-emitting episode hit graph"
    );
}

#[tokio::test]
async fn builder_metrics_override_records_lag_on_shared_counter() {
    let graph = Arc::new(RecordingGraph::default());
    let extractor = Arc::new(HintsOnlyExtractor::default());
    let metrics = Arc::new(RecallMetrics::default());
    let projector = Arc::new(
        EpisodeProjectorBuilder::default()
            .metrics(Arc::clone(&metrics))
            .build(
                Arc::clone(&graph) as Arc<dyn KnowledgeGraph>,
                extractor,
                workspace_scope(),
            ),
    );

    let (tx, rx) = broadcast::channel::<(RunId, Episode)>(2);
    let run = RunId::new();
    for i in 0..5 {
        tx.send((
            run,
            Episode::Started {
                agent: format!("agent-{i}"),
            },
        ))
        .unwrap();
    }

    let cancel = CancellationToken::new();
    let handle = Arc::clone(&projector).spawn(rx, cancel.clone());
    assert!(
        wait_for(200, || metrics.snapshot().projector_lag_episodes >= 3).await,
        "shared metrics counter must register projector_lag_episodes from the builder-produced projector",
    );
    cancel.cancel();
    handle.await.unwrap();
}

#[tokio::test]
async fn spawn_processes_broadcast_events() {
    let graph = Arc::new(RecordingGraph::default());
    let extractor = Arc::new(HintsOnlyExtractor::default());
    let metrics = Arc::new(RecallMetrics::default());
    let projector = Arc::new(EpisodeProjector::new(
        Arc::clone(&graph) as Arc<dyn KnowledgeGraph>,
        extractor,
        workspace_scope(),
        Arc::clone(&metrics),
    ));

    let inner: Arc<dyn EpisodicMemory> = Arc::new(InMemoryEpisodic::default());
    let (wrapper, rx) = ProjectableEpisodic::new(inner);

    let cancel = CancellationToken::new();
    let handle = Arc::clone(&projector).spawn(rx, cancel.clone());

    let run = RunId::new();
    wrapper
        .record(
            run,
            Episode::Started {
                agent: "alpha".into(),
            },
        )
        .await
        .unwrap();
    wrapper
        .record(
            run,
            Episode::ToolCall {
                name: "WebSearch".into(),
                args: serde_json::json!({}),
                result: klieo_core::memory::ToolResult::ok(serde_json::json!("hit")),
            },
        )
        .await
        .unwrap();

    let ok = wait_for(100, || graph.snapshot().len() >= 2).await;
    assert!(ok, "expected 2 index calls within timeout");

    cancel.cancel();
    handle.await.unwrap();

    let indexed = graph.snapshot();
    assert_eq!(indexed.len(), 2);
    assert!(indexed[0].entity_names.iter().any(|n| n == "alpha"));
    assert!(indexed[1].entity_names.iter().any(|n| n == "websearch"));
}

#[tokio::test]
async fn project_run_replays_and_indexes_every_episode() {
    let graph = Arc::new(RecordingGraph::default());
    let extractor = Arc::new(HintsOnlyExtractor::default());
    let metrics = Arc::new(RecallMetrics::default());
    let projector = EpisodeProjector::new(
        Arc::clone(&graph) as Arc<dyn KnowledgeGraph>,
        extractor,
        workspace_scope(),
        metrics,
    );

    let episodic = InMemoryEpisodic::default();
    let run = RunId::new();
    episodic
        .record(
            run,
            Episode::Started {
                agent: "alpha".into(),
            },
        )
        .await
        .unwrap();
    episodic
        .record(
            run,
            Episode::ToolCall {
                name: "search".into(),
                args: serde_json::json!({}),
                result: klieo_core::memory::ToolResult::ok(serde_json::json!("ok")),
            },
        )
        .await
        .unwrap();
    // Episode::Completed produces no entities — counted but not indexed.
    episodic.record(run, Episode::Completed).await.unwrap();

    let processed = projector.project_run(&episodic, run).await.unwrap();
    assert_eq!(processed, 3, "all replayed episodes counted");

    let indexed = graph.snapshot();
    assert_eq!(
        indexed.len(),
        2,
        "only entity-emitting episodes hit the graph",
    );
}

/// Counts `index_many` invocations + remembers the largest batch size
/// observed. Used to assert `project_run` flushes the entire run in
/// one batched call rather than N per-fact `index` calls.
#[derive(Default)]
struct BatchCountingGraph {
    index_many_calls: AtomicUsize,
    largest_batch: AtomicUsize,
    index_calls: AtomicUsize,
}

#[async_trait]
impl KnowledgeGraph for BatchCountingGraph {
    async fn index(
        &self,
        _scope: Scope,
        _fact_id: &FactId,
        _entities: &[EntityRef],
        _text: &str,
        _valid_from: Option<DateTime<Utc>>,
    ) -> Result<(), MemoryError> {
        self.index_calls.fetch_add(1, Ordering::SeqCst);
        Ok(())
    }

    async fn index_many(&self, _scope: Scope, batch: &[IndexEntry]) -> Result<(), MemoryError> {
        self.index_many_calls.fetch_add(1, Ordering::SeqCst);
        self.largest_batch.fetch_max(batch.len(), Ordering::SeqCst);
        Ok(())
    }

    async fn neighbors(
        &self,
        _scope: &Scope,
        _entities: &[EntityRef],
    ) -> Result<Vec<FactId>, MemoryError> {
        Ok(Vec::new())
    }
}

#[tokio::test]
async fn project_run_flushes_entire_run_in_one_index_many_call() {
    // project_run must drive `index_many` exactly once per call
    // regardless of episode count. Falling back to per-fact
    // `index` defeats N+1 elimination on backends that override
    // `index_many` for bulk-insert (e.g. Neo4j UNWIND).
    let graph = Arc::new(BatchCountingGraph::default());
    let extractor = Arc::new(HintsOnlyExtractor::default());
    let metrics = Arc::new(RecallMetrics::default());
    let projector = EpisodeProjector::new(
        Arc::clone(&graph) as Arc<dyn KnowledgeGraph>,
        extractor,
        workspace_scope(),
        metrics,
    );

    let episodic = InMemoryEpisodic::default();
    let run = RunId::new();
    for i in 0..5 {
        episodic
            .record(
                run,
                Episode::ToolCall {
                    name: format!("tool-{i}"),
                    args: serde_json::json!({}),
                    result: klieo_core::memory::ToolResult::ok(serde_json::json!("ok")),
                },
            )
            .await
            .unwrap();
    }

    projector.project_run(&episodic, run).await.unwrap();

    assert_eq!(
        graph.index_many_calls.load(Ordering::SeqCst),
        1,
        "project_run must flush in exactly one index_many call",
    );
    assert_eq!(
        graph.largest_batch.load(Ordering::SeqCst),
        5,
        "all five entity-emitting episodes must land in the same batch",
    );
    assert_eq!(
        graph.index_calls.load(Ordering::SeqCst),
        0,
        "project_run must NOT fall back to per-fact index",
    );
}

#[tokio::test]
async fn project_run_with_no_indexable_episodes_skips_index_many() {
    // Run with only zero-entity episodes (Episode::Completed) must
    // not invoke `index_many` at all — empty batches are not a
    // graph round-trip.
    let graph = Arc::new(BatchCountingGraph::default());
    let extractor = Arc::new(HintsOnlyExtractor::default());
    let metrics = Arc::new(RecallMetrics::default());
    let projector = EpisodeProjector::new(
        Arc::clone(&graph) as Arc<dyn KnowledgeGraph>,
        extractor,
        workspace_scope(),
        metrics,
    );

    let episodic = InMemoryEpisodic::default();
    let run = RunId::new();
    episodic.record(run, Episode::Completed).await.unwrap();
    episodic.record(run, Episode::Completed).await.unwrap();

    let processed = projector.project_run(&episodic, run).await.unwrap();
    assert_eq!(processed, 2);
    assert_eq!(
        graph.index_many_calls.load(Ordering::SeqCst),
        0,
        "empty-batch runs must not call index_many",
    );
}

#[tokio::test]
async fn lagged_receiver_records_metric_and_continues() {
    let graph = Arc::new(RecordingGraph::default());
    let extractor = Arc::new(HintsOnlyExtractor::default());
    let metrics = Arc::new(RecallMetrics::default());
    let projector = Arc::new(EpisodeProjector::new(
        Arc::clone(&graph) as Arc<dyn KnowledgeGraph>,
        extractor,
        workspace_scope(),
        Arc::clone(&metrics),
    ));

    let (tx, rx) = broadcast::channel::<(RunId, Episode)>(2);
    let run = RunId::new();
    // Push 5 events into a capacity-2 buffer before spawning the
    // subscriber. The first recv() then yields Lagged(3); the next
    // two yield the two retained events.
    for i in 0..5 {
        tx.send((
            run,
            Episode::Started {
                agent: format!("agent-{i}"),
            },
        ))
        .unwrap();
    }

    let cancel = CancellationToken::new();
    let handle = Arc::clone(&projector).spawn(rx, cancel.clone());
    assert!(
        wait_for(200, || metrics.snapshot().projector_lag_episodes >= 3).await,
        "expected projector_lag_episodes >= 3 within timeout",
    );
    cancel.cancel();
    handle.await.unwrap();

    let snap = metrics.snapshot();
    assert_eq!(snap.projector_lag_episodes, 3, "exactly 3 dropped");
    let indexed = graph.snapshot();
    assert_eq!(indexed.len(), 2, "two retained events still indexed");
}

#[tokio::test]
async fn graph_failure_does_not_stop_the_spawn_loop() {
    let failing = Arc::new(AlwaysFailGraph::default());
    let extractor = Arc::new(HintsOnlyExtractor::default());
    let metrics = Arc::new(RecallMetrics::default());
    let projector = Arc::new(EpisodeProjector::new(
        Arc::clone(&failing) as Arc<dyn KnowledgeGraph>,
        extractor,
        workspace_scope(),
        metrics,
    ));

    let inner: Arc<dyn EpisodicMemory> = Arc::new(InMemoryEpisodic::default());
    let (wrapper, rx) = ProjectableEpisodic::new(inner);

    let cancel = CancellationToken::new();
    let handle = Arc::clone(&projector).spawn(rx, cancel.clone());

    let run = RunId::new();
    for i in 0..3 {
        wrapper
            .record(
                run,
                Episode::Started {
                    agent: format!("agent-{i}"),
                },
            )
            .await
            .unwrap();
    }

    let ok = wait_for(100, || failing.attempts.load(Ordering::SeqCst) >= 3).await;
    assert!(ok, "spawn loop must keep attempting after each failure");

    cancel.cancel();
    handle.await.unwrap();

    assert_eq!(failing.attempts.load(Ordering::SeqCst), 3);
}

#[tokio::test]
async fn project_run_returns_replay_error_when_episodic_fails() {
    let graph = Arc::new(RecordingGraph::default());
    let projector = EpisodeProjector::new(
        Arc::clone(&graph) as Arc<dyn KnowledgeGraph>,
        Arc::new(HintsOnlyExtractor::default()),
        workspace_scope(),
        Arc::new(RecallMetrics::default()),
    );

    let failing = ReplayFailingEpisodic;
    let err = projector
        .project_run(&failing, RunId::new())
        .await
        .expect_err("replay failure must propagate");
    match err {
        klieo_graph_projector::ProjectionError::Replay { source } => {
            let inner = source
                .downcast_ref::<MemoryError>()
                .expect("source downcasts to MemoryError");
            assert!(
                matches!(inner, MemoryError::Store(s) if s.contains("simulated replay failure"))
            );
        }
        other => panic!("expected ProjectionError::Replay, got {other:?}"),
    }
    assert!(
        graph.snapshot().is_empty(),
        "replay failure must not index anything",
    );
}

#[tokio::test]
async fn project_run_returns_graph_error_when_index_fails() {
    let failing_graph = Arc::new(AlwaysFailGraph::default());
    let projector = EpisodeProjector::new(
        Arc::clone(&failing_graph) as Arc<dyn KnowledgeGraph>,
        Arc::new(HintsOnlyExtractor::default()),
        workspace_scope(),
        Arc::new(RecallMetrics::default()),
    );

    let episodic = InMemoryEpisodic::default();
    let run = RunId::new();
    episodic
        .record(
            run,
            Episode::Started {
                agent: "alpha".into(),
            },
        )
        .await
        .unwrap();

    let err = projector
        .project_run(&episodic, run)
        .await
        .expect_err("graph failure must propagate");
    match err {
        klieo_graph_projector::ProjectionError::Graph { source } => {
            let inner = source
                .downcast_ref::<MemoryError>()
                .expect("source downcasts to MemoryError");
            assert!(
                matches!(inner, MemoryError::Store(s) if s.contains("simulated graph backend failure"))
            );
        }
        other => panic!("expected ProjectionError::Graph, got {other:?}"),
    }
    assert_eq!(failing_graph.attempts.load(Ordering::SeqCst), 1);
}

#[tokio::test]
async fn spawn_exits_cleanly_on_immediate_cancel() {
    let graph = Arc::new(RecordingGraph::default());
    let projector = Arc::new(EpisodeProjector::new(
        Arc::clone(&graph) as Arc<dyn KnowledgeGraph>,
        Arc::new(HintsOnlyExtractor::default()),
        workspace_scope(),
        Arc::new(RecallMetrics::default()),
    ));

    let inner: Arc<dyn EpisodicMemory> = Arc::new(InMemoryEpisodic::default());
    let (_wrapper, rx) = ProjectableEpisodic::new(inner);

    let cancel = CancellationToken::new();
    cancel.cancel();
    let handle = Arc::clone(&projector).spawn(rx, cancel);
    handle
        .await
        .expect("spawn must exit cleanly on pre-cancelled token");
    assert!(graph.snapshot().is_empty(), "no events; no index calls");
}

#[tokio::test]
async fn same_kind_episodes_in_one_run_get_unique_fact_ids() {
    let graph = Arc::new(RecordingGraph::default());
    let extractor = Arc::new(HintsOnlyExtractor::default());
    let metrics = Arc::new(RecallMetrics::default());
    let projector = EpisodeProjector::new(
        Arc::clone(&graph) as Arc<dyn KnowledgeGraph>,
        extractor,
        workspace_scope(),
        metrics,
    );

    let episodic = InMemoryEpisodic::default();
    let run = RunId::new();
    for name in ["search", "lookup"] {
        episodic
            .record(
                run,
                Episode::ToolCall {
                    name: name.into(),
                    args: serde_json::json!({}),
                    result: klieo_core::memory::ToolResult::ok(serde_json::json!("ok")),
                },
            )
            .await
            .unwrap();
    }

    projector.project_run(&episodic, run).await.unwrap();

    let indexed = graph.snapshot();
    assert_eq!(indexed.len(), 2);
    assert_ne!(
        indexed[0].fact_id, indexed[1].fact_id,
        "FactIds must be unique across same-kind episodes in one run",
    );
    let first = format!("{}", indexed[0].fact_id);
    let second = format!("{}", indexed[1].fact_id);
    assert!(first.contains("-toolcall-0"), "first FactId: {first}");
    assert!(second.contains("-toolcall-1"), "second FactId: {second}");
}
