pe-core 0.1.0

//! Synthesizer — merges lobe outputs into enriched context.
//!
//! The synthesizer is the final step in the cognitive graph. It receives
//! outputs from all active lobes and produces a single enriched context
//! string that gets prepended to the LLM call.
//!
//! The library provides three preset synthesizers. Users can implement
//! the [`Synthesizer`] trait for custom synthesis strategies.

use crate::cognitive_signal::CognitiveSignal;
use crate::error::PeError;
use crate::lobe::LobeOutput;
use serde::{Deserialize, Serialize};
use std::cmp::Ordering;
use std::future::Future;
use std::pin::Pin;

/// Trait for combining lobe outputs into enriched context.
///
/// # Example
///
/// ```ignore
/// struct MySynthesizer;
///
/// impl Synthesizer for MySynthesizer {
///     fn synthesize<'a>(
///         &'a self,
///         outputs: &'a [LobeOutput],
///         input: &'a str,
///     ) -> SynthesisFuture<'a> {
///         Box::pin(async move {
///             let merged = outputs.iter().map(|o| o.content.as_str()).collect::<Vec<_>>().join("\n");
///             Ok(SynthesisResult::new(merged))
///         })
///     }
/// }
/// ```
pub trait Synthesizer: Send + Sync {
    /// Merge lobe outputs into a single enriched context.
    fn synthesize<'a>(&'a self, outputs: &'a [LobeOutput], input: &'a str) -> SynthesisFuture<'a>;
}

/// Future type for synthesis — async, Send, returns Result.
pub type SynthesisFuture<'a> =
    Pin<Box<dyn Future<Output = Result<SynthesisResult, PeError>> + Send + 'a>>;

/// Result from synthesis — enriched context + metadata.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SynthesisResult {
    /// The enriched context string to prepend to the LLM call.
    pub enriched_context: String,

    /// Overall confidence from synthesis (average of lobe confidences).
    pub confidence: f64,

    /// Aggregated signals from all lobes.
    #[serde(default)]
    pub signals: Vec<CognitiveSignal>,

    /// Detected conflicts between lobe outputs.
    #[serde(default)]
    pub conflicts: Vec<String>,
}

impl SynthesisResult {
    /// Create a new synthesis result with enriched context and confidence.
    pub fn new(enriched_context: impl Into<String>, confidence: f64) -> Self {
        Self {
            enriched_context: enriched_context.into(),
            confidence,
            signals: Vec::new(),
            conflicts: Vec::new(),
        }
    }
}

/// Concatenates all lobe outputs into a structured prompt section.
///
/// Each lobe's output is labeled with its name. Simple, deterministic,
/// no additional LLM call needed.
pub struct IntegrateSynthesizer;

impl Synthesizer for IntegrateSynthesizer {
    fn synthesize<'a>(&'a self, outputs: &'a [LobeOutput], _input: &'a str) -> SynthesisFuture<'a> {
        Box::pin(async move {
            if outputs.is_empty() {
                return Ok(SynthesisResult::new("", 0.0));
            }

            let mut sections = Vec::with_capacity(outputs.len());
            let mut all_signals = Vec::new();
            let mut total_confidence = 0.0;

            for output in outputs {
                if output.lobe_name.is_empty() {
                    sections.push(output.content.clone());
                } else {
                    sections.push(format!("[{}]\n{}", output.lobe_name, output.content));
                }
                all_signals.extend(output.signals.iter().cloned());
                total_confidence += output.confidence;
            }

            let avg_confidence = total_confidence / outputs.len() as f64;
            let enriched = sections.join("\n\n");

            Ok(SynthesisResult {
                enriched_context: enriched,
                confidence: avg_confidence,
                signals: all_signals,
                conflicts: Vec::new(),
            })
        })
    }
}

/// Weights lobe outputs by their confidence scores.
///
/// Higher-confidence lobes get more representation in the final output.
/// Outputs below 0.3 confidence are noted as low-confidence.
pub struct WeightedSynthesizer;

