#[cfg(feature = "remote")]
use crate::agent::remote::NodeId;
use crate::model::{AgentMessage, MessagePart};
use crate::model_info::get_model_info;
use crate::session::error::{SessionError, SessionResult};
use crate::session::provider_config::{ProviderConfigMode, resolve_provider_config};
use crate::session::store::{LLMConfig, Session, SessionExecutionConfig, SessionStore};
use arc_swap::ArcSwap;
use querymt::LLMParams;
use querymt::plugin::host::PluginRegistry;
use querymt::providers::ModelPricing;
use querymt::{
    LLMProvider,
    chat::{ChatMessage, ChatResponse, ChatRole, Content},
    error::LLMError,
};
use std::sync::Arc;
#[cfg(feature = "remote")]
use std::sync::atomic::{AtomicBool, Ordering};

/// Cached provider entry: `(config_id, provider_node_id, allow_mesh_fallback, provider)`.
///
/// `provider_node_id` is stored separately from `llm_configs` (in the `sessions`
/// table) so the same `config_id` can resolve to different providers depending
/// on which mesh node owns the session. `allow_mesh_fallback` is included so
/// toggling policy does not reuse a provider built under different routing rules.
type ProviderCacheEntry = (i64, Option<String>, bool, Arc<dyn LLMProvider>);

/// Lock-free single-entry provider cache.  Reads (every turn) use `ArcSwap::load`,
/// writes (model switch, mesh fallback toggle) use `ArcSwap::store`.
type ProviderCache = Arc<ArcSwap<Option<ProviderCacheEntry>>>;

fn is_non_billable_oauth_provider(provider: &str) -> bool {
    matches!(provider, "codex" | "kimi-code")
}

/// A wrapper around a `SessionStore` that resolves providers dynamically.
///
/// # Provider Caching
///
/// Provider construction is cached at two levels:
///
/// 1. **Global cache** (`cached_provider`): A single-entry cache keyed on
///    `LLMConfig.id`. Avoids rebuilding identical providers across turns.
///    For multi-model scenarios consider an LRU at the factory level.
///
/// 2. **Per-turn cache** (`SessionHandle::cached_llm_provider`): Each
///    `SessionHandle` lazily caches its resolved provider for the lifetime of
///    the turn. This means the global cache is consulted at most once per turn
///    instead of on every state-machine transition (~8× reduction).
///
/// Model switches via the dashboard write a new `llm_config_id` to the DB but
/// take effect only when the **next** `SessionHandle` is constructed (i.e. the
/// next `run_prompt` call), not mid-turn.
pub struct SessionProvider {
    plugin_registry: Arc<PluginRegistry>,
    history_store: Arc<dyn SessionStore>,
    initial_config: LLMParams,
    /// Lock-free single-entry provider cache keyed on `LLMConfig.id`.
    /// Uses a single-entry cache to ensure safe VRAM management for GPU models.
    cached_provider: ProviderCache,
    /// Agent identifier used for `{{ agent_id }}` in system prompt templates.
    /// Set at build time via `with_agent_id` so that session creation can
    /// resolve per-agent template variables.
    pub(crate) agent_id: Option<String>,
    /// Optional mesh handle — present when this node participates in a kameo mesh.
    /// Passed through to `build_provider` to enable `MeshChatProvider`
    /// routing and mesh-fallback discovery.
    ///
    /// Uses `ArcSwap` so the mesh can be injected *after* the
    /// `Arc<SessionProvider>` is already shared (e.g. via `AgentHandle::set_mesh`
    /// which is called after `AgentConfigBuilder::build()`). Lock-free reads.
    #[cfg(feature = "remote")]
    mesh: Arc<arc_swap::ArcSwapOption<crate::agent::remote::MeshHandle>>,
    /// Whether to scan the mesh when a provider is missing locally and
    /// `provider_node_id` is not explicitly set. Defaults to false.
    #[cfg(feature = "remote")]
    allow_mesh_fallback: Arc<AtomicBool>,
}

impl SessionProvider {
    pub fn new(
        plugin_registry: Arc<PluginRegistry>,
        store: Arc<dyn SessionStore>,
        initial_config: LLMParams,
    ) -> Self {
        Self {
            plugin_registry,
            history_store: store,
            initial_config,
            cached_provider: Arc::new(ArcSwap::from_pointee(None)),
            agent_id: None,
            #[cfg(feature = "remote")]
            mesh: Arc::new(arc_swap::ArcSwapOption::empty()),
            #[cfg(feature = "remote")]
            allow_mesh_fallback: Arc::new(AtomicBool::new(false)),
        }
    }

    /// Set the agent identifier used for `{{ agent_id }}` in system prompt
    /// templates.  Call at build time before wrapping in `Arc`.
    pub fn with_agent_id(mut self, id: Option<String>) -> Self {
        self.agent_id = id;
        self
    }

    /// Returns the initial `LLMParams` this provider was constructed with.
    ///
    /// Used by `ProviderHostActor` to forward the agent's custom parameters
    /// (e.g. `model_path`, `n_ctx`) when building a local provider on behalf
    /// of a remote mesh peer.  Without these params the friendly model name
    /// (e.g. `"qwen3-coder"`) cannot be resolved to a real GGUF path.
    pub fn initial_params(&self) -> &LLMParams {
        &self.initial_config
    }

