//! Python bindings for LLM Gateway (Phase-81: T81.1 + T81.2)
//!
//! Provides `PyLlmGateway` class with:
//! - Provider/Model registration and CRUD
//! - Conversation management
//! - Chat (non-streaming and streaming)
//! - Prompt template management
//! - Usage statistics

use std::collections::HashMap;
use std::sync::Arc;

use pyo3::prelude::*;
use pyo3::types::PyDict;
use pyo3::exceptions::{PyRuntimeError, PyValueError};

use aifastdb_core::llm::{
    LlmStore, ProviderRouter, ConversationManager,
    LlmProvider, LlmModel, ProviderBrand, ProviderProtocol,
    LlmConversation, ConversationQuery, MessageQuery, StoredMessage,
    PromptTemplate, PromptTemplateStore,
};

use crate::types::json_value_to_py;

/// Register LLM classes in the Python module.
pub fn register_llm(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<PyLlmGateway>()?;
    Ok(())
}

// ============================================================================
// PyLlmGateway
// ============================================================================

/// LLM Gateway — unified interface for LLM providers, conversations, and chat.
///
/// Usage:
/// ```python
/// from aifastdb import LlmGateway
///
/// gw = LlmGateway("/path/to/data")
/// gw.register_provider("openai", "OpenAI", "openai", "https://api.openai.com/v1", api_key="sk-...")
/// gw.register_model("gpt-4", "openai", "GPT-4", "gpt-4")
/// conv = gw.create_conversation("user1", "openai", "gpt-4", "gpt-4")
/// ```
#[pyclass(name = "LlmGateway")]
pub struct PyLlmGateway {
    store: Arc<LlmStore>,
    router: Arc<ProviderRouter>,
    conversations: Arc<ConversationManager>,
    prompts: Arc<PromptTemplateStore>,
    rt: tokio::runtime::Runtime,
}

#[pymethods]
impl PyLlmGateway {
    /// Create a new LLM Gateway.
    ///
    /// Args:
    ///     data_dir: Path to the data directory for persistent storage.
    #[new]
    fn new(data_dir: &str) -> PyResult<Self> {
        let dir = std::path::Path::new(data_dir);
        std::fs::create_dir_all(dir)
            .map_err(|e| PyRuntimeError::new_err(format!("Failed to create data dir: {e}")))?;

        let store = Arc::new(
            LlmStore::open(dir)
                .map_err(|e| PyRuntimeError::new_err(format!("LlmStore open failed: {e}")))?,
        );
        let router = Arc::new(ProviderRouter::default_router());
        let conversations = Arc::new(
            ConversationManager::new(dir, store.clone(), router.clone())
                .map_err(|e| PyRuntimeError::new_err(format!("ConversationManager init failed: {e}")))?,
        );
        let prompts = Arc::new(
            PromptTemplateStore::open(dir)
                .map_err(|e| PyRuntimeError::new_err(format!("PromptTemplateStore init failed: {e}")))?,
        );
        let rt = tokio::runtime::Runtime::new()
            .map_err(|e| PyRuntimeError::new_err(format!("Tokio runtime failed: {e}")))?;

        Ok(Self {
            store,
            router,
            conversations,
            prompts,
            rt,
        })
    }

    // ========================================================================
    // Provider CRUD
    // ========================================================================

    /// Register a new LLM provider.
    ///
    /// Args:
    ///     id: Unique provider ID (e.g., "openai")
    ///     name: Display name (e.g., "OpenAI")
    ///     brand: Provider brand ("openai", "anthropic", "gemini", "deepseek", "ollama", or custom)
    ///     base_url: API base URL
    ///     api_key: Optional API key
    ///     protocol: Optional protocol ("openai_compat", "anthropic", "gemini")
    #[pyo3(signature = (id, name, brand, base_url, api_key=None, protocol=None))]
    fn register_provider(
        &self,
        id: &str,
        name: &str,
        brand: &str,
        base_url: &str,
        api_key: Option<String>,
        protocol: Option<String>,
    ) -> PyResult<()> {
        let brand_enum = match brand {
            "openai" => ProviderBrand::OpenAI,
            "anthropic" => ProviderBrand::Anthropic,
            "gemini" => ProviderBrand::Gemini,
            "deepseek" => ProviderBrand::DeepSeek,
            "ollama" => ProviderBrand::Ollama,
            other => ProviderBrand::Custom(other.to_string()),
        };

        let protocol_enum = match protocol.as_deref().unwrap_or("openai_compat") {
            "openai_compat" => ProviderProtocol::OpenAiCompat,
            "anthropic" => ProviderProtocol::Anthropic,
            "gemini" => ProviderProtocol::Gemini,
            _ => ProviderProtocol::Custom,
        };

        let mut provider = LlmProvider::new(id, name, brand_enum, base_url)
            .with_protocol(protocol_enum);

        if let Some(key) = api_key {
            provider = provider.with_api_key(key);
        }

        self.store
            .upsert_provider(provider)
            .map_err(|e| PyRuntimeError::new_err(e.to_string()))?;

        Ok(())
    }