fn compare_confidence_desc(a: f64, b: f64) -> Ordering {
    sanitize_confidence(b).total_cmp(&sanitize_confidence(a))
}

fn compare_confidence_asc(a: f64, b: f64) -> Ordering {
    sanitize_confidence(a).total_cmp(&sanitize_confidence(b))
}

fn sanitize_confidence(value: f64) -> f64 {
    if value.is_nan() {
        f64::NEG_INFINITY
    } else {
        value
    }
}

impl Synthesizer for WeightedSynthesizer {
    fn synthesize<'a>(&'a self, outputs: &'a [LobeOutput], _input: &'a str) -> SynthesisFuture<'a> {
        Box::pin(async move {
            if outputs.is_empty() {
                return Ok(SynthesisResult::new("", 0.0));
            }

            let mut sorted: Vec<_> = outputs.iter().collect();
            sorted.sort_by(|a, b| compare_confidence_desc(a.confidence, b.confidence));

            let mut sections = Vec::new();
            let mut all_signals = Vec::new();
            let mut total_confidence = 0.0;

            for output in &sorted {
                sections.push(output.content.clone());
                all_signals.extend(output.signals.iter().cloned());
                total_confidence += output.confidence;
            }

            // Simple average — content ordering already reflects confidence weighting
            let avg_confidence = total_confidence / sorted.len() as f64;

            Ok(SynthesisResult {
                enriched_context: sections.join("\n\n"),
                confidence: avg_confidence,
                signals: all_signals,
                conflicts: Vec::new(),
            })
        })
    }
}

/// Majority vote — lobes produce discrete choices, most common wins.
///
/// Treats each lobe output as a "vote". The most frequent content wins.
/// Useful when lobes produce categorical outputs (yes/no, option A/B/C).
pub struct VoteSynthesizer;