    /// Attach a mesh handle at construction time (consuming builder).
    ///
    /// Use this when the mesh is available before the `SessionProvider` is
    /// wrapped in an `Arc`. For late injection (mesh bootstrapped after
    /// `Arc<SessionProvider>` is shared), use `set_mesh` instead.
    #[cfg(feature = "remote")]
    pub fn with_mesh(self, mesh: Option<crate::agent::remote::MeshHandle>) -> Self {
        match mesh {
            Some(m) => self.mesh.store(Some(Arc::new(m))),
            None => self.mesh.store(None),
        }
        self
    }

    /// Inject or replace the mesh handle after construction.
    ///
    /// Because `mesh` uses `ArcSwapOption`, this works even when the
    /// `SessionProvider` is already shared behind an `Arc`. All clones of
    /// this `SessionProvider` share the same `Arc` and therefore see the update
    /// immediately — no restart or rebuild is required.
    ///
    /// Called by `AgentHandle::set_mesh` so that sessions created by a
    /// `RemoteNodeManager` (which holds `Arc<AgentConfig>` with this provider)
    /// can route LLM calls through the mesh even though the mesh was bootstrapped
    /// after the config was built.
    #[cfg(feature = "remote")]
    pub fn set_mesh(&self, mesh: Option<crate::agent::remote::MeshHandle>) {
        match mesh {
            Some(m) => self.mesh.store(Some(Arc::new(m))),
            None => self.mesh.store(None),
        }
    }

    /// Enable/disable automatic mesh fallback when `provider_node_id` is not set.
    ///
    /// When disabled (default), `provider_node_id = None` means local-only lookup.
    /// When enabled, unresolved providers may be discovered on mesh peers.
    #[cfg(feature = "remote")]
    pub fn set_mesh_fallback(&self, enabled: bool) {
        self.allow_mesh_fallback.store(enabled, Ordering::Relaxed);
    }

    /// Attach mesh fallback policy at construction time (consuming builder).
    #[cfg(feature = "remote")]
    pub fn with_mesh_fallback(self, enabled: bool) -> Self {
        self.allow_mesh_fallback.store(enabled, Ordering::Relaxed);
        self
    }

    /// Fetch an existing session by ID
    pub async fn get_session(&self, session_id: &str) -> SessionResult<Option<Session>> {
        self.history_store.get_session(session_id).await
    }

    /// Load an existing session by ID
    pub async fn with_session(&self, session_id: &str) -> SessionResult<SessionHandle> {
        let session = self
            .get_session(session_id)
            .await?
            .ok_or_else(|| SessionError::SessionNotFound(session_id.to_string()))?;
        SessionHandle::new(Arc::new(self.clone()), session).await
    }

    /// Create a new session with optional cwd and parent
    pub async fn create_session(
        &self,
        cwd: Option<std::path::PathBuf>,
        parent_session_id: Option<&str>,
        execution_config: &SessionExecutionConfig,
    ) -> SessionResult<SessionHandle> {
        let fork_origin = if parent_session_id.is_some() {
            Some(crate::session::domain::ForkOrigin::Delegation)
        } else {
            None
        };
        let mut session = self
            .history_store
            .create_session(
                None,
                cwd.clone(),
                parent_session_id.map(|s| s.to_string()),
                fork_origin,
            )
            .await?;

        // Resolve session-specific template variables (cwd, datetime, …) in
        // system prompt strings before writing the config to the database.
        let resolved_config = {
            let mut builder = crate::template::SessionTemplateContext::builder();
            if let Some(ref p) = cwd {
                builder = builder.cwd(p);
            }
            if let Some(ref p) = self.initial_config.provider {
                builder = builder.provider(p.as_str());
            }
            if let Some(ref m) = self.initial_config.model {
                builder = builder.model(m.as_str());
            }
            if let Some(ref id) = self.agent_id {
                builder = builder.agent_id(id.as_str());
            }
            #[cfg(feature = "remote")]
            {
                if self.mesh.load().is_some() {
                    builder = builder.has_mesh(true);
                }
            }
            let ctx = builder.build();
            ctx.resolve_params(&self.initial_config).map_err(|e| {
                SessionError::InvalidOperation(format!(
                    "Failed to resolve system prompt templates: {e}"
                ))
            })?
        };

        let llm_config = self
            .history_store
            .create_or_get_llm_config(&resolved_config)
            .await?;
        self.history_store
            .set_session_llm_config(&session.public_id, llm_config.id)
            .await?;
        self.history_store
            .set_session_execution_config(&session.public_id, execution_config)
            .await?;
        session.llm_config_id = Some(llm_config.id);
        SessionHandle::new(Arc::new(self.clone()), session).await
    }

    pub fn history_store(&self) -> Arc<dyn SessionStore> {
        self.history_store.clone()
    }

    pub fn plugin_registry(&self) -> Arc<PluginRegistry> {
        self.plugin_registry.clone()
    }

    pub fn initial_config(&self) -> &LLMParams {
        &self.initial_config
    }

