//! Python bindings for LLM types

use pyo3::prelude::*;
use pyo3::exceptions::PyValueError;

use flowgentra_ai::core::llm::{
    LLMConfig, LLMProvider, Message, MessageRole, ResponseFormat, TokenUsage, ToolCall,
    ToolDefinition, model_pricing,
};

use crate::{json_to_py, py_to_json};

// ─── PyLLMConfig ────────────────────────────────────────────────────────────

/// LLM provider configuration.
///
/// Example:
///     config = LLMConfig("openai", "gpt-4", api_key="sk-...")
///     config = LLMConfig("anthropic", "claude-3-opus-20240229", api_key="sk-ant-...")
#[pyclass(name = "LLMConfig")]
#[derive(Clone)]
pub struct PyLLMConfig {
    pub(crate) inner: LLMConfig,
}

#[pymethods]
impl PyLLMConfig {
    /// Create a new LLM config.
    ///
    /// Args:
    ///     provider: Provider name ("openai", "anthropic", "mistral", "groq", "ollama", "huggingface", "azure")
    ///     model: Model identifier (e.g. "gpt-4", "claude-3-opus-20240229")
    ///     api_key: API key for the provider
    ///     temperature: Response randomness (0.0-2.0, optional)
    ///     max_tokens: Maximum response tokens (optional)
    ///     top_p: Nucleus sampling parameter (0.0-1.0, optional)
    #[new]
    #[pyo3(signature = (provider, model, api_key="".to_string(), temperature=None, max_tokens=None, top_p=None))]
    fn new(
        provider: &str,
        model: String,
        api_key: String,
        temperature: Option<f32>,
        max_tokens: Option<usize>,
        top_p: Option<f32>,
    ) -> PyResult<Self> {
        let llm_provider = match provider.to_lowercase().as_str() {
            "openai" => LLMProvider::OpenAI,
            "anthropic" => LLMProvider::Anthropic,
            "mistral" => LLMProvider::Mistral,
            "groq" => LLMProvider::Groq,
            "huggingface" => LLMProvider::HuggingFace,
            "ollama" => LLMProvider::Ollama,
            "azure" => LLMProvider::Azure,
            other => LLMProvider::Custom(other.to_string()),
        };

        let mut config = LLMConfig::new(llm_provider, model, api_key);
        config.temperature = temperature;
        config.max_tokens = max_tokens;
        config.top_p = top_p;

        Ok(PyLLMConfig { inner: config })
    }

    #[getter]
    fn provider(&self) -> String {
        self.inner.provider.name().to_string()
    }

    #[getter]
    fn model(&self) -> String {
        self.inner.model.clone()
    }

    #[getter]
    fn temperature(&self) -> Option<f32> {
        self.inner.temperature
    }

    #[setter]
    fn set_temperature(&mut self, val: Option<f32>) {
        self.inner.temperature = val;
    }

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

    #[setter]
    fn set_max_tokens(&mut self, val: Option<usize>) {
        self.inner.max_tokens = val;
    }

    #[getter]
    fn top_p(&self) -> Option<f32> {
        self.inner.top_p
    }

    #[setter]
    fn set_top_p(&mut self, val: Option<f32>) {
        self.inner.top_p = val;
    }

    #[getter]
    fn api_key(&self) -> String {
        self.inner.api_key.clone()
    }

    /// Set the response format for structured output.
    fn set_response_format(&mut self, format: &PyResponseFormat) {
        self.inner.response_format = format.inner.clone();
    }

    fn __repr__(&self) -> String {
        format!(
            "LLMConfig(provider='{}', model='{}')",
            self.inner.provider.name(),
            self.inner.model
        )
    }
}

// ─── PyResponseFormat ──────────────────────────────────────────────────────

/// Response format for LLM output.
///
/// Example:
///     fmt = ResponseFormat.text()
///     fmt = ResponseFormat.json()
///     fmt = ResponseFormat.json_schema("person", {"type": "object", ...})
#[pyclass(name = "ResponseFormat")]
#[derive(Clone)]
pub struct PyResponseFormat {
    pub(crate) inner: ResponseFormat,
}

#[pymethods]
impl PyResponseFormat {
    /// Plain text output (default).
    #[staticmethod]
    fn text() -> Self {
        PyResponseFormat { inner: ResponseFormat::Text }
    }

