//! Python bindings for text splitters

use pyo3::prelude::*;

use flowgentra_ai::core::rag::{
    RecursiveCharacterTextSplitter, MarkdownTextSplitter, HTMLTextSplitter,
    TokenTextSplitter, CodeTextSplitter, Language, TextSplitter, TextChunk,
};

use crate::rag::PyTextChunk;

// ─── Helper to convert Vec<TextChunk> → Vec<PyTextChunk> ───────────────────

fn to_py_chunks(chunks: Vec<TextChunk>) -> Vec<PyTextChunk> {
    chunks
        .into_iter()
        .map(|c| PyTextChunk { inner: c })
        .collect()
}

// ─── PyRecursiveCharacterTextSplitter ──────────────────────────────────────

/// Split text by recursively trying different separators.
///
/// Example:
///     splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
///     chunks = splitter.split_text("Very long document text...")
#[pyclass(name = "RecursiveCharacterTextSplitter")]
pub struct PyRecursiveCharacterTextSplitter {
    inner: RecursiveCharacterTextSplitter,
}

#[pymethods]
impl PyRecursiveCharacterTextSplitter {
    #[new]
    #[pyo3(signature = (chunk_size=1000, chunk_overlap=200))]
    fn new(chunk_size: usize, chunk_overlap: usize) -> Self {
        PyRecursiveCharacterTextSplitter {
            inner: RecursiveCharacterTextSplitter::new(chunk_size, chunk_overlap),
        }
    }

    /// Split text into chunks.
    fn split_text(&self, text: &str) -> Vec<PyTextChunk> {
        to_py_chunks(self.inner.split_text(text))
    }

    /// Split text with source metadata.
    fn split_with_source(&self, text: &str, source: &str) -> Vec<PyTextChunk> {
        to_py_chunks(self.inner.split_with_source(text, source))
    }

    fn __repr__(&self) -> String {
        format!(
            "RecursiveCharacterTextSplitter(chunk_size={}, overlap={})",
            self.inner.chunk_size, self.inner.chunk_overlap
        )
    }
}

// ─── PyMarkdownTextSplitter ────────────────────────────────────────────────

/// Split markdown text by headers and code blocks.
#[pyclass(name = "MarkdownTextSplitter")]
pub struct PyMarkdownTextSplitter {
    inner: MarkdownTextSplitter,
}

#[pymethods]
impl PyMarkdownTextSplitter {
    #[new]
    #[pyo3(signature = (chunk_size=1000, chunk_overlap=200))]
    fn new(chunk_size: usize, chunk_overlap: usize) -> Self {
        PyMarkdownTextSplitter {
            inner: MarkdownTextSplitter::new(chunk_size, chunk_overlap),
        }
    }

    fn split_text(&self, text: &str) -> Vec<PyTextChunk> {
        to_py_chunks(self.inner.split_text(text))
    }

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

// ─── PyHTMLTextSplitter ────────────────────────────────────────────────────

/// Split HTML text by tags.
#[pyclass(name = "HTMLTextSplitter")]
pub struct PyHTMLTextSplitter {
    inner: HTMLTextSplitter,
}

#[pymethods]
impl PyHTMLTextSplitter {
    #[new]
    #[pyo3(signature = (chunk_size=1000, chunk_overlap=200))]
    fn new(chunk_size: usize, chunk_overlap: usize) -> Self {
        PyHTMLTextSplitter {
            inner: HTMLTextSplitter::new(chunk_size, chunk_overlap),
        }
    }

    fn split_text(&self, text: &str) -> Vec<PyTextChunk> {
        to_py_chunks(self.inner.split_text(text))
    }

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

// ─── PyTokenTextSplitter ───────────────────────────────────────────────────

/// Split text by estimated token count.
#[pyclass(name = "TokenTextSplitter")]
pub struct PyTokenTextSplitter {
    inner: TokenTextSplitter,
}

#[pymethods]
impl PyTokenTextSplitter {
    #[new]
    #[pyo3(signature = (max_tokens=500, overlap_tokens=50))]
    fn new(max_tokens: usize, overlap_tokens: usize) -> Self {
        PyTokenTextSplitter {
            inner: TokenTextSplitter::new(max_tokens, overlap_tokens),
        }
    }

    fn split_text(&self, text: &str) -> Vec<PyTextChunk> {
        to_py_chunks(self.inner.split_text(text))
    }

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

// ─── PyCodeTextSplitter ────────────────────────────────────────────────────

/// Split source code by language-specific syntax.
///
/// Example:
///     splitter = CodeTextSplitter("python", chunk_size=500)
///     chunks = splitter.split_text(python_code)
#[pyclass(name = "CodeTextSplitter")]
pub struct PyCodeTextSplitter {
    inner: CodeTextSplitter,
}

#[pymethods]
impl PyCodeTextSplitter {
    /// Create a code text splitter for a given language.
    ///
    /// Supported languages: "python", "rust", "javascript", "typescript",
    /// "java", "go", "cpp", "c", "ruby", "swift", "kotlin"
    #[new]
    #[pyo3(signature = (language, chunk_size=1000, chunk_overlap=200))]
    fn new(language: &str, chunk_size: usize, chunk_overlap: usize) -> Self {
        let lang = match language.to_lowercase().as_str() {
            "python" => Language::Python,
            "rust" => Language::Rust,
            "javascript" | "js" => Language::JavaScript,
            "typescript" | "ts" => Language::TypeScript,
            "java" => Language::Java,
            "go" => Language::Go,
            "cpp" | "c++" => Language::Cpp,
            "csharp" | "c#" => Language::CSharp,
            "ruby" => Language::Ruby,
            _ => Language::Generic,
        };
        PyCodeTextSplitter {
            inner: CodeTextSplitter::new(chunk_size, chunk_overlap, lang),
        }
    }

    fn split_text(&self, text: &str) -> Vec<PyTextChunk> {
        to_py_chunks(self.inner.split_text(text))
    }

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