    pub async fn build_provider_for_session(
        &self,
        session_id: &str,
    ) -> SessionResult<Arc<dyn LLMProvider>> {
        let config = self
            .history_store
            .get_session_llm_config(session_id)
            .await?
            .ok_or_else(|| {
                SessionError::InvalidOperation("Session has no LLM config".to_string())
            })?;

        // `parse_llm_config_row` always returns `provider_node_id: None` because
        // provider_node_id lives in the `sessions` table, not `llm_configs`.
        // Read it separately so mesh routing (Case 1) is actually triggered
        // for remote sessions.
        #[cfg(feature = "remote")]
        let provider_node_id: Option<String> = self
            .history_store
            .get_session_provider_node_id(session_id)
            .await
            .unwrap_or(None);
        #[cfg(not(feature = "remote"))]
        let provider_node_id: Option<String> = None;

        #[cfg(feature = "remote")]
        let allow_mesh_fallback = self.allow_mesh_fallback.load(Ordering::Relaxed);
        #[cfg(not(feature = "remote"))]
        let allow_mesh_fallback = false;

        // Check cache first - lock-free fast path for same config + provider_node_id + fallback policy
        {
            let cache = self.cached_provider.load();
            if let Some(ref entry) = **cache {
                let (
                    ref cached_config_id,
                    ref cached_provider_node_id,
                    cached_allow_mesh_fallback,
                    ref cached_provider,
                ) = *entry;
                if *cached_config_id == config.id
                    && cached_provider_node_id.as_deref() == provider_node_id.as_deref()
                    && cached_allow_mesh_fallback == allow_mesh_fallback
                {
                    log::trace!(
                        "Provider cache hit for config_id={} ({}:{}) provider_node_id={:?} allow_mesh_fallback={}",
                        config.id,
                        config.provider,
                        config.model,
                        provider_node_id,
                        allow_mesh_fallback,
                    );
                    return Ok(Arc::clone(cached_provider));
                }
            }
        }

        // Cache miss or config/provider_node_id changed - build new provider
        log::debug!(
            "Provider cache miss for config_id={} ({}:{}) provider_node_id={:?} allow_mesh_fallback={}, building new provider",
            config.id,
            config.provider,
            config.model,
            provider_node_id,
            allow_mesh_fallback,
        );

        let req = ProviderRequest::new(&config.provider, &config.model)
            .with_params(config.params.as_ref())
            .with_session_id(session_id);
        #[cfg(feature = "remote")]
        let req = req.with_provider_node_id(provider_node_id.as_deref());

        let provider = self.build_provider(req).await?;

        // Update cache - this will drop the old provider if config changed.
        // For GPU models (llama_cpp), this frees VRAM before loading the new model.
        self.cached_provider.store(Arc::new(Some((
            config.id,
            provider_node_id,
            allow_mesh_fallback,
            Arc::clone(&provider),
        ))));

        Ok(provider)
    }

    /// Clear the global provider cache.
    ///
    /// This forces the next provider resolution to rebuild from current config/credentials.
    pub async fn clear_provider_cache(&self) {
        self.cached_provider.store(Arc::new(None));
    }

    /// Get pricing information for a session's model
    ///
    /// Returns `None` if:
    /// - The session doesn't have an LLM config
    /// - Pricing information is not available for the model
    pub async fn get_session_pricing(
        &self,
        session_id: &str,
    ) -> SessionResult<Option<ModelPricing>> {
        let llm_config = self
            .history_store
            .get_session_llm_config(session_id)
            .await?;

        Ok(llm_config.and_then(|config| {
            if is_non_billable_oauth_provider(&config.provider) {
                return None;
            }
            get_model_info(&config.provider, &config.model).map(|info| info.pricing)
        }))
    }

    /// Get pricing information for a specific provider and model
    pub fn get_pricing(provider: &str, model: &str) -> Option<ModelPricing> {
        if is_non_billable_oauth_provider(provider) {
            return None;
        }
        get_model_info(provider, model).map(|info| info.pricing)
    }

    /// Get LLM config by ID
    pub async fn get_llm_config(&self, config_id: i64) -> SessionResult<Option<LLMConfig>> {
        self.history_store.get_llm_config(config_id).await
    }

