//! LLM Provider abstraction.
//!
//! Alle store LLM-APIer er varianter av "send messages + tools, få svar tilbake".
//! OpenAI-kompatibelt format (OpenRouter, xAI, OpenAI, Groq, Ollama) bruker
//! identisk wire format. Anthropic og Gemini har egne formater.
//!
//! Vi støtter alle via en felles trait.

use serde::{Deserialize, Serialize};

// ============================================================================
// Felles typer
// ============================================================================

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Message {
    pub role: String, // "system", "user", "assistant", "tool"
    pub content: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tool_calls: Option<Vec<ToolCall>>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tool_call_id: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolCall {
    pub id: String,
    pub r#type: String, // "function"
    pub function: FunctionCall,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FunctionCall {
    pub name: String,
    pub arguments: String, // JSON string
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolDef {
    pub r#type: String, // "function"
    pub function: FunctionDef,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FunctionDef {
    pub name: String,
    pub description: String,
    pub parameters: serde_json::Value, // JSON Schema
}

#[derive(Debug, Clone)]
pub struct CompletionResponse {
    pub message: Message,
    pub usage: TokenUsage,
    pub model: String,
}

#[derive(Debug, Clone, Default)]
pub struct TokenUsage {
    pub input_tokens: u64,
    pub output_tokens: u64,
}

// ============================================================================
// Provider trait
// ============================================================================

#[async_trait::async_trait]
pub trait LlmProvider: Send + Sync {
    /// Send messages + tools, get completion back.
    async fn complete(
        &self,
        messages: &[Message],
        tools: &[ToolDef],
    ) -> Result<CompletionResponse, ProviderError>;

    /// Provider name (for logging)
    fn name(&self) -> &str;

    /// Model identifier
    fn model_id(&self) -> &str;

    /// Context window size in tokens.
    fn context_window(&self) -> u64 {
        crate::context::estimate_context_window(self.model_id())
    }
}

#[derive(Debug, thiserror::Error)]
pub enum ProviderError {
    #[error("HTTP error: {0}")]
    Http(#[from] reqwest::Error),
    #[error("API error ({status}): {body}")]
    Api { status: u16, body: String },
    #[error("Parse error: {0}")]
    Parse(String),
    #[error("No API key configured for {provider}")]
    NoApiKey { provider: String },
    #[allow(dead_code)]
    #[error("Budget exhausted: {0}")]
    BudgetExhausted(String),
}

impl ProviderError {
    /// Whether this error is retryable (transient API failures).
    pub fn is_retryable(&self) -> bool {
        match self {
            ProviderError::Api { status, .. } => matches!(status, 429 | 500 | 502 | 503),
            ProviderError::Http(e) => {
                // Retry on connection/timeout errors
                e.is_timeout() || e.is_connect()
            }
            _ => false,
        }
    }

    /// Extract Retry-After hint (seconds) from error body if present.
    pub fn retry_after_hint(&self) -> Option<u64> {
        if let ProviderError::Api { status: 429, body } = self {
            // Some APIs include retry_after in JSON error body
            if let Ok(v) = serde_json::from_str::<serde_json::Value>(body) {
                if let Some(secs) = v["error"]["retry_after"].as_u64()
                    .or_else(|| v["retry_after"].as_u64())
                {
                    return Some(secs);
                }
            }
        }
        None
    }
}

// ============================================================================
// OpenAI-compatible provider (OpenRouter, xAI, OpenAI, Groq, Ollama)
// ============================================================================

pub struct OpenAiCompatible {
    client: reqwest::Client,
    base_url: String,
    api_key: String,
    model: String,
    provider_name: String,
    max_tokens: u32,
}

impl OpenAiCompatible {
    pub fn new(
        base_url: impl Into<String>,
        api_key: impl Into<String>,
        model: impl Into<String>,
        provider_name: impl Into<String>,
    ) -> Self {
        Self {
            client: reqwest::Client::new(),
            base_url: base_url.into(),
            api_key: api_key.into(),
            model: model.into(),
            provider_name: provider_name.into(),
            max_tokens: 4096,
        }
    }

    pub fn with_max_tokens(mut self, max_tokens: u32) -> Self {
        self.max_tokens = max_tokens;
        self
    }

    /// OpenRouter
    pub fn openrouter(api_key: impl Into<String>, model: impl Into<String>) -> Self {
        Self::new(
            "https://openrouter.ai/api/v1",
            api_key,
            model,
            "openrouter",
        )
    }

    /// OpenAI
    pub fn openai(api_key: impl Into<String>, model: impl Into<String>) -> Self {
        Self::new(
            "https://api.openai.com/v1",
            api_key,
            model,
            "openai",
        )
    }

    /// xAI (Grok)
    pub fn xai(api_key: impl Into<String>, model: impl Into<String>) -> Self {
        Self::new(
            "https://api.x.ai/v1",
            api_key,
            model,
            "xai",
        )
    }

    /// Ollama (lokal)
    pub fn ollama(model: impl Into<String>) -> Self {
        Self::new(
            "http://localhost:11434/v1",
            "ollama", // Ollama trenger ikke key
            model,
            "ollama",
        )
    }
}

#[async_trait::async_trait]
impl LlmProvider for OpenAiCompatible {
    async fn complete(
        &self,
        messages: &[Message],
        tools: &[ToolDef],
    ) -> Result<CompletionResponse, ProviderError> {
        let mut body = serde_json::json!({
            "model": self.model,
            "messages": messages,
            "max_tokens": self.max_tokens,
        });

        if !tools.is_empty() {
            body["tools"] = serde_json::to_value(tools).unwrap();
        }

        let res = self
            .client
            .post(format!("{}/chat/completions", self.base_url))
            .header("Authorization", format!("Bearer {}", self.api_key))
            .header("Content-Type", "application/json")
            .json(&body)
            .send()
            .await?;

        let status = res.status().as_u16();
        if status >= 400 {
            let body = res.text().await.unwrap_or_default();
            return Err(ProviderError::Api { status, body });
        }

        let data: serde_json::Value = res.json().await?;

        let choice = data["choices"]
            .get(0)
            .ok_or_else(|| ProviderError::Parse("No choices in response".into()))?;

        let msg = &choice["message"];
        let tool_calls: Option<Vec<ToolCall>> = msg
            .get("tool_calls")
            .and_then(|tc| serde_json::from_value(tc.clone()).ok());

        let message = Message {
            role: "assistant".into(),
            content: msg["content"].as_str().map(|s| s.to_string()),
            tool_calls,
            tool_call_id: None,
        };

        let usage = TokenUsage {
            input_tokens: data["usage"]["prompt_tokens"].as_u64().unwrap_or(0),
            output_tokens: data["usage"]["completion_tokens"].as_u64().unwrap_or(0),
        };

        let model = data["model"]
            .as_str()
            .unwrap_or(&self.model)
            .to_string();

        Ok(CompletionResponse {
            message,
            usage,
            model,
        })
    }

    fn name(&self) -> &str {
        &self.provider_name
    }

    fn model_id(&self) -> &str {
        &self.model
    }
}

// ============================================================================
// Anthropic provider
// ============================================================================

pub struct Anthropic {
    client: reqwest::Client,
    api_key: String,
    model: String,
    max_tokens: u32,
}

impl Anthropic {
    pub fn new(api_key: impl Into<String>, model: impl Into<String>) -> Self {
        Self {
            client: reqwest::Client::new(),
            api_key: api_key.into(),
            model: model.into(),
            max_tokens: 8192,
        }
    }

    pub fn with_max_tokens(mut self, max_tokens: u32) -> Self {
        self.max_tokens = max_tokens;
        self
    }
}

#[async_trait::async_trait]
impl LlmProvider for Anthropic {
    async fn complete(
        &self,
        messages: &[Message],
        tools: &[ToolDef],
    ) -> Result<CompletionResponse, ProviderError> {
        // Convert OpenAI-style messages to Anthropic format
        let system = messages
            .iter()
            .filter(|m| m.role == "system")
            .filter_map(|m| m.content.as_deref())
            .collect::<Vec<_>>()
            .join("\n\n");

        let anthropic_messages: Vec<serde_json::Value> = messages
            .iter()
            .filter(|m| m.role != "system")
            .map(|m| {
                if m.role == "tool" {
                    serde_json::json!({
                        "role": "user",
                        "content": [{
                            "type": "tool_result",
                            "tool_use_id": m.tool_call_id,
                            "content": m.content,
                        }]
                    })
                } else if let Some(ref tool_calls) = m.tool_calls {
                    let content: Vec<serde_json::Value> = tool_calls
                        .iter()
                        .map(|tc| {
                            serde_json::json!({
                                "type": "tool_use",
                                "id": tc.id,
                                "name": tc.function.name,
                                "input": serde_json::from_str::<serde_json::Value>(&tc.function.arguments)
                                    .unwrap_or(serde_json::Value::Object(Default::default())),
                            })
                        })
                        .collect();
                    serde_json::json!({ "role": "assistant", "content": content })
                } else {
                    serde_json::json!({
                        "role": m.role,
                        "content": m.content.as_deref().unwrap_or(""),
                    })
                }
            })
            .collect();

        let anthropic_tools: Vec<serde_json::Value> = tools
            .iter()
            .map(|t| {
                serde_json::json!({
                    "name": t.function.name,
                    "description": t.function.description,
                    "input_schema": t.function.parameters,
                })
            })
            .collect();

        let mut body = serde_json::json!({
            "model": self.model,
            "max_tokens": self.max_tokens,
            "messages": anthropic_messages,
        });

        if !system.is_empty() {
            body["system"] = serde_json::Value::String(system);
        }
        if !anthropic_tools.is_empty() {
            body["tools"] = serde_json::to_value(&anthropic_tools).unwrap();
        }

        let res = self
            .client
            .post("https://api.anthropic.com/v1/messages")
            .header("x-api-key", &self.api_key)
            .header("anthropic-version", "2023-06-01")
            .header("Content-Type", "application/json")
            .json(&body)
            .send()
            .await?;

        let status = res.status().as_u16();
        if status >= 400 {
            let body = res.text().await.unwrap_or_default();
            return Err(ProviderError::Api { status, body });
        }

        let data: serde_json::Value = res.json().await?;

        // Parse Anthropic response → our common format
        let content_blocks = data["content"].as_array();
        let mut text_parts = Vec::new();
        let mut tool_calls = Vec::new();

        if let Some(blocks) = content_blocks {
            for block in blocks {
                match block["type"].as_str() {
                    Some("text") => {
                        if let Some(t) = block["text"].as_str() {
                            text_parts.push(t.to_string());
                        }
                    }
                    Some("tool_use") => {
                        tool_calls.push(ToolCall {
                            id: block["id"].as_str().unwrap_or("").to_string(),
                            r#type: "function".into(),
                            function: FunctionCall {
                                name: block["name"].as_str().unwrap_or("").to_string(),
                                arguments: serde_json::to_string(&block["input"])
                                    .unwrap_or_default(),
                            },
                        });
                    }
                    _ => {}
                }
            }
        }

        let message = Message {
            role: "assistant".into(),
            content: if text_parts.is_empty() {
                None
            } else {
                Some(text_parts.join("\n"))
            },
            tool_calls: if tool_calls.is_empty() {
                None
            } else {
                Some(tool_calls)
            },
            tool_call_id: None,
        };

        let usage = TokenUsage {
            input_tokens: data["usage"]["input_tokens"].as_u64().unwrap_or(0),
            output_tokens: data["usage"]["output_tokens"].as_u64().unwrap_or(0),
        };

        Ok(CompletionResponse {
            message,
            usage,
            model: self.model.clone(),
        })
    }

    fn name(&self) -> &str {
        "anthropic"
    }

    fn model_id(&self) -> &str {
        &self.model
    }
}

// ============================================================================
// Google Gemini provider
// ============================================================================

pub struct Gemini {
    client: reqwest::Client,
    api_key: String,
    model: String,
    max_tokens: u32,
}

impl Gemini {
    pub fn new(api_key: impl Into<String>, model: impl Into<String>) -> Self {
        Self {
            client: reqwest::Client::new(),
            api_key: api_key.into(),
            model: model.into(),
            max_tokens: 8192,
        }
    }

    pub fn with_max_tokens(mut self, max_tokens: u32) -> Self {
        self.max_tokens = max_tokens;
        self
    }
}

#[async_trait::async_trait]
impl LlmProvider for Gemini {
    async fn complete(
        &self,
        messages: &[Message],
        tools: &[ToolDef],
    ) -> Result<CompletionResponse, ProviderError> {
        // Convert to Gemini format
        let system = messages
            .iter()
            .filter(|m| m.role == "system")
            .filter_map(|m| m.content.as_deref())
            .collect::<Vec<_>>()
            .join("\n\n");

        let contents: Vec<serde_json::Value> = messages
            .iter()
            .filter(|m| m.role != "system")
            .map(|m| {
                let role = match m.role.as_str() {
                    "assistant" => "model",
                    "tool" => "function",
                    _ => "user",
                };

                if m.role == "tool" {
                    serde_json::json!({
                        "role": "function",
                        "parts": [{
                            "functionResponse": {
                                "name": m.tool_call_id.as_deref().unwrap_or("unknown"),
                                "response": { "result": m.content.as_deref().unwrap_or("") }
                            }
                        }]
                    })
                } else if let Some(ref tool_calls) = m.tool_calls {
                    let parts: Vec<serde_json::Value> = tool_calls
                        .iter()
                        .map(|tc| {
                            serde_json::json!({
                                "functionCall": {
                                    "name": tc.function.name,
                                    "args": serde_json::from_str::<serde_json::Value>(&tc.function.arguments)
                                        .unwrap_or_default(),
                                }
                            })
                        })
                        .collect();
                    serde_json::json!({ "role": role, "parts": parts })
                } else {
                    serde_json::json!({
                        "role": role,
                        "parts": [{ "text": m.content.as_deref().unwrap_or("") }]
                    })
                }
            })
            .collect();

        let gemini_tools: Vec<serde_json::Value> = if tools.is_empty() {
            vec![]
        } else {
            let declarations: Vec<serde_json::Value> = tools
                .iter()
                .map(|t| {
                    serde_json::json!({
                        "name": t.function.name,
                        "description": t.function.description,
                        "parameters": t.function.parameters,
                    })
                })
                .collect();
            vec![serde_json::json!({ "functionDeclarations": declarations })]
        };

        let mut body = serde_json::json!({
            "contents": contents,
            "generationConfig": {
                "maxOutputTokens": self.max_tokens,
            },
        });

        if !system.is_empty() {
            body["systemInstruction"] = serde_json::json!({
                "parts": [{ "text": system }]
            });
        }
        if !gemini_tools.is_empty() {
            body["tools"] = serde_json::to_value(&gemini_tools).unwrap();
        }

        let url = format!(
            "https://generativelanguage.googleapis.com/v1beta/models/{}:generateContent?key={}",
            self.model, self.api_key
        );

        let res = self
            .client
            .post(&url)
            .header("Content-Type", "application/json")
            .json(&body)
            .send()
            .await?;

        let status = res.status().as_u16();
        if status >= 400 {
            let body = res.text().await.unwrap_or_default();
            return Err(ProviderError::Api { status, body });
        }

        let data: serde_json::Value = res.json().await?;

        // Parse Gemini response
        let parts = data["candidates"][0]["content"]["parts"]
            .as_array();

        let mut text_parts = Vec::new();
        let mut tool_calls = Vec::new();

        if let Some(parts) = parts {
            for (i, part) in parts.iter().enumerate() {
                if let Some(text) = part["text"].as_str() {
                    text_parts.push(text.to_string());
                }
                if let Some(fc) = part.get("functionCall") {
                    tool_calls.push(ToolCall {
                        id: format!("call_{}", i),
                        r#type: "function".into(),
                        function: FunctionCall {
                            name: fc["name"].as_str().unwrap_or("").to_string(),
                            arguments: serde_json::to_string(&fc["args"])
                                .unwrap_or_default(),
                        },
                    });
                }
            }
        }

        let message = Message {
            role: "assistant".into(),
            content: if text_parts.is_empty() {
                None
            } else {
                Some(text_parts.join("\n"))
            },
            tool_calls: if tool_calls.is_empty() {
                None
            } else {
                Some(tool_calls)
            },
            tool_call_id: None,
        };

        let usage = TokenUsage {
            input_tokens: data["usageMetadata"]["promptTokenCount"]
                .as_u64()
                .unwrap_or(0),
            output_tokens: data["usageMetadata"]["candidatesTokenCount"]
                .as_u64()
                .unwrap_or(0),
        };

        Ok(CompletionResponse {
            message,
            usage,
            model: self.model.clone(),
        })
    }

    fn name(&self) -> &str {
        "gemini"
    }

    fn model_id(&self) -> &str {
        &self.model
    }
}

// ============================================================================
// Provider factory
// ============================================================================

/// Create a provider from a model string.
///
/// Format: "provider/model" or just "model" (defaults to openrouter).
///
/// Examples:
///   "openrouter/anthropic/claude-sonnet-4"
///   "anthropic/claude-sonnet-4"
///   "gemini/gemini-2.5-flash"
///   "xai/grok-3"
///   "openai/gpt-4o"
///   "ollama/llama3"
pub fn create_provider(
    model_spec: &str,
    api_keys: &ApiKeys,
    max_tokens: Option<u32>,
) -> Result<Box<dyn LlmProvider>, ProviderError> {
    let (provider, model) = if let Some(idx) = model_spec.find('/') {
        (&model_spec[..idx], &model_spec[idx + 1..])
    } else {
        ("openrouter", model_spec)
    };

    match provider {
        "openrouter" => {
            let key = api_keys.openrouter.as_deref()
                .ok_or_else(|| ProviderError::NoApiKey { provider: "openrouter".into() })?;
            let mut p = OpenAiCompatible::openrouter(key, model);
            if let Some(mt) = max_tokens { p = p.with_max_tokens(mt); }
            Ok(Box::new(p))
        }
        "anthropic" => {
            let key = api_keys.anthropic.as_deref()
                .ok_or_else(|| ProviderError::NoApiKey { provider: "anthropic".into() })?;
            let mut p = Anthropic::new(key, model);
            if let Some(mt) = max_tokens { p = p.with_max_tokens(mt); }
            Ok(Box::new(p))
        }
        "gemini" | "google" => {
            let key = api_keys.gemini.as_deref()
                .ok_or_else(|| ProviderError::NoApiKey { provider: "gemini".into() })?;
            let mut p = Gemini::new(key, model);
            if let Some(mt) = max_tokens { p = p.with_max_tokens(mt); }
            Ok(Box::new(p))
        }
        "xai" | "grok" => {
            let key = api_keys.xai.as_deref()
                .ok_or_else(|| ProviderError::NoApiKey { provider: "xai".into() })?;
            let mut p = OpenAiCompatible::xai(key, model);
            if let Some(mt) = max_tokens { p = p.with_max_tokens(mt); }
            Ok(Box::new(p))
        }
        "openai" => {
            let key = api_keys.openai.as_deref()
                .ok_or_else(|| ProviderError::NoApiKey { provider: "openai".into() })?;
            let mut p = OpenAiCompatible::openai(key, model);
            if let Some(mt) = max_tokens { p = p.with_max_tokens(mt); }
            Ok(Box::new(p))
        }
        "ollama" | "local" => {
            let mut p = OpenAiCompatible::ollama(model);
            if let Some(mt) = max_tokens { p = p.with_max_tokens(mt); }
            Ok(Box::new(p))
        }
        _ => Err(ProviderError::Parse(format!("Unknown provider: {}", provider))),
    }
}

// ============================================================================
// Cost estimation
// ============================================================================

/// Cost per million tokens (input, output) for known models.
/// Returns (0.0, 0.0) for unknown models.
pub fn cost_per_million_tokens(model: &str) -> (f64, f64) {
    let m = model.to_lowercase();

    // Anthropic
    if m.contains("opus") { return (15.0, 75.0); }
    if m.contains("sonnet") { return (3.0, 15.0); }
    if m.contains("haiku") { return (0.25, 1.25); }

    // Gemini
    if m.contains("gemini") && m.contains("flash") { return (0.075, 0.30); }
    if m.contains("gemini") && m.contains("pro") { return (1.25, 5.0); }

    // Grok
    if m.contains("grok-3") && m.contains("mini") { return (0.30, 0.50); }
    if m.contains("grok-3") { return (3.0, 15.0); }
    if m.contains("grok") { return (2.0, 10.0); }

    // OpenAI
    if m.contains("gpt-4o-mini") { return (0.15, 0.60); }
    if m.contains("gpt-4o") { return (2.50, 10.0); }
    if m.contains("o1") || m.contains("o3") { return (10.0, 40.0); }

    // Local / unknown — free
    (0.0, 0.0)
}

/// Calculate cost from token usage and model.
pub fn calculate_cost(model: &str, usage: &TokenUsage) -> f64 {
    let (input_cpm, output_cpm) = cost_per_million_tokens(model);
    (usage.input_tokens as f64 * input_cpm + usage.output_tokens as f64 * output_cpm) / 1_000_000.0
}

// ============================================================================
// Retry with exponential backoff
// ============================================================================

/// Retry configuration for API calls.
pub struct RetryConfig {
    /// Maximum number of retries (default: 3).
    pub max_retries: u32,
    /// Base delay in milliseconds (default: 1000).
    pub base_delay_ms: u64,
    /// Maximum delay in milliseconds (default: 30000).
    pub max_delay_ms: u64,
}

impl Default for RetryConfig {
    fn default() -> Self {
        Self {
            max_retries: 3,
            base_delay_ms: 1_000,
            max_delay_ms: 30_000,
        }
    }
}

/// Call provider.complete() with retry + exponential backoff on retryable errors.
pub async fn complete_with_retry(
    provider: &dyn LlmProvider,
    messages: &[Message],
    tools: &[ToolDef],
    config: &RetryConfig,
) -> Result<CompletionResponse, ProviderError> {
    let mut last_err = None;

    for attempt in 0..=config.max_retries {
        match provider.complete(messages, tools).await {
            Ok(resp) => return Ok(resp),
            Err(e) => {
                if !e.is_retryable() || attempt == config.max_retries {
                    return Err(e);
                }

                // Calculate delay: exponential backoff with jitter
                let base = config.base_delay_ms * 2u64.pow(attempt);
                let delay = if let Some(retry_after) = e.retry_after_hint() {
                    // Respect Retry-After header from 429
                    retry_after * 1000
                } else {
                    base.min(config.max_delay_ms)
                };

                tracing::warn!(
                    attempt = attempt + 1,
                    max_retries = config.max_retries,
                    delay_ms = delay,
                    error = %e,
                    "Retryable API error, backing off"
                );

                tokio::time::sleep(std::time::Duration::from_millis(delay)).await;
                last_err = Some(e);
            }
        }
    }

    Err(last_err.unwrap())
}

/// API keys loaded from environment
#[derive(Debug, Clone, Default)]
pub struct ApiKeys {
    pub openrouter: Option<String>,
    pub anthropic: Option<String>,
    pub gemini: Option<String>,
    pub xai: Option<String>,
    pub openai: Option<String>,
}

impl ApiKeys {
    pub fn from_env() -> Self {
        Self {
            openrouter: std::env::var("OPENROUTER_API_KEY").ok().filter(|s| s != "PLACEHOLDER" && !s.is_empty()),
            anthropic: std::env::var("ANTHROPIC_API_KEY").ok().filter(|s| !s.is_empty()),
            gemini: std::env::var("GEMINI_API_KEY").ok().filter(|s| !s.is_empty()),
            xai: std::env::var("XAI_API_KEY").ok().filter(|s| !s.is_empty()),
            openai: std::env::var("OPENAI_API_KEY").ok().filter(|s| !s.is_empty()),
        }
    }
}
