a3s 0.8.2

use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::{Arc, RwLock};

use a3s_code_core::{Agent, AgentSession, CodeConfig, LlmClient};
use tokio::sync::Mutex;

use crate::budget::DEFAULT_CODE_WEB_EFFORT_ID;

#[derive(Debug, Clone)]
pub(in crate::api::code_web) struct CodeWebSessionControls {
    pub(in crate::api::code_web) effort: String,
    pub(in crate::api::code_web) goal: Option<String>,
}

impl Default for CodeWebSessionControls {
    fn default() -> Self {
        Self {
            effort: DEFAULT_CODE_WEB_EFFORT_ID.to_string(),
            goal: None,
        }
    }
}

#[derive(Clone, Default)]
pub(in crate::api::code_web) struct CodeWebSessionContext {
    pub(in crate::api::code_web) compact_summary: Option<String>,
    pub(in crate::api::code_web) auto_compact:
        Option<crate::compact::auto_compact::AutoCompactController>,
    pub(in crate::api::code_web) llm_client: Option<Arc<dyn LlmClient>>,
}

impl CodeWebSessionContext {
    pub(in crate::api::code_web) fn set_llm_client(&mut self, client: Arc<dyn LlmClient>) {
        self.llm_client = Some(client);
    }

    pub(in crate::api::code_web) fn llm_client(&self) -> Option<Arc<dyn LlmClient>> {
        self.llm_client.clone()
    }
}

#[derive(Debug, Clone)]
pub(in crate::api::code_web) struct CodeWebSessionSettings {
    pub(in crate::api::code_web) model: Option<String>,
    pub(in crate::api::code_web) follow_default_model: bool,
    pub(in crate::api::code_web) permission_mode: String,
    pub(in crate::api::code_web) planning_mode: Option<String>,
    pub(in crate::api::code_web) goal_tracking: Option<bool>,
}

impl Default for CodeWebSessionSettings {
    fn default() -> Self {
        Self {
            model: None,
            follow_default_model: true,
            permission_mode: "auto".to_string(),
            planning_mode: None,
            goal_tracking: None,
        }
    }
}

pub(in crate::api) struct CodeWebState {
    pub(in crate::api::code_web) agent: Arc<Agent>,
    pub(in crate::api::code_web) config_path: PathBuf,
    pub(in crate::api::code_web) auto_compact_threshold: f64,
    pub(in crate::api::code_web) default_workspace: PathBuf,
    pub(in crate::api::code_web) code_config: RwLock<CodeConfig>,
    pub(in crate::api::code_web) sessions: Mutex<HashMap<String, Arc<AgentSession>>>,
    pub(in crate::api::code_web) messages: Mutex<HashMap<String, Vec<serde_json::Value>>>,
    pub(in crate::api::code_web) session_controls: Mutex<HashMap<String, CodeWebSessionControls>>,
    pub(in crate::api::code_web) session_contexts: Mutex<HashMap<String, CodeWebSessionContext>>,
    pub(in crate::api::code_web) session_settings: Mutex<HashMap<String, CodeWebSessionSettings>>,
}

impl CodeWebState {
    pub(in crate::api) fn new(
        agent: Arc<Agent>,
        config_path: PathBuf,
        default_workspace: PathBuf,
        code_config: CodeConfig,
    ) -> Self {
        let auto_compact_threshold = crate::config::auto_compact_threshold_for_path(&config_path);
        Self {
            agent,
            config_path,
            auto_compact_threshold,
            default_workspace,
            code_config: RwLock::new(code_config),
            sessions: Mutex::new(HashMap::new()),
            messages: Mutex::new(HashMap::new()),
            session_controls: Mutex::new(HashMap::new()),
            session_contexts: Mutex::new(HashMap::new()),
            session_settings: Mutex::new(HashMap::new()),
        }
    }

    pub(in crate::api::code_web) fn code_config_snapshot(&self) -> CodeConfig {
        self.code_config
            .read()
            .map(|config| config.clone())
            .unwrap_or_default()
    }

    pub(in crate::api::code_web) fn current_default_model(&self) -> Option<String> {
        self.code_config_snapshot().default_model
    }

    pub(in crate::api::code_web) async fn close(&self) {
        self.agent.close().await;
    }
}

#[cfg(test)]
mod tests {
    use a3s_code_core::llm::{StreamEvent, ToolDefinition};
    use a3s_code_core::{LlmClient, LlmResponse, Message};
    use async_trait::async_trait;
    use tokio::sync::mpsc;
    use tokio_util::sync::CancellationToken;

    use super::*;

    struct TestClient;

    #[async_trait]
    impl LlmClient for TestClient {
        async fn complete(
            &self,
            _messages: &[Message],
            _system: Option<&str>,
            _tools: &[ToolDefinition],
        ) -> anyhow::Result<LlmResponse> {
            unreachable!("client identity test does not send requests")
        }

        async fn complete_streaming(
            &self,
            _messages: &[Message],
            _system: Option<&str>,
            _tools: &[ToolDefinition],
            _cancel_token: CancellationToken,
        ) -> anyhow::Result<mpsc::Receiver<StreamEvent>> {
            unreachable!("client identity test does not send requests")
        }
    }

    #[test]
    fn session_context_replaces_and_clears_llm_client_by_identity() {
        let first: Arc<dyn LlmClient> = Arc::new(TestClient);
        let second: Arc<dyn LlmClient> = Arc::new(TestClient);
        let mut context = CodeWebSessionContext::default();

        context.set_llm_client(Arc::clone(&first));
        assert!(Arc::ptr_eq(
            &first,
            &context.llm_client().expect("first client")
        ));

        context.set_llm_client(Arc::clone(&second));
        assert!(Arc::ptr_eq(
            &second,
            &context.llm_client().expect("replacement client")
        ));

        let mut contexts = HashMap::from([("session".to_string(), context)]);
        contexts.remove("session");
        assert!(!contexts.contains_key("session"));
    }
}