    /// List all registered providers.
    fn list_providers(&self, py: Python<'_>) -> PyResult<Vec<PyObject>> {
        let providers = self.store.list_providers(None, None);
        let mut result = Vec::new();
        for p in &providers {
            let dict = PyDict::new_bound(py);
            dict.set_item("id", &p.id)?;
            dict.set_item("name", &p.name)?;
            dict.set_item("brand", format!("{:?}", p.brand))?;
            dict.set_item("base_url", &p.base_url)?;
            dict.set_item("status", format!("{:?}", p.status))?;
            result.push(dict.into());
        }
        Ok(result)
    }

    /// Delete a provider by ID.
    fn delete_provider(&self, id: &str) -> PyResult<bool> {
        self.store
            .delete_provider(id)
            .map_err(|e| PyRuntimeError::new_err(e.to_string()))
    }

    // ========================================================================
    // Model CRUD
    // ========================================================================

    /// Register a new LLM model.
    ///
    /// Args:
    ///     id: Unique model ID (e.g., "gpt-4")
    ///     provider_id: Provider this model belongs to
    ///     name: Model name
    ///     display_name: Human-readable name
    #[pyo3(signature = (id, provider_id, name, display_name=None))]
    fn register_model(
        &self,
        id: &str,
        provider_id: &str,
        name: &str,
        display_name: Option<&str>,
    ) -> PyResult<()> {
        let model = LlmModel::new(id, provider_id, name, display_name.unwrap_or(name));
        self.store
            .upsert_model(model)
            .map_err(|e| PyRuntimeError::new_err(e.to_string()))?;
        Ok(())
    }

    /// List all registered models.
    fn list_models(&self, py: Python<'_>) -> PyResult<Vec<PyObject>> {
        let models = self.store.list_models(None, None);
        let mut result = Vec::new();
        for m in &models {
            let dict = PyDict::new_bound(py);
            dict.set_item("id", &m.id)?;
            dict.set_item("name", &m.name)?;
            dict.set_item("provider_id", &m.provider_id)?;
            dict.set_item("is_default", m.is_default)?;
            result.push(dict.into());
        }
        Ok(result)
    }

    /// Delete a model by ID.
    fn delete_model(&self, id: &str) -> PyResult<bool> {
        self.store
            .delete_model(id)
            .map_err(|e| PyRuntimeError::new_err(e.to_string()))
    }

    // ========================================================================
    // Conversation Management
    // ========================================================================

    /// Create a new conversation.
    ///
    /// Returns a dict with the conversation info.
    #[pyo3(signature = (user_id, provider_id, model_id, model_name, title=None, system_prompt=None))]
    fn create_conversation(
        &self,
        py: Python<'_>,
        user_id: &str,
        provider_id: &str,
        model_id: &str,
        model_name: &str,
        title: Option<String>,
        system_prompt: Option<String>,
    ) -> PyResult<PyObject> {
        let mut conv = self
            .conversations
            .create_conversation(user_id, provider_id, model_id, model_name)
            .map_err(|e| PyRuntimeError::new_err(e.to_string()))?;

        if let Some(t) = title {
            conv.title = t;
        }
        if let Some(prompt) = system_prompt {
            conv.system_prompt = Some(prompt);
        }

        conv_to_dict(py, &conv)
    }

    /// List conversations for a user.
    fn list_conversations(&self, py: Python<'_>, user_id: &str) -> PyResult<Vec<PyObject>> {
        let query = ConversationQuery::for_user(user_id);
        let convs = self.conversations.list_conversations(&query);
        let mut result = Vec::new();
        for c in &convs {
            result.push(conv_to_dict(py, c)?);
        }
        Ok(result)
    }