    /// Build an LLM provider from a [`ProviderRequest`].
    ///
    /// Infrastructure (plugin registry, mesh handle, mesh-fallback policy) is
    /// read from `&self`. The [`ProviderRequest`] carries per-call values:
    /// provider name, model, params, optional API-key override, session ID
    /// (for prompt-cache key substitution), and optional mesh node routing.
    ///
    /// Provider construction includes:
    /// - Routing to a remote `MeshChatProvider` when `provider_node_id` names a mesh peer
    /// - Looking up the factory from the plugin registry
    /// - Merging model and params into a builder config
    /// - Resolving API keys (respecting user's preferred auth method from dashboard)
    /// - Applying model-specific heuristic defaults
    /// - Falling back to the mesh when the provider is unavailable locally (requires `remote` feature)
    pub async fn build_provider(
        &self,
        req: ProviderRequest<'_>,
    ) -> SessionResult<Arc<dyn LLMProvider>> {
        let provider_name = req.provider_name;
        let model = req.model;
        let params = req.params;
        let api_key_override = req.api_key_override;
        let session_id = req.session_id;

        // ── Case 1: Explicit remote node requested ─────────────────────────────
        #[cfg(feature = "remote")]
        if let Some(node_id) = req.provider_node_id
            && node_id != "local"
        {
            let mesh_guard = self.mesh.load();
            let mesh = mesh_guard.as_deref().ok_or_else(|| {
                SessionError::InvalidOperation(format!(
                    "provider_node_id='{}' specified but no mesh handle available",
                    node_id
                ))
            })?;
            let node_id = NodeId::parse(node_id).map_err(SessionError::InvalidOperation)?;
            log::debug!(
                "build_provider: routing {}/{} to mesh node '{}'",
                provider_name,
                model,
                node_id
            );
            let remote_params = params
                .cloned()
                .map(|mut p| {
                    if let Some(obj) = p.as_object_mut() {
                        obj.insert(
                            "_remote_session_id".to_string(),
                            serde_json::Value::String(session_id.to_string()),
                        );
                    }
                    p
                })
                .or_else(|| Some(serde_json::json!({"_remote_session_id": session_id})));
            return Ok(Arc::new(
                querymt_remote::MeshChatProvider::from_node_id(
                    mesh,
                    &node_id,
                    provider_name,
                    model,
                )
                .with_params(remote_params),
            ));
        }

        // ── Case 2: Try local provider ─────────────────────────────────────────
        let plugin_registry = &self.plugin_registry;
        let factory = plugin_registry.get(provider_name).await;

        #[cfg(not(feature = "remote"))]
        let factory = factory.ok_or_else(|| {
            SessionError::InvalidOperation(format!("Unknown provider: {}", provider_name))
        })?;

        // With the remote feature enabled we may optionally fall back to the mesh
        // (Case 3 below), so we don't error-out immediately.
        #[cfg(feature = "remote")]
        let factory = match factory {
            Some(f) => f,
            None => {
                // ── Case 3: Not available locally → optional mesh fallback ─────
                let allow_mesh_fallback = self.allow_mesh_fallback.load(Ordering::Relaxed);
                let mesh_guard = self.mesh.load();
                let mesh_handle = mesh_guard.as_deref();

                if allow_mesh_fallback && let Some(mesh) = mesh_handle {
                    log::debug!(
                        "build_provider: provider '{}' not found locally, searching mesh",
                        provider_name
                    );
                    if let Some(node_id) =
                        querymt_remote::find_provider_on_mesh(mesh, provider_name).await
                    {
                        log::info!(
                            "build_provider: found '{}' on mesh peer '{}', using MeshChatProvider",
                            provider_name,
                            node_id
                        );
                        let remote_params = params
                            .cloned()
                            .map(|mut p| {
                                if let Some(obj) = p.as_object_mut() {
                                    obj.insert(
                                        "_remote_session_id".to_string(),
                                        serde_json::Value::String(session_id.to_string()),
                                    );
                                }
                                p
                            })
                            .or_else(|| {
                                Some(serde_json::json!({"_remote_session_id": session_id}))
                            });
                        return Ok(Arc::new(
                            querymt_remote::MeshChatProvider::from_node_id(
                                mesh,
                                &node_id,
                                provider_name,
                                model,
                            )
                            .with_params(remote_params),
                        ));
                    }
                }
                return Err(SessionError::InvalidOperation(format!(
                    "Unknown provider: {}",
                    provider_name
                )));
            }
        };

        let resolved_cfg = resolve_provider_config(
            plugin_registry,
            &self.initial_config,
            &factory,
            provider_name,
            ProviderConfigMode::Runtime {
                model,
                params,
                api_key_override,
                session_id,
            },
        )
        .await?;

        #[cfg(feature = "oauth")]
        let builder_config = resolved_cfg.builder_config;
        let pruned_config_str = resolved_cfg.pruned_config_str;
        let pruned_keys = resolved_cfg.pruned_keys;
        let _use_oauth_resolver = resolved_cfg.use_oauth_resolver;

        if !pruned_keys.is_empty() {
            const MAX_KEYS_TO_LOG: usize = 50;
            let shown = pruned_keys.len().min(MAX_KEYS_TO_LOG);
            let suffix = if pruned_keys.len() > shown {
                format!(" (+{} more)", pruned_keys.len() - shown)
            } else {
                String::new()
            };

            log::warn!(
                "Pruned unsupported config keys for provider '{}': {}{}",
                provider_name,
                pruned_keys[..shown].join(", "),
                suffix
            );
        }

        // If OAuth was used, attach a resolver so expired tokens are refreshed
        // transparently on each request.
        //
        // We try the generic LLMProviderFactory path first: Extism providers
        // handle HTTP internally and accept set_key_resolver directly on the
        // LLMProvider. For native HTTP providers the generic path wraps the
        // inner HTTP provider in Arc before we can set the resolver, so we
        // fall back to the HTTPLLMProviderFactory path to set it on the inner
        // provider before wrapping.
        #[cfg(feature = "oauth")]
        if _use_oauth_resolver {
            let initial_key = builder_config
                .get("api_key")
                .and_then(|v| v.as_str())
                .unwrap_or("")
                .to_string();

            let resolver = std::sync::Arc::new(crate::auth::OAuthKeyResolver::new(
                provider_name,
                &initial_key,
            ));

            // Generic path — works for Extism providers that implement set_key_resolver.
            let mut provider = factory.from_config(&pruned_config_str).map_err(|e| {
                LLMError::PluginError(format!("provider '{}': {}", provider_name, e))
            })?;
            provider.set_key_resolver(resolver.clone());

            if provider.key_resolver().is_some() {
                log::debug!(
                    "Attached OAuthKeyResolver via LLMProvider for '{}' (model: {})",
                    provider_name,
                    model
                );
                return Ok(Arc::from(provider));
            }

            // Fallback for native HTTP providers: set the resolver on the inner
            // HTTPLLMProvider before it gets wrapped in Arc by the adapter.
            if let Some(http_factory) = factory.as_http() {
                let mut http_provider =
                    http_factory.from_config(&pruned_config_str).map_err(|e| {
                        LLMError::PluginError(format!("provider '{}': {}", provider_name, e))
                    })?;
                http_provider.set_key_resolver(resolver);

                log::debug!(
                    "Attached OAuthKeyResolver via HTTPLLMProvider for '{}' (model: {})",
                    provider_name,
                    model
                );

                let adapter = querymt::adapters::LLMProviderFromHTTP::new(http_provider);
                return Ok(Arc::from(Box::new(adapter) as Box<dyn LLMProvider>));
            }

            // Neither path attached a resolver — return the provider as-is.
            log::warn!(
                "OAuthKeyResolver could not be attached for provider '{}' (model: {})",
                provider_name,
                model
            );
            return Ok(Arc::from(provider));
        }

        let provider = factory
            .from_config(&pruned_config_str)
            .map_err(|e| LLMError::PluginError(format!("provider '{}': {}", provider_name, e)))?;
        Ok(Arc::from(provider))
    }
}

