//! Python bindings for FlowgentraAI via PyO3
//!
//! Exposes the core Rust library to Python through native extension module.

use pyo3::prelude::*;
use pyo3::types::{PyDict, PyList};
use serde_json::Value;

mod state;
mod agent;
mod agents;
mod llm;
mod llm_client;
mod config;
mod graph;
mod rag;
mod vector_store;
mod memory;
mod conversation_memory;
mod error;
mod tools;
mod tool_registry;
mod evaluation;
mod text_splitter;
mod prompt_parser;
mod message_graph;
mod supervisor;
mod routing;
mod observability;
mod mcp;
mod document_loader;
mod ingestion;
mod rag_eval;
mod reranker;
mod token_buffer_memory;
mod file_checkpointer;
mod chroma;
mod builtin_tools;
mod tool_node;
mod visualization;
mod advanced_nodes;
mod llm_reranker;
mod memory_aware_agent;
mod rag_config;
mod eval_config;
mod remaining;

use state::*;
use agent::*;
use agents::*;
use llm::*;
use llm_client::*;
use config::*;
use graph::*;
use rag::*;
use vector_store::*;
use memory::*;
use conversation_memory::*;
use tools::*;
use tool_registry::*;
use evaluation::*;
use text_splitter::*;
use prompt_parser::*;
use message_graph::*;
use supervisor::*;
use routing::*;
use observability::*;
use mcp::*;
use document_loader::*;
use ingestion::*;
use rag_eval::*;
use reranker::*;
use token_buffer_memory::*;
use file_checkpointer::*;
use chroma::*;
use builtin_tools::*;
use tool_node::*;
use visualization::*;
use advanced_nodes::*;
use llm_reranker::*;
use memory_aware_agent::*;
use rag_config::*;
use eval_config::*;
use remaining::*;

// ─── Helpers ────────────────────────────────────────────────────────────────

/// Convert a serde_json::Value to a Python object
pub fn json_to_py(py: Python<'_>, val: &Value) -> PyResult<PyObject> {
    match val {
        Value::Null => Ok(py.None()),
        Value::Bool(b) => Ok(b.to_object(py)),
        Value::Number(n) => {
            if let Some(i) = n.as_i64() {
                Ok(i.to_object(py))
            } else if let Some(f) = n.as_f64() {
                Ok(f.to_object(py))
            } else {
                Ok(py.None())
            }
        }
        Value::String(s) => Ok(s.to_object(py)),
        Value::Array(arr) => {
            let list = PyList::empty_bound(py);
            for item in arr {
                list.append(json_to_py(py, item)?)?;
            }
            Ok(list.into())
        }
        Value::Object(map) => {
            let dict = PyDict::new_bound(py);
            for (k, v) in map {
                dict.set_item(k, json_to_py(py, v)?)?;
            }
            Ok(dict.into())
        }
    }
}

/// Convert a Python object to serde_json::Value
pub fn py_to_json(obj: &Bound<'_, PyAny>) -> PyResult<Value> {
    if obj.is_none() {
        Ok(Value::Null)
    } else if let Ok(b) = obj.extract::<bool>() {
        Ok(Value::Bool(b))
    } else if let Ok(i) = obj.extract::<i64>() {
        Ok(Value::Number(i.into()))
    } else if let Ok(f) = obj.extract::<f64>() {
        Ok(serde_json::Number::from_f64(f)
            .map(Value::Number)
            .unwrap_or(Value::Null))
    } else if let Ok(s) = obj.extract::<String>() {
        Ok(Value::String(s))
    } else if let Ok(list) = obj.downcast::<PyList>() {
        let mut arr = Vec::new();
        for item in list.iter() {
            arr.push(py_to_json(&item)?);
        }
        Ok(Value::Array(arr))
    } else if let Ok(dict) = obj.downcast::<PyDict>() {
        let mut map = serde_json::Map::new();
        for (k, v) in dict.iter() {
            let key: String = k.extract()?;
            map.insert(key, py_to_json(&v)?);
        }
        Ok(Value::Object(map))
    } else {
        // Fallback: try string representation
        let s = obj.str()?.to_string();
        Ok(Value::String(s))
    }
}

/// Get or create a tokio runtime for blocking async calls
pub fn get_runtime() -> &'static tokio::runtime::Runtime {
    use std::sync::OnceLock;
    static RT: OnceLock<tokio::runtime::Runtime> = OnceLock::new();
    RT.get_or_init(|| {
        tokio::runtime::Runtime::new().expect("Failed to create Tokio runtime")
    })
}

// ─── Python Module ──────────────────────────────────────────────────────────

