//! RAG Pipeline core traits
//!
//! Defines the abstract interfaces for the RAG pipeline components:
//!
//! - **Retriever**: Finds relevant documents for a query
//! - **Generator**: Produces natural language answers from context
//! - **ContextAssembler**: Assembles retrieval results into a prompt
//!
//! Note: For embedding, we reuse the existing `PerceptionEngine` trait
//! from `crate::perception::traits`.

use super::error::RagResult;
use super::types::{ChatMessage, RagConfig, RagResponse, RetrievalResult};

// ============================================================================
// Retriever Trait
// ============================================================================

/// Retrieves relevant documents for a given query.
///
/// Implementations can combine vector search, full-text search, metadata filters,
/// and reranking to find the most relevant content.
///
/// # Example
///
/// ```ignore
/// struct MyRetriever { /* ... */ }
///
/// impl Retriever for MyRetriever {
///     fn retrieve(&self, query: &str, embedding: Option<&[f32]>, config: &RagConfig)
///         -> RagResult<Vec<RetrievalResult>>
///     {
///         // Perform hybrid search ...
///     }
/// }
/// ```
pub trait Retriever: Send + Sync {
    /// Retrieve relevant documents.
    ///
    /// # Arguments
    /// * `query` - The natural language query text
    /// * `embedding` - Pre-computed query embedding (if available; retriever may compute its own)
    /// * `config` - RAG configuration controlling search parameters
    ///
    /// # Returns
    /// A list of retrieval results sorted by relevance (highest first).
    fn retrieve(
        &self,
        query: &str,
        embedding: Option<&[f32]>,
        config: &RagConfig,
    ) -> RagResult<Vec<RetrievalResult>>;
}

// ============================================================================
// Generator Trait
// ============================================================================

/// Generates natural language answers from context and conversation history.
///
/// Implementations can use local models (MicroGPT) or external LLM APIs.
pub trait Generator: Send + Sync {
    /// Generate an answer given assembled messages.
    ///
    /// # Arguments
    /// * `messages` - The full message list (system + history + user), already assembled
    /// * `config` - RAG configuration controlling generation parameters
    ///
    /// # Returns
    /// A `RagResponse` containing the answer, token usage, etc.
    fn generate(
        &self,
        messages: &[ChatMessage],
        config: &RagConfig,
    ) -> RagResult<RagResponse>;
}

// Blanket impl: Arc<dyn Generator> also implements Generator,
// so shared generators can be used directly with the pipeline builder.
impl Generator for std::sync::Arc<dyn Generator> {
    fn generate(
        &self,
        messages: &[ChatMessage],
        config: &RagConfig,
    ) -> RagResult<RagResponse> {
        (**self).generate(messages, config)
    }
}

// ============================================================================
// StreamGenerator Trait (Phase-82: T82.2)
// ============================================================================

/// Streaming generator — async variant of [`Generator`].
///
/// Instead of returning a complete `RagResponse`, this trait returns a
/// `Stream<Item = StreamEvent>` that yields incremental events:
///
/// - `TextDelta` — text content fragments
/// - `ThinkingDelta` — reasoning process fragments
/// - `Usage` — token usage information
/// - `Done` — completion signal
///
/// This enables real-time progressive output (SSE, gRPC streaming, WebSocket).
///
/// # Requires
///
/// - `llm-gateway` feature enabled
///
/// # Example
///
/// ```ignore
/// use aifastdb_core::rag::StreamGenerator;
/// use tokio_stream::StreamExt;
///
/// let stream = generator.generate_stream(&messages, &config).await?;
/// tokio::pin!(stream);
/// while let Some(event) = stream.next().await {
///     match event? {
///         StreamEvent::TextDelta { content } => print!("{}", content),
///         StreamEvent::Done { .. } => break,
///         _ => {}
///     }
/// }
/// ```
#[cfg(feature = "llm-gateway")]
#[async_trait::async_trait]
pub trait StreamGenerator: Send + Sync {
    /// Generate a streaming response given assembled messages.
    ///
    /// # Arguments
    /// * `messages` - The full message list (system + history + user), already assembled
    /// * `config` - RAG configuration controlling generation parameters
    ///
    /// # Returns
    /// A `Pin<Box<dyn Stream<...>>>` that yields `StreamEvent` items.
    /// The stream ends with `StreamEvent::Done { ... }` on success.
    async fn generate_stream(
        &self,
        messages: &[ChatMessage],
        config: &RagConfig,
    ) -> RagResult<
        std::pin::Pin<
            Box<
                dyn futures_core::Stream<Item = RagResult<crate::llm::types::StreamEvent>>
                    + Send,
            >,
        >,
    >;
}

// ============================================================================
// SessionManager Trait (Phase-82: T82.4)
// ============================================================================