impl Clone for SessionProvider {
    fn clone(&self) -> Self {
        Self {
            plugin_registry: self.plugin_registry.clone(),
            history_store: Arc::clone(&self.history_store),
            initial_config: self.initial_config.clone(),
            cached_provider: Arc::clone(&self.cached_provider),
            agent_id: self.agent_id.clone(),
            // Share the same Arcs/atomics so all clones see runtime updates.
            #[cfg(feature = "remote")]
            mesh: Arc::clone(&self.mesh),
            #[cfg(feature = "remote")]
            allow_mesh_fallback: Arc::clone(&self.allow_mesh_fallback),
        }
    }
}

/// A handle to an active session, providing access to the session's LLM provider
/// and message history.
///
/// This handle encapsulates the session state and provides methods for interacting
/// with the LLM and managing conversation history.
///
/// # Turn-pinned semantics
///
/// A `SessionHandle` captures the session's `LLMConfig` at construction time and
/// caches the resolved `LLMProvider` for reuse within a turn. Model switches
/// (e.g. via the dashboard) take effect on the **next** `run_prompt`, not mid-turn.
/// This avoids redundant DB lookups and ensures consistent provider/model identity
/// throughout a single execution cycle.
pub struct SessionHandle {
    provider: Arc<SessionProvider>,
    session: Session,
    /// LLM config resolved once at construction time (turn-pinned).
    llm_config: Option<LLMConfig>,
    /// Session execution config resolved once at construction time (turn-pinned).
    execution_config: Option<SessionExecutionConfig>,
    /// Lazily cached LLM provider for this turn.
    cached_llm_provider: tokio::sync::OnceCell<Arc<dyn LLMProvider>>,
}

impl Clone for SessionHandle {
    fn clone(&self) -> Self {
        Self {
            provider: self.provider.clone(),
            session: self.session.clone(),
            llm_config: self.llm_config.clone(),
            execution_config: self.execution_config.clone(),
            // Each clone gets its own OnceCell; the first `.provider()` call
            // will still hit the global cache (cheap Arc::clone on hit) so
            // this is fine — it just won't share the local cell.
            cached_llm_provider: tokio::sync::OnceCell::new(),
        }
    }
}

impl SessionHandle {
    pub async fn new(provider: Arc<SessionProvider>, session: Session) -> SessionResult<Self> {
        // Eagerly resolve the LLM config so callers never need a separate DB fetch.
        let llm_config = if let Some(config_id) = session.llm_config_id {
            provider.get_llm_config(config_id).await?
        } else {
            None
        };
        let execution_config = provider
            .history_store
            .get_session_execution_config(&session.public_id)
            .await?;
        Ok(Self {
            provider,
            session,
            llm_config,
            execution_config,
            cached_llm_provider: tokio::sync::OnceCell::new(),
        })
    }

    /// Get the session information
    pub fn session(&self) -> &Session {
        &self.session
    }

    /// Get the LLM config captured at handle creation time (turn-pinned).
    pub fn llm_config(&self) -> Option<&LLMConfig> {
        self.llm_config.as_ref()
    }

    /// Get the session execution config captured at handle creation time (turn-pinned).
    pub fn execution_config(&self) -> Option<&SessionExecutionConfig> {
        self.execution_config.as_ref()
    }

    pub async fn provider(&self) -> SessionResult<Arc<dyn LLMProvider>> {
        self.cached_llm_provider
            .get_or_try_init(|| async {
                self.provider
                    .build_provider_for_session(&self.session.public_id)
                    .await
            })
            .await
            .map(Arc::clone)
    }

    /// Get the session history as rich AgentMessages
    pub async fn get_agent_history(&self) -> SessionResult<Vec<AgentMessage>> {
        self.provider
            .history_store
            .get_history(&self.session.public_id)
            .await
    }

    /// Get the effective session history (post-compaction, no snapshot-only messages)
    pub async fn get_effective_agent_history(&self) -> SessionResult<Vec<AgentMessage>> {
        self.provider
            .history_store
            .get_effective_history(&self.session.public_id)
            .await
    }

    /// Get the session history converted to standard ChatMessages for the LLM
    pub async fn history(&self) -> Vec<ChatMessage> {
        match self.get_effective_agent_history().await {
            Ok(agent_msgs) => agent_msgs
                .iter()
                .map(|m| {
                    m.to_chat_message_with_target(
                        self.llm_config.as_ref().map(|cfg| cfg.provider.as_str()),
                        self.llm_config.as_ref().map(|cfg| cfg.model.as_str()),
                        self.execution_config
                            .as_ref()
                            .and_then(|cfg| cfg.max_prompt_bytes),
                    )
                })
                .collect(),
            Err(err) => {
                log::warn!("Failed to load session history: {}", err);
                Vec::new()
            }
        }
    }

    /// Persist an AgentMessage to the store
    pub async fn add_message(&self, message: AgentMessage) -> SessionResult<()> {
        self.provider
            .history_store
            .add_message(&self.session.public_id, message)
            .await
    }

    /// Execute a raw tool call without side effects
    pub async fn call_tool(
        &self,
        name: &str,
        args: serde_json::Value,
    ) -> Result<Vec<Content>, LLMError> {
        let provider = self.provider().await?;
        provider.call_tool(name, args).await
    }