    /// Get a conversation by ID.
    fn get_conversation(&self, py: Python<'_>, conv_id: &str, user_id: &str) -> PyResult<PyObject> {
        let conv = self
            .conversations
            .get_conversation(conv_id, user_id)
            .map_err(|e| PyRuntimeError::new_err(e.to_string()))?;
        conv_to_dict(py, &conv)
    }

    /// Delete a conversation.
    fn delete_conversation(&self, conv_id: &str, user_id: &str) -> PyResult<bool> {
        self.conversations
            .delete_conversation(conv_id, user_id)
            .map_err(|e| PyRuntimeError::new_err(e.to_string()))
    }

    /// Get messages for a conversation.
    #[pyo3(signature = (conv_id, limit=100, offset=0))]
    fn get_messages(
        &self,
        py: Python<'_>,
        conv_id: &str,
        limit: usize,
        offset: usize,
    ) -> PyResult<Vec<PyObject>> {
        let query = MessageQuery {
            conversation_id: conv_id.to_string(),
            limit,
            offset,
            after_timestamp: None,
            before_timestamp: None,
        };
        let messages = self
            .conversations
            .get_messages(&query)
            .map_err(|e| PyRuntimeError::new_err(e.to_string()))?;

        let mut result = Vec::new();
        for m in &messages {
            result.push(msg_to_dict(py, m)?);
        }
        Ok(result)
    }

    // ========================================================================
    // Chat
    // ========================================================================

    /// Send a message and get a response (non-streaming, blocks until complete).
    ///
    /// Returns a dict with the assistant's response message.
    fn chat(&self, py: Python<'_>, conv_id: &str, user_id: &str, content: &str) -> PyResult<PyObject> {
        let stored = self.rt.block_on(async {
            self.conversations
                .send_message(conv_id, user_id, content)
                .await
        })
        .map_err(|e| PyRuntimeError::new_err(e.to_string()))?;

        msg_to_dict(py, &stored)
    }

    /// Stop an active generation.
    fn stop_generation(&self, conv_id: &str) -> PyResult<bool> {
        Ok(self.conversations.stop_generation(conv_id))
    }

    // ========================================================================
    // Usage Statistics
    // ========================================================================

    /// Get usage statistics for a user.
    fn get_usage_stats(&self, py: Python<'_>, user_id: &str) -> PyResult<PyObject> {
        let stats = self.conversations.get_usage_stats(user_id);
        let dict = PyDict::new_bound(py);
        dict.set_item("total_prompt_tokens", stats.total_prompt_tokens)?;
        dict.set_item("total_completion_tokens", stats.total_completion_tokens)?;
        dict.set_item("total_tokens", stats.total_tokens)?;
        dict.set_item("total_cost", stats.total_cost)?;
        dict.set_item("message_count", stats.message_count)?;
        dict.set_item("conversation_count", stats.conversation_count)?;
        Ok(dict.into())
    }

    // ========================================================================
    // Prompt Templates
    // ========================================================================

    /// Create or update a prompt template.
    ///
    /// Returns a dict with the template info.
    #[pyo3(signature = (id, name, content, description=None, category=None, is_default=false))]
    fn upsert_template(
        &self,
        py: Python<'_>,
        id: &str,
        name: &str,
        content: &str,
        description: Option<&str>,
        category: Option<&str>,
        is_default: bool,
    ) -> PyResult<PyObject> {
        let mut template = PromptTemplate::new(id, name, content);
        if let Some(desc) = description {
            template.description = desc.to_string();
        }
        if let Some(cat) = category {
            template.category = cat.to_string();
        }
        if is_default {
            template.is_default = true;
        }

        let created = self
            .prompts
            .upsert(template)
            .map_err(|e| PyRuntimeError::new_err(e.to_string()))?;

        template_to_dict(py, &created)
    }

    /// Get a prompt template by ID.
    fn get_template(&self, py: Python<'_>, id: &str) -> PyResult<Option<PyObject>> {
        match self.prompts.get(id) {
            Some(t) => Ok(Some(template_to_dict(py, &t)?)),
            None => Ok(None),
        }
    }

    /// List prompt templates, optionally filtered by category.
    #[pyo3(signature = (category=None))]
    fn list_templates(&self, py: Python<'_>, category: Option<&str>) -> PyResult<Vec<PyObject>> {
        let templates = self.prompts.list(category);
        let mut result = Vec::new();
        for t in &templates {
            result.push(template_to_dict(py, t)?);
        }
        Ok(result)
    }