impl Synthesizer for VoteSynthesizer {
    fn synthesize<'a>(&'a self, outputs: &'a [LobeOutput], _input: &'a str) -> SynthesisFuture<'a> {
        Box::pin(async move {
            if outputs.is_empty() {
                return Ok(SynthesisResult::new("", 0.0));
            }

            let mut votes: std::collections::HashMap<&str, (usize, f64)> =
                std::collections::HashMap::new();
            let mut all_signals = Vec::new();

            for output in outputs {
                let entry = votes.entry(output.content.as_str()).or_insert((0, 0.0));
                entry.0 += 1;
                entry.1 += output.confidence;
                all_signals.extend(output.signals.iter().cloned());
            }

            // Deterministic tie-breaking: vote count → accumulated confidence → lexicographic
            let winner_content = votes
                .iter()
                .max_by(|(key_a, (count_a, conf_a)), (key_b, (count_b, conf_b))| {
                    count_a
                        .cmp(count_b)
                        .then(compare_confidence_asc(*conf_a, *conf_b))
                        .then(key_a.cmp(key_b))
                })
                .map(|(content, _)| *content)
                .unwrap_or("");

            let total_votes = outputs.len();
            let winner_votes = votes.get(winner_content).map(|(c, _)| *c).unwrap_or(0);
            let consensus = winner_votes as f64 / total_votes as f64;

            let mut conflicts = Vec::new();
            if consensus < 1.0 {
                let dissent: Vec<_> = votes
                    .keys()
                    .filter(|k| **k != winner_content)
                    .map(|k| k.to_string())
                    .collect();
                if !dissent.is_empty() {
                    conflicts.push(format!("Dissenting views: {}", dissent.join(", ")));
                }
            }

            Ok(SynthesisResult {
                enriched_context: winner_content.to_string(),
                confidence: consensus,
                signals: all_signals,
                conflicts,
            })
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn mock_outputs() -> Vec<LobeOutput> {
        vec![
            LobeOutput::new("Facts: X is true", 0.9),
            LobeOutput::new("Risk: Y might fail", 0.7),
        ]
    }

    #[tokio::test]
    async fn test_integrate_synthesizer() {
        let synth = IntegrateSynthesizer;
        let result = synth.synthesize(&mock_outputs(), "test").await.unwrap();
        assert!(result.enriched_context.contains("Facts: X is true"));
        assert!(result.enriched_context.contains("Risk: Y might fail"));
        assert!((result.confidence - 0.8).abs() < f64::EPSILON);
    }

    #[tokio::test]
    async fn test_integrate_empty() {
        let synth = IntegrateSynthesizer;
        let result = synth.synthesize(&[], "test").await.unwrap();
        assert!(result.enriched_context.is_empty());
    }

    #[tokio::test]
    async fn test_weighted_synthesizer_orders_by_confidence() {
        let synth = WeightedSynthesizer;
        let result = synth.synthesize(&mock_outputs(), "test").await.unwrap();
        // Higher confidence output should appear first
        let pos_facts = result.enriched_context.find("Facts").unwrap_or(usize::MAX);
        let pos_risk = result.enriched_context.find("Risk").unwrap_or(usize::MAX);
        assert!(pos_facts < pos_risk);
    }

    #[tokio::test]
    async fn test_vote_synthesizer_majority() {
        let synth = VoteSynthesizer;
        let outputs = vec![
            LobeOutput::new("yes", 0.8),
            LobeOutput::new("yes", 0.6),
            LobeOutput::new("no", 0.9),
        ];
        let result = synth.synthesize(&outputs, "test").await.unwrap();
        assert_eq!(result.enriched_context, "yes");
        assert!((result.confidence - 2.0 / 3.0).abs() < f64::EPSILON);
        assert!(!result.conflicts.is_empty());
    }

    #[tokio::test]
    async fn test_vote_synthesizer_unanimous() {
        let synth = VoteSynthesizer;
        let outputs = vec![
            LobeOutput::new("approve", 0.9),
            LobeOutput::new("approve", 0.8),
        ];
        let result = synth.synthesize(&outputs, "test").await.unwrap();
        assert_eq!(result.enriched_context, "approve");
        assert!((result.confidence - 1.0).abs() < f64::EPSILON);
        assert!(result.conflicts.is_empty());
    }

    #[tokio::test]
    async fn test_weighted_synthesizer_handles_nan_confidence() {
        let synth = WeightedSynthesizer;
        let outputs = vec![
            LobeOutput::new("bad sensor", f64::NAN),
            LobeOutput::new("stable output", 0.9),
        ];
        let result = synth.synthesize(&outputs, "test").await.unwrap();
        let pos_stable = result
            .enriched_context
            .find("stable output")
            .unwrap_or(usize::MAX);
        let pos_bad = result
            .enriched_context
            .find("bad sensor")
            .unwrap_or(usize::MAX);
        assert!(pos_stable < pos_bad);
    }

    #[tokio::test]
    async fn test_vote_synthesizer_handles_nan_tiebreak_confidence() {
        let synth = VoteSynthesizer;
        let outputs = vec![
            LobeOutput::new("approve", f64::NAN),
            LobeOutput::new("reject", 0.6),
        ];
        let result = synth.synthesize(&outputs, "test").await.unwrap();
        assert_eq!(result.enriched_context, "reject");
    }

    #[tokio::test]
    async fn test_signals_aggregated() {
        let synth = IntegrateSynthesizer;
        let outputs = vec![
            LobeOutput::new("ok", 0.9).with_signal(CognitiveSignal::Proceed),
            LobeOutput::new("careful", 0.5).with_signal(CognitiveSignal::ProceedWithCaution {
                concern: "low confidence".into(),
            }),
        ];
        let result = synth.synthesize(&outputs, "test").await.unwrap();
        assert_eq!(result.signals.len(), 2);
    }
}