    /// Submit messages to the LLM without auto-saving
    pub async fn submit_request(
        &self,
        messages: &[ChatMessage],
    ) -> Result<Box<dyn ChatResponse>, LLMError> {
        let provider = self.provider().await?;
        provider.chat(messages).await
    }

    /// Higher-level chat interface (used by CLI) that handles conversion and storage
    pub async fn chat(&self, messages: &[ChatMessage]) -> SessionResult<Box<dyn ChatResponse>> {
        // 1. Store incoming messages (User or Tool Result)
        for msg in messages {
            let agent_msg = self.convert_chat_to_agent(msg);
            self.add_message(agent_msg).await?;
        }

        // 2. Fetch full history for context
        let llm_messages = self.history().await;

        // 3. Call LLM
        let response = self.submit_request(&llm_messages).await?;

        // 4. Store response
        let response_msg: ChatMessage = response.as_ref().into();
        let agent_response = self.convert_chat_to_agent(&response_msg);
        self.add_message(agent_response).await?;

        Ok(response)
    }

    /// Get pricing information for this session's model.
    ///
    /// Returns `None` for non-billable OAuth providers or OAuth sessions
    /// (no per-token cost), or when pricing information is unavailable for the model.
    ///
    /// OAuth is detected by checking whether the cached provider has a key
    /// resolver attached — OAuth providers always do, static-key providers
    /// never do. This avoids threading a boolean through the provider stack.
    pub fn get_pricing(&self) -> Option<ModelPricing> {
        if self
            .llm_config
            .as_ref()
            .is_some_and(|config| is_non_billable_oauth_provider(&config.provider))
        {
            return None;
        }

        if let Some(provider) = self.cached_llm_provider.get()
            && provider.key_resolver().is_some()
        {
            return None;
        }

        self.llm_config
            .as_ref()
            .and_then(|config| get_model_info(&config.provider, &config.model))
            .map(|info| info.pricing)
    }

    pub fn convert_chat_to_agent(&self, msg: &ChatMessage) -> AgentMessage {
        let mut parts = Vec::new();

        for block in &msg.content {
            match block {
                Content::Text { text } => {
                    parts.push(MessagePart::Text {
                        content: text.clone(),
                    });
                }
                Content::Thinking { text, signature } => {
                    parts.push(MessagePart::Reasoning {
                        content: text.clone(),
                        signature: signature.clone(),
                        time_ms: None,
                    });
                }
                Content::ToolUse {
                    id,
                    name,
                    arguments,
                } => {
                    parts.push(MessagePart::ToolUse(querymt::ToolCall {
                        id: id.clone(),
                        call_type: "function".to_string(),
                        function: querymt::FunctionCall {
                            name: name.clone(),
                            arguments: arguments.to_string(),
                        },
                    }));
                }
                Content::ToolResult {
                    id,
                    name,
                    is_error,
                    content,
                } => {
                    parts.push(MessagePart::ToolResult {
                        call_id: id.clone(),
                        content: content.clone(),
                        is_error: *is_error,
                        tool_name: name.clone(),
                        tool_arguments: None,
                        compacted_at: None,
                    });
                }
                _ => {
                    // Image, ImageUrl, Pdf, Audio, ResourceLink — store as text description
                    let desc = block.to_string();
                    if !desc.is_empty() {
                        parts.push(MessagePart::Text { content: desc });
                    }
                }
            }
        }

        let source_provider = if msg.role == ChatRole::Assistant {
            self.llm_config.as_ref().map(|cfg| cfg.provider.clone())
        } else {
            None
        };
        let source_model = if msg.role == ChatRole::Assistant {
            self.llm_config.as_ref().map(|cfg| cfg.model.clone())
        } else {
            None
        };

        AgentMessage {
            id: uuid::Uuid::now_v7().to_string(),
            session_id: self.session.public_id.clone(),
            role: msg.role.clone(),
            parts,
            created_at: time::OffsetDateTime::now_utc().unix_timestamp(),
            parent_message_id: None,
            source_provider,
            source_model,
        }
    }
}

/// Preferred authentication method for a provider.
///
/// Stored in the system keyring under `auth_method_{provider}` and used by
/// the credential resolution logic to determine the order in which OAuth,
/// stored API key, and environment variable sources are tried.
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
#[cfg_attr(feature = "dashboard", typeshare::typeshare)]
#[serde(rename_all = "snake_case")]
pub enum AuthMethod {
    #[serde(rename = "oauth")]
    OAuth,
    ApiKey,
    EnvVar,
}

impl std::fmt::Display for AuthMethod {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            AuthMethod::OAuth => write!(f, "oauth"),
            AuthMethod::ApiKey => write!(f, "api_key"),
            AuthMethod::EnvVar => write!(f, "env_var"),
        }
    }
}

impl std::str::FromStr for AuthMethod {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "oauth" => Ok(AuthMethod::OAuth),
            "api_key" => Ok(AuthMethod::ApiKey),
            "env_var" => Ok(AuthMethod::EnvVar),
            other => Err(format!("unknown auth method: {}", other)),
        }
    }
}

/// Read the preferred auth method for a provider from SecretStore.
pub(crate) fn preferred_auth_method(provider_name: &str) -> Option<AuthMethod> {
    querymt_utils::secret_store::SecretStore::new()
        .ok()
        .and_then(|s| s.get(&format!("auth_method_{}", provider_name)))
        .and_then(|v| v.parse::<AuthMethod>().ok())
}