    /// Delete a prompt template by ID.
    fn delete_template(&self, id: &str) -> PyResult<bool> {
        self.prompts
            .delete(id)
            .map_err(|e| PyRuntimeError::new_err(e.to_string()))
    }

    /// Render a prompt template with variables.
    ///
    /// Args:
    ///     id: Template ID
    ///     variables: Dict of variable name → value
    ///
    /// Returns the rendered string.
    fn render_template(&self, id: &str, variables: HashMap<String, String>) -> PyResult<String> {
        self.prompts
            .render(id, &variables)
            .map_err(|e| PyRuntimeError::new_err(e.to_string()))
    }

    // ========================================================================
    // Info
    // ========================================================================

    /// Get gateway summary stats.
    fn info(&self, py: Python<'_>) -> PyResult<PyObject> {
        let dict = PyDict::new_bound(py);
        dict.set_item("provider_count", self.store.provider_count())?;
        dict.set_item("model_count", self.store.model_count())?;
        dict.set_item("conversation_count", self.conversations.conversation_count())?;
        dict.set_item("message_count", self.conversations.message_count())?;
        dict.set_item("template_count", self.prompts.len())?;
        Ok(dict.into())
    }
}

// ============================================================================
// Helpers
// ============================================================================

fn conv_to_dict(py: Python<'_>, conv: &LlmConversation) -> PyResult<PyObject> {
    let dict = PyDict::new_bound(py);
    dict.set_item("id", &conv.id)?;
    dict.set_item("title", &conv.title)?;
    dict.set_item("provider_id", &conv.provider_id)?;
    dict.set_item("model_id", &conv.model_id)?;
    dict.set_item("model_name", &conv.model_name)?;
    dict.set_item("system_prompt", &conv.system_prompt)?;
    dict.set_item("owner_id", &conv.owner_id)?;
    dict.set_item("message_count", conv.message_count)?;
    dict.set_item("total_tokens", conv.total_tokens)?;
    dict.set_item("total_cost", conv.total_cost)?;
    dict.set_item("is_favorite", conv.is_favorite)?;
    dict.set_item("is_archived", conv.is_archived)?;
    dict.set_item("created_at", conv.created_at)?;
    dict.set_item("updated_at", conv.updated_at)?;
    Ok(dict.into())
}

fn msg_to_dict(py: Python<'_>, msg: &StoredMessage) -> PyResult<PyObject> {
    let dict = PyDict::new_bound(py);
    dict.set_item("id", &msg.id)?;
    dict.set_item("conversation_id", &msg.conversation_id)?;
    dict.set_item("role", format!("{:?}", msg.role).to_lowercase())?;
    dict.set_item("content", &msg.content)?;
    dict.set_item("model_id", &msg.model_id)?;
    dict.set_item("status", format!("{:?}", msg.status).to_lowercase())?;
    dict.set_item("created_at", msg.created_at)?;

    if let Some(ref usage) = msg.token_usage {
        let usage_dict = PyDict::new_bound(py);
        usage_dict.set_item("prompt_tokens", usage.prompt_tokens)?;
        usage_dict.set_item("completion_tokens", usage.completion_tokens)?;
        usage_dict.set_item("total_tokens", usage.total_tokens)?;
        dict.set_item("token_usage", usage_dict)?;
    }
    if let Some(duration) = msg.duration_ms {
        dict.set_item("duration_ms", duration)?;
    }

    Ok(dict.into())
}

fn template_to_dict(py: Python<'_>, t: &PromptTemplate) -> PyResult<PyObject> {
    let dict = PyDict::new_bound(py);
    dict.set_item("id", &t.id)?;
    dict.set_item("name", &t.name)?;
    dict.set_item("content", &t.content)?;
    dict.set_item("description", &t.description)?;
    dict.set_item("variables", &t.variables)?;
    dict.set_item("category", &t.category)?;
    dict.set_item("version", &t.version)?;
    dict.set_item("is_builtin", t.is_builtin)?;
    dict.set_item("is_default", t.is_default)?;
    dict.set_item("created_at", t.created_at)?;
    dict.set_item("updated_at", t.updated_at)?;
    Ok(dict.into())
}