/// Abstraction for RAG session (conversation history) management.
///
/// Implementations:
/// - [`InMemorySessionManager`](super::session::InMemorySessionManager) — ephemeral, in-memory (default)
/// - [`ConversationSessionAdapter`](super::session::ConversationSessionAdapter) — delegates to
///   `ConversationManager` for persistent, ownership-aware storage (requires `llm-gateway`)
///
/// # Phase-82 T82.4
///
/// Introduced to decouple `RagPipeline` from the concrete `InMemorySessionManager`,
/// enabling unified session persistence through `ConversationManager`.
pub trait SessionManager: Send + Sync {
    /// Get or create a session by ID (idempotent).
    fn get_or_create(&self, session_id: &str) -> RagResult<super::types::ChatSession>;

    /// Push a message to a session's history.
    fn push_message(
        &self,
        session_id: &str,
        message: super::types::ChatMessage,
    ) -> RagResult<()>;

    /// Get the most recent conversation turns (user + assistant pairs).
    fn get_recent_history(
        &self,
        session_id: &str,
        max_turns: usize,
    ) -> RagResult<Vec<super::types::ChatMessage>>;

    /// Get the full conversation history for a session.
    fn get_history(&self, session_id: &str) -> RagResult<Vec<super::types::ChatMessage>>;

    /// Delete a session.
    fn delete_session(&self, session_id: &str) -> RagResult<()>;

    /// Get the number of active sessions (optional, for monitoring).
    fn session_count(&self) -> usize {
        0 // Default: not all backends can count efficiently
    }
}

// ============================================================================
// ContextAssembler Trait
// ============================================================================

/// Assembles retrieval results and conversation history into a prompt.
///
/// Responsible for:
/// 1. Formatting retrieved documents into a context string
/// 2. Applying token budget constraints (truncating if necessary)
/// 3. Combining system prompt, context, history, and user question
pub trait ContextAssembler: Send + Sync {
    /// Assemble the final message list for the generator.
    ///
    /// # Arguments
    /// * `question` - The current user question
    /// * `retrieved` - Retrieved documents (already ranked)
    /// * `history` - Previous conversation turns
    /// * `config` - RAG configuration
    ///
    /// # Returns
    /// A list of `ChatMessage` ready for the generator, plus the number
    /// of context tokens used.
    fn assemble(
        &self,
        question: &str,
        retrieved: &[RetrievalResult],
        history: &[ChatMessage],
        config: &RagConfig,
    ) -> RagResult<(Vec<ChatMessage>, usize)>;
}

#[cfg(test)]
mod tests {
    use super::*;
    use super::super::types::*;
    #[allow(unused_imports)]
    use super::super::error::*;

    // ── Dummy implementations for trait-object safety tests ──

    struct DummyRetriever;

    impl Retriever for DummyRetriever {
        fn retrieve(
            &self,
            _query: &str,
            _embedding: Option<&[f32]>,
            config: &RagConfig,
        ) -> RagResult<Vec<RetrievalResult>> {
            let mut results = Vec::new();
            for i in 0..config.top_k.min(3) {
                results.push(RetrievalResult::new(
                    format!("doc-{}", i),
                    format!("Content for document {}", i),
                    1.0 - (i as f64 * 0.1),
                ));
            }
            Ok(results)
        }
    }

    struct DummyGenerator;

    impl Generator for DummyGenerator {
        fn generate(
            &self,
            messages: &[ChatMessage],
            _config: &RagConfig,
        ) -> RagResult<RagResponse> {
            let answer = format!(
                "Generated from {} messages",
                messages.len()
            );
            Ok(RagResponse::new(answer))
        }
    }

    struct DummyAssembler;

    impl ContextAssembler for DummyAssembler {
        fn assemble(
            &self,
            question: &str,
            retrieved: &[RetrievalResult],
            _history: &[ChatMessage],
            config: &RagConfig,
        ) -> RagResult<(Vec<ChatMessage>, usize)> {
            let context: String = retrieved
                .iter()
                .map(|r| r.content.clone())
                .collect::<Vec<_>>()
                .join("\n");
            let system = config.system_template.replace("{context}", &context);
            let user = config.user_template.replace("{question}", question);

            let messages = vec![
                ChatMessage::system(system),
                ChatMessage::user(user),
            ];
            let tokens = messages.iter().map(|m| m.approx_tokens()).sum();
            Ok((messages, tokens))
        }
    }

    #[test]
    fn test_trait_objects() {
        let retriever: Box<dyn Retriever> = Box::new(DummyRetriever);
        let generator: Box<dyn Generator> = Box::new(DummyGenerator);
        let assembler: Box<dyn ContextAssembler> = Box::new(DummyAssembler);

        let config = RagConfig::default();

        // Retrieval
        let results = retriever.retrieve("test query", None, &config).unwrap();
        assert!(!results.is_empty());

        // Assembly
        let (messages, tokens) = assembler
            .assemble("test query", &results, &[], &config)
            .unwrap();
        assert!(tokens > 0);
        assert!(!messages.is_empty());

        // Generation
        let response = generator.generate(&messages, &config).unwrap();
        assert!(!response.answer.is_empty());
    }
}