/// Resolve an API key for a provider, respecting the user's preferred auth method.
///
/// The three credential sources are:
/// - **OAuth**: token from the system keyring (requires `oauth` feature)
/// - **Stored key**: API key set via the dashboard UI (stored in SecretStore under `env_var_name`)
/// - **Env var**: environment variable (e.g. `OPENAI_API_KEY`)
///
/// `env_var_name` may be `None` for OAuth-only providers (e.g. Codex). In that case
/// only the OAuth path is attempted; stored-key and env-var are skipped.
///
/// When `preferred` is `None`, the default order is: OAuth → stored key → env var.
/// When set, the preferred source is tried first, then the remaining sources in order.
pub(crate) async fn resolve_api_key_with_preference(
    provider_name: &str,
    env_var_name: Option<&str>,
    preferred: Option<&AuthMethod>,
    use_oauth_resolver: &mut bool,
) -> Option<String> {
    #[cfg(not(feature = "oauth"))]
    let _ = use_oauth_resolver;

    // Build the resolution order based on preference
    let order: Vec<AuthMethod> = match preferred {
        Some(AuthMethod::OAuth) | None => {
            vec![AuthMethod::OAuth, AuthMethod::ApiKey, AuthMethod::EnvVar]
        }
        Some(AuthMethod::ApiKey) => {
            vec![AuthMethod::ApiKey, AuthMethod::EnvVar, AuthMethod::OAuth]
        }
        Some(AuthMethod::EnvVar) => {
            vec![AuthMethod::EnvVar, AuthMethod::ApiKey, AuthMethod::OAuth]
        }
    };

    for method in &order {
        match method {
            AuthMethod::OAuth => {
                #[cfg(feature = "oauth")]
                {
                    use crate::auth::get_or_refresh_token;
                    match get_or_refresh_token(provider_name).await {
                        Ok(token) => {
                            log::debug!("Using OAuth token for provider: {}", provider_name);
                            *use_oauth_resolver = true;
                            return Some(token);
                        }
                        Err(e) => {
                            log::debug!("OAuth unavailable for {}: {}", provider_name, e);
                        }
                    }
                }
            }
            AuthMethod::ApiKey => {
                // Stored key requires an env_var_name (used as keyring key)
                if let Some(name) = env_var_name {
                    if let Some(key) = querymt_utils::secret_store::SecretStore::new()
                        .ok()
                        .and_then(|s| s.get(name))
                    {
                        log::debug!(
                            "Using stored API key for provider: {} ({})",
                            provider_name,
                            name
                        );
                        return Some(key);
                    }
                    log::debug!("No stored API key for {} ({})", provider_name, name);
                }
            }
            AuthMethod::EnvVar => {
                if let Some(name) = env_var_name {
                    if let Ok(key) = std::env::var(name) {
                        log::debug!("Using env var for provider: {} ({})", provider_name, name);
                        return Some(key);
                    }
                    log::debug!("Env var {} not set", name);
                }
            }
        }
    }

    None
}

/// Per-call parameters for building an LLM provider via
/// [`SessionProvider::build_provider`].
///
/// Infrastructure concerns (plugin registry, mesh handle, mesh-fallback policy)
/// live on [`SessionProvider`] and are accessed through `&self`. This struct
/// carries only the values that vary per request.
///
/// Use [`ProviderRequest::new`] for the common case — optional fields default
/// to sensible values (`params: None`, `api_key_override: None`,
/// `session_id: "standalone"`, `provider_node_id: None`).
pub struct ProviderRequest<'a> {
    /// Provider name (e.g. `"anthropic"`, `"openai"`).
    pub provider_name: &'a str,
    /// Model name (e.g. `"claude-sonnet-4-20250514"`, `"gpt-5"`).
    pub model: &'a str,
    /// Optional extra params JSON blob merged into the provider config.
    pub params: Option<&'a serde_json::Value>,
    /// Override the API key resolved from env/OAuth/keyring.
    pub api_key_override: Option<&'a str>,
    /// Session identifier for prompt-cache key substitution.
    /// Defaults to `"standalone"` for session-less contexts.
    pub session_id: &'a str,
    /// Route the call to a specific mesh node's `ProviderHostActor`.
    /// `None` (default) or `"local"` means local provider resolution.
    #[cfg(feature = "remote")]
    pub provider_node_id: Option<&'a str>,
}

impl<'a> ProviderRequest<'a> {
    /// Create a request with only the required fields; optional fields use defaults.
    pub fn new(provider_name: &'a str, model: &'a str) -> Self {
        Self {
            provider_name,
            model,
            params: None,
            api_key_override: None,
            session_id: "standalone",
            #[cfg(feature = "remote")]
            provider_node_id: None,
        }
    }

    pub fn with_params(mut self, params: Option<&'a serde_json::Value>) -> Self {
        self.params = params;
        self
    }

    pub fn with_api_key_override(mut self, api_key_override: Option<&'a str>) -> Self {
        self.api_key_override = api_key_override;
        self
    }

    pub fn with_session_id(mut self, session_id: &'a str) -> Self {
        self.session_id = session_id;
        self
    }

    #[cfg(feature = "remote")]
    pub fn with_provider_node_id(mut self, provider_node_id: Option<&'a str>) -> Self {
        self.provider_node_id = provider_node_id;
        self
    }
}

#[cfg(test)]
pub mod tests {
    use super::*;
    use querymt::chat::{ChatMessage, ChatResponse, FinishReason, StreamChunk};
    use querymt::completion::{CompletionRequest, CompletionResponse};
    use querymt::error::LLMError;
    use std::pin::Pin;
    use tokio::sync::mpsc;
    use tokio_stream::wrappers::ReceiverStream;

    /// Mock LLM provider for testing
    pub struct MockProvider {
        response_text: String,
    }