/// FlowgentraAI Python module
#[pymodule]
fn _native(m: &Bound<'_, PyModule>) -> PyResult<()> {
    // ── State ──
    m.add_class::<PySharedState>()?;
    m.add_class::<PyPlainState>()?;

    // ── Agent (config-driven) ──
    m.add_class::<PyAgent>()?;

    // ── Prebuilt Agents ──
    m.add_class::<PyAgentType>()?;
    m.add_class::<PyToolSpec>()?;
    m.add_class::<PyAgentBuilder>()?;
    m.add_class::<PyGraphBasedAgent>()?;

    // ── Config ──
    m.add_class::<PyAgentConfig>()?;
    m.add_class::<PyStateField>()?;

    // ── Graph (code-driven) ──
    m.add_class::<PyStateGraph>()?;
    m.add_class::<PyStateGraphBuilder>()?;
    m.add("END", graph::PY_END)?;

    // ── LLM ──
    m.add_class::<PyLLMConfig>()?;
    m.add_class::<PyResponseFormat>()?;
    m.add_class::<PyMessage>()?;
    m.add_class::<PyToolCall>()?;
    m.add_class::<PyToolDefinition>()?;
    m.add_class::<PyTokenUsage>()?;
    m.add_class::<PyLLMClient>()?;

    // ── Tools & Registry ──
    m.add_class::<PyToolCallRequest>()?;
    m.add_class::<PyToolCallResult>()?;
    m.add_class::<PyToolRegistry>()?;
    m.add_class::<PyJsonSchema>()?;

    // ── RAG ──
    m.add_class::<PyDocument>()?;
    m.add_class::<PySearchResult>()?;
    m.add_class::<PyTextChunk>()?;

    // ── Vector Store & Embeddings ──
    m.add_class::<PyEmbeddings>()?;
    m.add_class::<PyInMemoryVectorStore>()?;
    m.add_class::<PyRetrievalConfig>()?;
    m.add_class::<PyRetriever>()?;

    // ── Text Splitters ──
    m.add_class::<PyRecursiveCharacterTextSplitter>()?;
    m.add_class::<PyMarkdownTextSplitter>()?;
    m.add_class::<PyHTMLTextSplitter>()?;
    m.add_class::<PyTokenTextSplitter>()?;
    m.add_class::<PyCodeTextSplitter>()?;

    // ── Memory ──
    m.add_class::<PyCheckpoint>()?;
    m.add_class::<PyCheckpointMetadata>()?;
    m.add_class::<PyConversationMemory>()?;

    // ── Evaluation ──
    m.add_class::<PyEvaluationResult>()?;

    // ── Prompt & Parsers ──
    m.add_class::<PyPromptTemplate>()?;
    m.add_class::<PyJsonOutputParser>()?;
    m.add_class::<PyListOutputParser>()?;

    // ── Message Graph ──
    m.add_class::<PyMessageGraph>()?;
    m.add_class::<PyMessageGraphBuilder>()?;

    // ── Supervisor ──
    m.add_class::<PySupervisor>()?;

    // ── Routing Conditions ──
    m.add_class::<PyComparisonOp>()?;
    m.add_class::<PyCondition>()?;
    m.add_class::<PyConditionBuilder>()?;

    // ── Observability ──
    m.add_class::<PyExecutionTrace>()?;
    m.add_class::<PyExecutionTracer>()?;

    // ── MCP ──
    m.add_class::<PyMCPConfig>()?;

    // ── Document Loaders ──
    m.add_class::<PyFileType>()?;
    m.add_class::<PyLoadedDocument>()?;

    // ── Ingestion ──
    m.add_class::<PyIngestionPipeline>()?;
    m.add_class::<PyIngestionStats>()?;

    // ── RAG Evaluation ──
    m.add_class::<PyEvalQuery>()?;
    m.add_class::<PyEvalResults>()?;
    m.add_class::<PyQueryResult>()?;

    // ── PDF ──
    m.add_class::<rag::PyPdfDocument>()?;

    // ── Rerankers ──
    m.add_class::<PyNoopReranker>()?;
    m.add_class::<PyRRFReranker>()?;
    m.add_class::<PyCrossEncoderReranker>()?;

    // ── Advanced Memory ──
    m.add_class::<PyTokenBufferMemory>()?;
    m.add_class::<PySummaryConfig>()?;
    m.add_class::<PySummaryMemory>()?;

    // ── File Checkpointer ──
    m.add_class::<PyFileCheckpointer>()?;

    // ── ChromaDB ──
    m.add_class::<PyChromaStore>()?;

    // ── Built-in Tools ──
    m.add_class::<PyCalculatorTool>()?;
    m.add_class::<PySearchTool>()?;
    m.add_class::<PyWebRequestTool>()?;
    m.add_class::<PyFilesTool>()?;

    // ── Tool Node ──
    m.add_class::<PyToolNode>()?;

    // ── Visualization ──
    m.add_class::<PyVisualizationConfig>()?;

    // ── Advanced Nodes ──
    m.add_class::<PyJoinType>()?;
    m.add_class::<PyMergeStrategy>()?;
    m.add_class::<PyBranchConfig>()?;
    m.add_class::<PyLoopNodeConfig>()?;
    m.add_class::<PyParallelNodeConfig>()?;
    m.add_class::<PySubgraphNodeConfig>()?;
    m.add_class::<PyJoinNodeConfig>()?;

    // ── LLM Reranker ──
    m.add_class::<PyLLMReranker>()?;

    // ── Memory-Aware Agent ──
    m.add_class::<PyMemoryAwareAgent>()?;
    m.add_class::<PyMemoryStats>()?;

    // ── RAG Config Types ──
    m.add_class::<PyVectorStoreType>()?;
    m.add_class::<PyRAGConfig>()?;
    m.add_class::<PyVectorStoreConfig>()?;
    m.add_class::<PyEmbeddingsConfig>()?;
    m.add_class::<PyRetrievalSettings>()?;
    m.add_class::<PyPdfSettings>()?;
    m.add_class::<PyRAGGraphConfig>()?;

    // ── Evaluation Config & Reporting ──
    m.add_class::<PyScoringConfig>()?;
    m.add_class::<PyGradingConfig>()?;
    m.add_class::<PyEvaluationConfig>()?;
    m.add_class::<PyNodeResult>()?;
    m.add_class::<PyEvaluationReport>()?;

    // ── Remaining Types ──
    m.add_class::<PyChunkMetadata>()?;
    m.add_class::<PyRetrieverStrategy>()?;
    m.add_class::<PyRerankStrategy>()?;
    m.add_class::<PyVectorStore>()?;

    // ── Free functions ──
    m.add_function(wrap_pyfunction!(py_from_config_path, m)?)?;
    m.add_function(wrap_pyfunction!(py_chunk_text, m)?)?;
    m.add_function(wrap_pyfunction!(py_extract_text, m)?)?;
    m.add_function(wrap_pyfunction!(py_estimate_tokens, m)?)?;
    m.add_function(wrap_pyfunction!(py_model_pricing, m)?)?;
    m.add_function(wrap_pyfunction!(graph::graph_end, m)?)?;
    m.add_function(wrap_pyfunction!(observability::py_init_tracing, m)?)?;

    // Document loader functions
    m.add_function(wrap_pyfunction!(document_loader::py_load_document, m)?)?;
    m.add_function(wrap_pyfunction!(document_loader::py_load_directory, m)?)?;

    // RAG evaluation functions
    m.add_function(wrap_pyfunction!(rag_eval::py_hit_rate, m)?)?;
    m.add_function(wrap_pyfunction!(rag_eval::py_mrr, m)?)?;
    m.add_function(wrap_pyfunction!(rag_eval::py_mean_ndcg, m)?)?;
    m.add_function(wrap_pyfunction!(rag_eval::py_rag_evaluate, m)?)?;

    // PDF / chunk functions
    m.add_function(wrap_pyfunction!(rag::py_chunk_text_by_tokens, m)?)?;
    m.add_function(wrap_pyfunction!(rag::py_extract_and_chunk, m)?)?;
    m.add_function(wrap_pyfunction!(rag::py_extract_pdf, m)?)?;

    // Hybrid search & dedup functions
    m.add_function(wrap_pyfunction!(rag::py_bm25_score, m)?)?;
    m.add_function(wrap_pyfunction!(rag::py_hybrid_merge, m)?)?;
    m.add_function(wrap_pyfunction!(rag::py_dedup_by_id, m)?)?;
    m.add_function(wrap_pyfunction!(rag::py_dedup_by_similarity, m)?)?;
    m.add_function(wrap_pyfunction!(rag::py_decompose_query, m)?)?;

    // Tool node functions
    m.add_function(wrap_pyfunction!(tool_node::py_create_tool_node, m)?)?;
    m.add_function(wrap_pyfunction!(tool_node::py_store_tool_calls, m)?)?;
    m.add_function(wrap_pyfunction!(tool_node::py_check_tools_condition, m)?)?;

    // Evaluation utility
    m.add_function(wrap_pyfunction!(remaining::py_evaluate_output_score, m)?)?;

    // Visualization functions
    m.add_function(wrap_pyfunction!(visualization::py_visualize_graph, m)?)?;
    m.add_function(wrap_pyfunction!(visualization::py_graph_to_dot, m)?)?;
    m.add_function(wrap_pyfunction!(visualization::py_graph_to_mermaid, m)?)?;

    Ok(())
}