    /// Force JSON output.
    #[staticmethod]
    fn json() -> Self {
        PyResponseFormat { inner: ResponseFormat::Json }
    }

    /// JSON output constrained to a specific schema.
    #[staticmethod]
    fn json_schema(name: &str, schema: &Bound<'_, pyo3::types::PyAny>) -> PyResult<Self> {
        let json_schema = py_to_json(schema)?;
        Ok(PyResponseFormat {
            inner: ResponseFormat::JsonSchema {
                name: name.to_string(),
                schema: json_schema,
            },
        })
    }

    fn __repr__(&self) -> String {
        match &self.inner {
            ResponseFormat::Text => "ResponseFormat.Text".to_string(),
            ResponseFormat::Json => "ResponseFormat.Json".to_string(),
            ResponseFormat::JsonSchema { name, .. } => {
                format!("ResponseFormat.JsonSchema(name='{}')", name)
            }
        }
    }

    fn __str__(&self) -> String {
        self.__repr__()
    }
}

// ─── PyMessage ──────────────────────────────────────────────────────────────

/// A message in a conversation.
///
/// Example:
///     msg = Message.system("You are a helpful assistant")
///     msg = Message.user("What is Rust?")
///     msg = Message.assistant("Rust is a systems programming language")
#[pyclass(name = "Message")]
#[derive(Clone)]
pub struct PyMessage {
    pub(crate) inner: Message,
}

#[pymethods]
impl PyMessage {
    #[new]
    #[pyo3(signature = (role, content, tool_call_id=None))]
    fn new(role: &str, content: String, tool_call_id: Option<String>) -> PyResult<Self> {
        let msg_role = match role.to_lowercase().as_str() {
            "system" => MessageRole::System,
            "user" => MessageRole::User,
            "assistant" => MessageRole::Assistant,
            "tool" => MessageRole::Tool,
            _ => return Err(PyValueError::new_err(format!("Invalid role: {}", role))),
        };
        Ok(PyMessage {
            inner: Message {
                role: msg_role,
                content,
                tool_calls: None,
                tool_call_id,
            },
        })
    }

    /// Create a system message
    #[staticmethod]
    fn system(content: &str) -> Self {
        PyMessage {
            inner: Message::system(content),
        }
    }

    /// Create a user message
    #[staticmethod]
    fn user(content: &str) -> Self {
        PyMessage {
            inner: Message::user(content),
        }
    }

    /// Create an assistant message
    #[staticmethod]
    fn assistant(content: &str) -> Self {
        PyMessage {
            inner: Message::assistant(content),
        }
    }

    /// Create a tool result message
    #[staticmethod]
    #[pyo3(signature = (content, tool_call_id=None))]
    fn tool(content: &str, tool_call_id: Option<&str>) -> Self {
        let msg = if let Some(id) = tool_call_id {
            Message::tool_result(id, content)
        } else {
            Message::tool(content)
        };
        PyMessage { inner: msg }
    }

    #[getter]
    fn role(&self) -> String {
        match self.inner.role {
            MessageRole::System => "system".to_string(),
            MessageRole::User => "user".to_string(),
            MessageRole::Assistant => "assistant".to_string(),
            MessageRole::Tool => "tool".to_string(),
        }
    }

    #[getter]
    fn content(&self) -> String {
        self.inner.content.clone()
    }

    #[setter]
    fn set_content(&mut self, content: String) {
        self.inner.content = content;
    }

    #[getter]
    fn tool_call_id(&self) -> Option<String> {
        self.inner.tool_call_id.clone()
    }

    /// Whether this message has tool calls
    fn has_tool_calls(&self) -> bool {
        self.inner.has_tool_calls()
    }

    fn is_system(&self) -> bool {
        self.inner.is_system()
    }

    fn is_user(&self) -> bool {
        self.inner.is_user()
    }

    fn is_assistant(&self) -> bool {
        self.inner.is_assistant()
    }

    fn is_tool(&self) -> bool {
        self.inner.is_tool()
    }