    impl Default for MockProvider {
        fn default() -> Self {
            Self::new()
        }
    }

    impl MockProvider {
        pub fn new() -> Self {
            Self {
                response_text: "Mock response".to_string(),
            }
        }

        #[allow(dead_code)]
        pub fn with_response(response: String) -> Self {
            Self {
                response_text: response,
            }
        }
    }

    // ChatResponse implementation
    #[derive(Debug)]
    struct MockChatResponse {
        content: String,
    }

    impl std::fmt::Display for MockChatResponse {
        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
            write!(f, "{}", self.content)
        }
    }

    impl ChatResponse for MockChatResponse {
        fn text(&self) -> Option<String> {
            Some(self.content.clone())
        }

        fn thinking(&self) -> Option<String> {
            None
        }

        fn usage(&self) -> Option<querymt::Usage> {
            None
        }

        fn finish_reason(&self) -> Option<FinishReason> {
            Some(FinishReason::Stop)
        }

        fn tool_calls(&self) -> Option<Vec<querymt::ToolCall>> {
            None
        }
    }

    #[async_trait::async_trait]
    impl querymt::chat::ChatProvider for MockProvider {
        async fn chat_with_tools(
            &self,
            _messages: &[ChatMessage],
            _tools: Option<&[querymt::chat::Tool]>,
        ) -> Result<Box<dyn ChatResponse>, LLMError> {
            Ok(Box::new(MockChatResponse {
                content: self.response_text.clone(),
            }))
        }

        async fn chat_stream_with_tools(
            &self,
            _messages: &[ChatMessage],
            _tools: Option<&[querymt::chat::Tool]>,
        ) -> Result<
            Pin<Box<dyn tokio_stream::Stream<Item = Result<StreamChunk, LLMError>> + Send>>,
            LLMError,
        > {
            let (_tx, rx) = mpsc::channel(1);
            Ok(Box::pin(ReceiverStream::new(rx)))
        }
    }

    #[async_trait::async_trait]
    impl querymt::completion::CompletionProvider for MockProvider {
        async fn complete(&self, _req: &CompletionRequest) -> Result<CompletionResponse, LLMError> {
            Err(LLMError::NotImplemented("Not implemented".into()))
        }
    }

    #[async_trait::async_trait]
    impl querymt::embedding::EmbeddingProvider for MockProvider {
        async fn embed(&self, _input: Vec<String>) -> Result<Vec<Vec<f32>>, LLMError> {
            Err(LLMError::NotImplemented("Not implemented".into()))
        }
    }

    impl LLMProvider for MockProvider {}

    // ── resolve_api_key_with_preference tests ─────────────────────────────────

    /// Use a unique env var name per test to avoid cross-test interference.
    fn unique_env_var(suffix: &str) -> String {
        format!("QMT_TEST_API_KEY_{}", suffix)
    }

    #[tokio::test]
    async fn resolve_env_var_preferred() {
        let var = unique_env_var("ENV_PREF");
        // SAFETY: test-only; each test uses a unique env var name.
        unsafe { std::env::set_var(&var, "key-from-env") };

        let mut use_oauth = false;
        let result = super::resolve_api_key_with_preference(
            "test-provider",
            Some(&var),
            Some(&AuthMethod::EnvVar),
            &mut use_oauth,
        )
        .await;

        assert_eq!(result.as_deref(), Some("key-from-env"));
        assert!(!use_oauth);
        unsafe { std::env::remove_var(&var) };
    }

    #[tokio::test]
    async fn resolve_env_var_as_fallback_no_preference() {
        let var = unique_env_var("ENV_FALLBACK");
        // SAFETY: test-only; each test uses a unique env var name.
        unsafe { std::env::set_var(&var, "fallback-key") };

        let mut use_oauth = false;
        let result = super::resolve_api_key_with_preference(
            "test-provider",
            Some(&var),
            None, // default order: OAuth → stored key → env var
            &mut use_oauth,
        )
        .await;

        // OAuth and stored key will fail, should fall back to env var
        assert_eq!(result.as_deref(), Some("fallback-key"));
        unsafe { std::env::remove_var(&var) };
    }

    #[tokio::test]
    async fn resolve_returns_none_when_no_sources() {
        let var = unique_env_var("NONE_EXIST");
        // SAFETY: test-only; each test uses a unique env var name.
        unsafe { std::env::remove_var(&var) };

        let mut use_oauth = false;
        let result = super::resolve_api_key_with_preference(
            "no-such-provider",
            Some(&var),
            None,
            &mut use_oauth,
        )
        .await;

        assert!(result.is_none());
        assert!(!use_oauth);
    }

    #[tokio::test]
    async fn resolve_oauth_only_provider_no_env_var_name() {
        // OAuth-only providers pass env_var_name=None.
        // Without a valid OAuth token, should return None.
        let mut use_oauth = false;
        let result = super::resolve_api_key_with_preference(
            "codex-like",
            None, // no env var name
            None,
            &mut use_oauth,
        )
        .await;

        assert!(result.is_none());
    }

    #[test]
    fn static_pricing_omits_non_billable_oauth_providers() {
        assert!(SessionProvider::get_pricing("codex", "gpt-5").is_none());
        assert!(SessionProvider::get_pricing("kimi-code", "kimi-k2.5").is_none());
    }

    #[test]
    fn helper_marks_non_billable_oauth_providers() {
        assert!(is_non_billable_oauth_provider("codex"));
        assert!(is_non_billable_oauth_provider("kimi-code"));
        assert!(!is_non_billable_oauth_provider("openai"));
    }
}
