//! Python bindings for TokenBufferMemory and SummaryMemory

use pyo3::prelude::*;
use std::pin::Pin;

use flowgentra_ai::core::memory::{
    ConversationMemory, TokenBufferMemory, SummaryMemory, SummaryConfig,
};

use crate::error::to_py_err;
use crate::llm::PyMessage;

// ─── PyTokenBufferMemory ───────────────────────────────────────────────────

/// Token-aware conversation memory that trims to a token budget.
///
/// Example:
///     mem = TokenBufferMemory(max_tokens=4000)
///     mem.add_message("thread-1", Message.user("Hello"))
///     messages = mem.messages("thread-1")
#[pyclass(name = "TokenBufferMemory")]
pub struct PyTokenBufferMemory {
    inner: TokenBufferMemory,
}

#[pymethods]
impl PyTokenBufferMemory {
    #[new]
    #[pyo3(signature = (max_tokens=4000))]
    fn new(max_tokens: usize) -> Self {
        PyTokenBufferMemory {
            inner: TokenBufferMemory::new(max_tokens),
        }
    }

    /// Add a message to a thread.
    fn add_message(&self, thread_id: &str, message: &PyMessage) -> PyResult<()> {
        self.inner
            .add_message(thread_id, message.inner.clone())
            .map_err(to_py_err)
    }

    /// Get messages for a thread.
    #[pyo3(signature = (thread_id, limit=None))]
    fn messages(&self, thread_id: &str, limit: Option<usize>) -> PyResult<Vec<PyMessage>> {
        let msgs = self.inner.messages(thread_id, limit).map_err(to_py_err)?;
        Ok(msgs.into_iter().map(|m| PyMessage { inner: m }).collect())
    }

    /// Clear all messages for a thread.
    fn clear(&self, thread_id: &str) -> PyResult<()> {
        self.inner.clear(thread_id).map_err(to_py_err)
    }

    fn __repr__(&self) -> String {
        "TokenBufferMemory(...)".to_string()
    }
}

// ─── PySummaryConfig ───────────────────────────────────────────────────────

/// Configuration for SummaryMemory.
///
/// Example:
///     config = SummaryConfig(buffer_size=10, max_summary_tokens=200)
#[pyclass(name = "SummaryConfig")]
#[derive(Clone)]
pub struct PySummaryConfig {
    pub(crate) inner: SummaryConfig,
}

#[pymethods]
impl PySummaryConfig {
    #[new]
    #[pyo3(signature = (buffer_size=10, max_summary_tokens=200))]
    fn new(buffer_size: usize, max_summary_tokens: usize) -> Self {
        PySummaryConfig {
            inner: SummaryConfig {
                buffer_size,
                max_summary_tokens,
            },
        }
    }

    #[getter]
    fn buffer_size(&self) -> usize {
        self.inner.buffer_size
    }

    #[getter]
    fn max_summary_tokens(&self) -> usize {
        self.inner.max_summary_tokens
    }

    fn __repr__(&self) -> String {
        format!(
            "SummaryConfig(buffer_size={}, max_summary_tokens={})",
            self.inner.buffer_size, self.inner.max_summary_tokens
        )
    }
}

// ─── PySummaryMemory ───────────────────────────────────────────────────────

// Type alias for the summarize function used by SummaryMemory
type SummarizeFn = Box<
    dyn Fn(String) -> Pin<Box<dyn std::future::Future<Output = Result<String, flowgentra_ai::core::error::FlowgentraError>> + Send>>
        + Send
        + Sync,
>;

/// Conversation memory with automatic summarization.
///
/// When the message buffer exceeds `buffer_size`, older messages are
/// summarized using the provided function.
///
/// Example:
///     def summarize(text):
///         return "Summary: " + text[:100]
///
///     mem = SummaryMemory(SummaryConfig(), summarize)
///     mem.add_message("thread-1", Message.user("Hello"))
#[pyclass(name = "SummaryMemory")]
pub struct PySummaryMemory {
    inner: SummaryMemory<SummarizeFn>,
}

#[pymethods]
impl PySummaryMemory {
    #[new]
    fn new(config: &PySummaryConfig, summarize_fn: PyObject) -> Self {
        let func = Python::with_gil(|py| summarize_fn.clone_ref(py));

        let rust_fn: SummarizeFn = Box::new(move |text: String| {
            let func = Python::with_gil(|py| func.clone_ref(py));
            Box::pin(async move {
                let result = Python::with_gil(|py| -> Result<String, flowgentra_ai::core::error::FlowgentraError> {
                    let res = func.call1(py, (text,)).map_err(|e| {
                        flowgentra_ai::core::error::FlowgentraError::ToolError(format!("summarize_fn error: {}", e))
                    })?;
                    let s: String = res.extract(py).map_err(|e| {
                        flowgentra_ai::core::error::FlowgentraError::ToolError(format!("summarize_fn must return str: {}", e))
                    })?;
                    Ok(s)
                });
                result
            })
        });

        PySummaryMemory {
            inner: SummaryMemory::new(config.inner.clone(), rust_fn),
        }
    }

    /// Add a message to a thread.
    fn add_message(&self, thread_id: &str, message: &PyMessage) -> PyResult<()> {
        self.inner
            .add_message(thread_id, message.inner.clone())
            .map_err(to_py_err)
    }

    /// Get messages for a thread.
    #[pyo3(signature = (thread_id, limit=None))]
    fn messages(&self, thread_id: &str, limit: Option<usize>) -> PyResult<Vec<PyMessage>> {
        let msgs = self.inner.messages(thread_id, limit).map_err(to_py_err)?;
        Ok(msgs.into_iter().map(|m| PyMessage { inner: m }).collect())
    }

    /// Clear all messages for a thread.
    fn clear(&self, thread_id: &str) -> PyResult<()> {
        self.inner.clear(thread_id).map_err(to_py_err)
    }

    /// Get the current summary for a thread (if any).
    fn get_summary(&self, thread_id: &str) -> Option<String> {
        self.inner.get_summary(thread_id)
    }

    /// Trigger summarization if the buffer exceeds the configured size.
    fn summarize_if_needed(&self, thread_id: &str) -> PyResult<()> {
        let rt = crate::get_runtime();
        rt.block_on(self.inner.summarize_if_needed(thread_id))
            .map_err(to_py_err)
    }

    fn __repr__(&self) -> String {
        "SummaryMemory(...)".to_string()
    }
}