    /// Get tool calls as list of PyToolCall
    fn tool_calls(&self) -> Vec<PyToolCall> {
        self.inner
            .tool_calls
            .as_ref()
            .map(|tcs| {
                tcs.iter()
                    .map(|tc| PyToolCall {
                        inner: tc.clone(),
                    })
                    .collect()
            })
            .unwrap_or_default()
    }

    fn __repr__(&self) -> String {
        let preview_len = self.inner.content.len().min(80);
        format!("Message(role='{}', content='{}')", self.role(), &self.inner.content[..preview_len])
    }

    fn __str__(&self) -> String {
        format!("{}", self.inner)
    }
}

// ─── PyToolCall ─────────────────────────────────────────────────────────────

/// A tool call made by the LLM
#[pyclass(name = "ToolCall")]
#[derive(Clone)]
pub struct PyToolCall {
    pub(crate) inner: ToolCall,
}

#[pymethods]
impl PyToolCall {
    #[new]
    fn new(id: String, name: String, arguments: &Bound<'_, PyAny>) -> PyResult<Self> {
        let args = py_to_json(arguments)?;
        Ok(PyToolCall {
            inner: ToolCall {
                id,
                name,
                arguments: args,
            },
        })
    }

    #[getter]
    fn id(&self) -> String {
        self.inner.id.clone()
    }

    #[getter]
    fn name(&self) -> String {
        self.inner.name.clone()
    }

    #[getter]
    fn arguments(&self, py: Python<'_>) -> PyResult<PyObject> {
        json_to_py(py, &self.inner.arguments)
    }

    fn __repr__(&self) -> String {
        format!("ToolCall(id='{}', name='{}')", self.inner.id, self.inner.name)
    }
}

// ─── PyToolDefinition ───────────────────────────────────────────────────────

/// A tool definition passed to the LLM
#[pyclass(name = "ToolDefinition")]
#[derive(Clone)]
pub struct PyToolDefinition {
    pub(crate) inner: ToolDefinition,
}

#[pymethods]
impl PyToolDefinition {
    #[new]
    fn new(name: String, description: String, parameters: &Bound<'_, PyAny>) -> PyResult<Self> {
        let params = py_to_json(parameters)?;
        Ok(PyToolDefinition {
            inner: ToolDefinition {
                name,
                description,
                parameters: params,
            },
        })
    }

    #[getter]
    fn name(&self) -> String {
        self.inner.name.clone()
    }

    #[getter]
    fn description(&self) -> String {
        self.inner.description.clone()
    }

    #[getter]
    fn parameters(&self, py: Python<'_>) -> PyResult<PyObject> {
        json_to_py(py, &self.inner.parameters)
    }

    fn __repr__(&self) -> String {
        format!(
            "ToolDefinition(name='{}', description='{}')",
            self.inner.name, self.inner.description
        )
    }
}

// ─── PyTokenUsage ───────────────────────────────────────────────────────────

/// Token usage statistics from an LLM response
#[pyclass(name = "TokenUsage")]
#[derive(Clone)]
pub struct PyTokenUsage {
    pub(crate) inner: TokenUsage,
}

#[pymethods]
impl PyTokenUsage {
    #[new]
    fn new(prompt_tokens: u64, completion_tokens: u64) -> Self {
        PyTokenUsage {
            inner: TokenUsage::new(prompt_tokens, completion_tokens),
        }
    }

    #[getter]
    fn prompt_tokens(&self) -> u64 {
        self.inner.prompt_tokens
    }

    #[getter]
    fn completion_tokens(&self) -> u64 {
        self.inner.completion_tokens
    }

    #[getter]
    fn total_tokens(&self) -> u64 {
        self.inner.total_tokens
    }

    /// Estimate cost in USD for the given model
    fn estimated_cost(&self, model: &str) -> Option<f64> {
        self.inner.estimated_cost(model)
    }

    fn __repr__(&self) -> String {
        format!(
            "TokenUsage(prompt={}, completion={}, total={})",
            self.inner.prompt_tokens, self.inner.completion_tokens, self.inner.total_tokens
        )
    }
}

// ─── Free function ──────────────────────────────────────────────────────────

/// Get model pricing (input_price_per_million, output_price_per_million) in USD.
///
/// Returns None if the model is not in the pricing table.
#[pyfunction]
pub fn py_model_pricing(model: &str) -> Option<(f64, f64)> {
    model_pricing(model)
}
