//! In-process agent service implementation.

use std::collections::HashMap;

use rig::completion::Message;
use tauri::{AppHandle, Runtime};

use crate::agent::{
    AgentConfigStatus, AgentRuntimeConfig, AgentStoredSettings, ChatPromptRequest,
    LocalProviderSettings, PreparedChatTurn, ProviderMode, RemoteProviderSettings, RigChatGateway,
    SaveAgentRuntimeConfigRequest, TaskProfile, UpdateRemoteApiKeyRequest,
};
use crate::error::AppError;

use super::gateway::ChatGateway;
use super::routing::{
    compute_local_configured, compute_overall_configured, compute_remote_configured,
    normalize_routes, resolve_runtime, validate_settings,
};
use super::settings::AgentSettingsService;

#[derive(Debug, Default)]
struct SessionManager {
    sessions: HashMap<String, Vec<Message>>,
    next_session_id: u64,
}

impl SessionManager {
    fn prepare_turn(
        &mut self,
        request: ChatPromptRequest,
        runtime: AgentRuntimeConfig,
    ) -> Result<PreparedChatTurn, AppError> {
        let prompt = request.prompt.trim();
        if prompt.is_empty() {
            return Err(AppError::EmptyPrompt);
        }

        let session_id = request.session_id.unwrap_or_else(|| {
            self.next_session_id += 1;
            format!("session-{}", self.next_session_id)
        });

        let history = self.sessions.entry(session_id.clone()).or_default();
        let prior_history = history.clone();
        history.push(Message::user(prompt));

        Ok(PreparedChatTurn {
            workspace_id: request.workspace_id,
            session_id,
            prompt: prompt.to_string(),
            history: prior_history,
            runtime,
        })
    }

    fn record_assistant_reply(&mut self, session_id: &str, reply: &str) -> Result<(), AppError> {
        let history = self
            .sessions
            .get_mut(session_id)
            .ok_or_else(|| AppError::UnknownSession(session_id.to_string()))?;

        history.push(Message::assistant(reply));
        Ok(())
    }
}

/// Stateful backend agent service combining settings, plaintext key storage, and session history.
#[derive(Debug)]
pub struct AgentService<R: Runtime, G: ChatGateway = RigChatGateway> {
    session_manager: SessionManager,
    settings: AgentSettingsService<R>,
    gateway: G,
}

impl<R: Runtime> AgentService<R, RigChatGateway> {
    /// Create a new agent service bound to the current Tauri application.
    pub fn new(app_handle: &AppHandle<R>) -> Result<Self, AppError> {
        Self::new_with_gateway(app_handle, RigChatGateway)
    }
}

impl<R: Runtime, G: ChatGateway> AgentService<R, G> {
    /// Create a new agent service with a caller-supplied gateway.
    pub fn new_with_gateway(app_handle: &AppHandle<R>, gateway: G) -> Result<Self, AppError> {
        Ok(Self {
            session_manager: SessionManager::default(),
            settings: AgentSettingsService::new(app_handle),
            gateway,
        })
    }

    /// Clone the configured chat gateway for work that must outlive the state lock.
    pub fn gateway(&self) -> G {
        self.gateway.clone()
    }

    /// Prepare a new chat turn, resolving provider settings and the stored API key.
    pub fn prepare_turn(
        &mut self,
        request: ChatPromptRequest,
    ) -> Result<PreparedChatTurn, AppError> {
        let runtime = self.resolve_runtime(
            request.agent_profile,
            request.provider_override,
            request.model_override.clone(),
        )?;
        self.session_manager.prepare_turn(request, runtime)
    }

    /// Record the assistant reply after the stream completes successfully.
    pub fn record_assistant_reply(
        &mut self,
        session_id: &str,
        reply: &str,
    ) -> Result<(), AppError> {
        self.session_manager
            .record_assistant_reply(session_id, reply)
    }

    /// Return the current public agent configuration status.
    pub fn get_config_status(&self) -> Result<AgentConfigStatus, AppError> {
        let settings = self.settings.load()?;
        Ok(self.build_status(settings))
    }

    /// Persist the provider and task routing configuration.
    pub fn save_runtime_config(
        &mut self,
        request: SaveAgentRuntimeConfigRequest,
    ) -> Result<AgentConfigStatus, AppError> {
        let mut settings = self.settings.load()?;
        settings.remote = RemoteProviderSettings {
            enabled: request.remote_enabled,
            base_url: request.remote_base_url.trim().to_string(),
            api_key: settings.remote.api_key,
        };
        settings.local = LocalProviderSettings {
            enabled: request.local_enabled,
            base_url: request.local_base_url.trim().to_string(),
            available_models: normalize_models(request.local_available_models),
        };
        settings.default_remote_model = request.default_remote_model.trim().to_string();
        settings.task_defaults = request.task_defaults;
        normalize_routes(&mut settings)?;
        validate_settings(&settings)?;

        let persisted = self.settings.save(settings)?;
        Ok(self.build_status(persisted))
    }

    /// Save or replace the plaintext remote API key.
    pub fn update_remote_api_key(
        &mut self,
        request: UpdateRemoteApiKeyRequest,
    ) -> Result<AgentConfigStatus, AppError> {
        let api_key = request.api_key.trim().to_string();
        if api_key.is_empty() {
            return Err(AppError::RemoteApiKeyMissing);
        }

        let settings = self.settings.set_remote_api_key(api_key)?;
        Ok(self.build_status(settings))
    }

    /// Remove the stored remote API key.
    pub fn clear_remote_api_key(&mut self) -> Result<AgentConfigStatus, AppError> {
        let settings = self.settings.set_remote_api_key(String::new())?;
        Ok(self.build_status(settings))
    }

    fn build_status(&self, settings: AgentStoredSettings) -> AgentConfigStatus {
        AgentConfigStatus {
            configured: compute_overall_configured(&settings),
            remote_configured: compute_remote_configured(&settings),
            local_configured: compute_local_configured(&settings),
            remote_enabled: settings.remote.enabled,
            local_enabled: settings.local.enabled,
            has_remote_api_key: !settings.remote.api_key.trim().is_empty(),
            remote_base_url: settings.remote.base_url,
            local_base_url: settings.local.base_url,
            default_remote_model: settings.default_remote_model,
            local_available_models: settings.local.available_models,
            task_defaults: settings.task_defaults,
        }
    }

    fn resolve_runtime(
        &self,
        task_profile: Option<TaskProfile>,
        provider_override: Option<ProviderMode>,
        model_override: Option<String>,
    ) -> Result<AgentRuntimeConfig, AppError> {
        let settings = self.settings.load()?;
        if !compute_overall_configured(&settings) {
            return Err(AppError::AgentNotConfigured);
        }

        resolve_runtime(
            &settings,
            task_profile.unwrap_or(TaskProfile::InteractiveChat),
            provider_override,
            model_override,
        )
    }
}

fn normalize_models(models: Vec<String>) -> Vec<String> {
    let mut normalized = Vec::new();

    for model in models {
        let trimmed = model.trim();
        if trimmed.is_empty() {
            continue;
        }

        if !normalized.iter().any(|existing| existing == trimmed) {
            normalized.push(trimmed.to_string());
        }
    }

    normalized
}

#[cfg(test)]
mod tests {
    use std::env;
    use std::fs;
    use std::path::PathBuf;
    use std::sync::{Mutex, OnceLock};
    use std::time::{SystemTime, UNIX_EPOCH};

    use super::SessionManager;
    use crate::agent::{
        default_task_defaults, AgentRuntimeConfig, AgentService, ChatPromptRequest, ProviderMode,
        SaveAgentRuntimeConfigRequest, TaskProfile, UpdateRemoteApiKeyRequest,
        DEFAULT_LOCAL_BASE_URL, DEFAULT_REMOTE_BASE_URL, DEFAULT_REMOTE_MODEL,
    };
    use crate::error::AppError;

    use rig::completion::Message;
    use tauri::test::{mock_builder, mock_context, noop_assets, MockRuntime};

    #[test]
    fn returns_empty_prompt_error_when_request_contains_only_whitespace() {
        let mut sessions = SessionManager::default();

        let result = sessions.prepare_turn(
            ChatPromptRequest {
                workspace_id: "workspace-1".to_string(),
                session_id: None,
                prompt: "   ".to_string(),
                agent_profile: None,
                model_override: None,
                provider_override: None,
            },
            sample_runtime(),
        );

        assert_eq!(result.unwrap_err(), AppError::EmptyPrompt);
    }

    #[test]
    fn creates_new_session_when_request_does_not_provide_one() {
        let mut sessions = SessionManager::default();

        let result = sessions
            .prepare_turn(
                ChatPromptRequest {
                    workspace_id: "workspace-1".to_string(),
                    session_id: None,
                    prompt: "Summarize AAPL".to_string(),
                    agent_profile: None,
                    model_override: None,
                    provider_override: None,
                },
                sample_runtime(),
            )
            .unwrap();

        assert_eq!(result.session_id, "session-1");
        assert!(result.history.is_empty());
    }

    #[test]
    fn reuses_existing_session_and_returns_prior_history() {
        let mut sessions = SessionManager::default();
        let session_id = "session-42".to_string();

        sessions
            .prepare_turn(
                ChatPromptRequest {
                    workspace_id: "workspace-1".to_string(),
                    session_id: Some(session_id.clone()),
                    prompt: "First prompt".to_string(),
                    agent_profile: None,
                    model_override: None,
                    provider_override: None,
                },
                sample_runtime(),
            )
            .unwrap();
        sessions
            .record_assistant_reply(&session_id, "First reply")
            .unwrap();

        let result = sessions
            .prepare_turn(
                ChatPromptRequest {
                    workspace_id: "workspace-1".to_string(),
                    session_id: Some(session_id),
                    prompt: "Second prompt".to_string(),
                    agent_profile: None,
                    model_override: None,
                    provider_override: None,
                },
                sample_runtime(),
            )
            .unwrap();

        assert_eq!(result.history.len(), 2);
        assert_eq!(result.history[0], Message::user("First prompt"));
        assert_eq!(result.history[1], Message::assistant("First reply"));
    }

    #[test]
    fn persists_settings_and_plaintext_api_key_in_store() {
        with_test_home("config", || {
            let app = build_test_app();
            let mut service = AgentService::new(&app.handle()).unwrap();

            let initial = service.get_config_status().unwrap();
            assert!(!initial.configured);
            assert!(!initial.has_remote_api_key);
            assert_eq!(initial.remote_base_url, DEFAULT_REMOTE_BASE_URL);
            assert_eq!(initial.local_base_url, DEFAULT_LOCAL_BASE_URL);
            assert_eq!(initial.default_remote_model, DEFAULT_REMOTE_MODEL);

            let saved = service
                .save_runtime_config(SaveAgentRuntimeConfigRequest {
                    remote_enabled: true,
                    remote_base_url: "https://example.test/v4".to_string(),
                    default_remote_model: "glm-test".to_string(),
                    local_enabled: true,
                    local_base_url: "http://127.0.0.1:1234".to_string(),
                    local_available_models: vec!["qwen-small".to_string()],
                    task_defaults: default_task_defaults("glm-test"),
                })
                .unwrap();
            assert_eq!(saved.remote_base_url, "https://example.test/v4");
            assert_eq!(saved.default_remote_model, "glm-test");
            assert!(!saved.has_remote_api_key);

            let updated = service
                .update_remote_api_key(UpdateRemoteApiKeyRequest {
                    api_key: "z-ai-key-1".to_string(),
                })
                .unwrap();
            assert!(updated.configured);
            assert!(updated.has_remote_api_key);

            let prepared = service
                .prepare_turn(ChatPromptRequest {
                    workspace_id: "workspace-1".to_string(),
                    session_id: None,
                    prompt: "hello".to_string(),
                    agent_profile: None,
                    model_override: None,
                    provider_override: None,
                })
                .unwrap();
            assert_eq!(prepared.runtime.base_url, "https://example.test/v4");
            assert_eq!(prepared.runtime.model, "glm-test");
            assert_eq!(prepared.runtime.provider_mode, ProviderMode::Remote);
            assert_eq!(prepared.runtime.api_key.as_deref(), Some("z-ai-key-1"));
        });
    }

    #[test]
    fn clears_plaintext_remote_api_key_from_store() {
        with_test_home("clear", || {
            let app = build_test_app();
            let mut service = AgentService::new(&app.handle()).unwrap();

            service
                .save_runtime_config(SaveAgentRuntimeConfigRequest {
                    remote_enabled: true,
                    remote_base_url: "https://example.test/v4".to_string(),
                    default_remote_model: "glm-test".to_string(),
                    local_enabled: false,
                    local_base_url: DEFAULT_LOCAL_BASE_URL.to_string(),
                    local_available_models: Vec::new(),
                    task_defaults: default_task_defaults("glm-test"),
                })
                .unwrap();

            service
                .update_remote_api_key(UpdateRemoteApiKeyRequest {
                    api_key: "z-ai-key-1".to_string(),
                })
                .unwrap();

            let cleared = service.clear_remote_api_key().unwrap();
            assert!(!cleared.configured);
            assert!(!cleared.has_remote_api_key);

            let result = service.prepare_turn(ChatPromptRequest {
                workspace_id: "workspace-1".to_string(),
                session_id: None,
                prompt: "hello".to_string(),
                agent_profile: None,
                model_override: None,
                provider_override: None,
            });
            assert_eq!(result.unwrap_err(), AppError::AgentNotConfigured);
        });
    }

    #[test]
    fn local_task_route_resolves_without_remote_api_key() {
        with_test_home("local-route", || {
            let app = build_test_app();
            let mut service = AgentService::new(&app.handle()).unwrap();
            let mut task_defaults = default_task_defaults("glm-test");
            task_defaults.insert(
                TaskProfile::InteractiveChat,
                crate::agent::AgentTaskRoute {
                    provider_mode: ProviderMode::Local,
                    model: "qwen-local".to_string(),
                },
            );

            service
                .save_runtime_config(SaveAgentRuntimeConfigRequest {
                    remote_enabled: true,
                    remote_base_url: "https://example.test/v4".to_string(),
                    default_remote_model: "glm-test".to_string(),
                    local_enabled: true,
                    local_base_url: "http://127.0.0.1:1234".to_string(),
                    local_available_models: vec!["qwen-local".to_string()],
                    task_defaults,
                })
                .unwrap();

            let prepared = service
                .prepare_turn(ChatPromptRequest {
                    workspace_id: "workspace-1".to_string(),
                    session_id: None,
                    prompt: "hello".to_string(),
                    agent_profile: Some(TaskProfile::InteractiveChat),
                    model_override: None,
                    provider_override: None,
                })
                .unwrap();

            assert_eq!(prepared.runtime.provider_mode, ProviderMode::Local);
            assert_eq!(prepared.runtime.model, "qwen-local");
            assert_eq!(prepared.runtime.api_key, None);
        });
    }

    #[test]
    fn provider_override_replaces_task_default() {
        with_test_home("provider-override", || {
            let app = build_test_app();
            let mut service = AgentService::new(&app.handle()).unwrap();
            let mut task_defaults = default_task_defaults("glm-test");
            task_defaults.insert(
                TaskProfile::InteractiveChat,
                crate::agent::AgentTaskRoute {
                    provider_mode: ProviderMode::Local,
                    model: "qwen-local".to_string(),
                },
            );

            service
                .save_runtime_config(SaveAgentRuntimeConfigRequest {
                    remote_enabled: true,
                    remote_base_url: "https://example.test/v4".to_string(),
                    default_remote_model: "glm-test".to_string(),
                    local_enabled: true,
                    local_base_url: "http://127.0.0.1:1234".to_string(),
                    local_available_models: vec!["qwen-local".to_string()],
                    task_defaults,
                })
                .unwrap();
            service
                .update_remote_api_key(UpdateRemoteApiKeyRequest {
                    api_key: "z-ai-key-1".to_string(),
                })
                .unwrap();

            let prepared = service
                .prepare_turn(ChatPromptRequest {
                    workspace_id: "workspace-1".to_string(),
                    session_id: None,
                    prompt: "hello".to_string(),
                    agent_profile: Some(TaskProfile::InteractiveChat),
                    model_override: None,
                    provider_override: Some(ProviderMode::Remote),
                })
                .unwrap();

            assert_eq!(prepared.runtime.provider_mode, ProviderMode::Remote);
            assert_eq!(prepared.runtime.model, "glm-test");
            assert_eq!(prepared.runtime.api_key.as_deref(), Some("z-ai-key-1"));
        });
    }

    #[test]
    fn model_override_replaces_task_default() {
        with_test_home("model-override", || {
            let app = build_test_app();
            let mut service = AgentService::new(&app.handle()).unwrap();
            service
                .save_runtime_config(SaveAgentRuntimeConfigRequest {
                    remote_enabled: true,
                    remote_base_url: "https://example.test/v4".to_string(),
                    default_remote_model: "glm-test".to_string(),
                    local_enabled: false,
                    local_base_url: DEFAULT_LOCAL_BASE_URL.to_string(),
                    local_available_models: Vec::new(),
                    task_defaults: default_task_defaults("glm-test"),
                })
                .unwrap();
            service
                .update_remote_api_key(UpdateRemoteApiKeyRequest {
                    api_key: "z-ai-key-1".to_string(),
                })
                .unwrap();

            let prepared = service
                .prepare_turn(ChatPromptRequest {
                    workspace_id: "workspace-1".to_string(),
                    session_id: None,
                    prompt: "hello".to_string(),
                    agent_profile: None,
                    model_override: Some("glm-override".to_string()),
                    provider_override: None,
                })
                .unwrap();

            assert_eq!(prepared.runtime.model, "glm-override");
        });
    }

    #[test]
    fn local_task_without_model_fails_validation() {
        with_test_home("local-validation", || {
            let app = build_test_app();
            let mut service = AgentService::new(&app.handle()).unwrap();
            let mut task_defaults = default_task_defaults("glm-test");
            task_defaults.insert(
                TaskProfile::InteractiveChat,
                crate::agent::AgentTaskRoute {
                    provider_mode: ProviderMode::Local,
                    model: String::new(),
                },
            );

            let result = service.save_runtime_config(SaveAgentRuntimeConfigRequest {
                remote_enabled: true,
                remote_base_url: "https://example.test/v4".to_string(),
                default_remote_model: "glm-test".to_string(),
                local_enabled: true,
                local_base_url: "http://127.0.0.1:1234".to_string(),
                local_available_models: Vec::new(),
                task_defaults,
            });

            assert_eq!(
                result.unwrap_err(),
                AppError::ModelMissing(TaskProfile::InteractiveChat)
            );
        });
    }

    #[test]
    fn remote_task_without_api_key_fails_validation_at_prepare_time() {
        with_test_home("remote-validation", || {
            let app = build_test_app();
            let mut service = AgentService::new(&app.handle()).unwrap();
            service
                .save_runtime_config(SaveAgentRuntimeConfigRequest {
                    remote_enabled: true,
                    remote_base_url: "https://example.test/v4".to_string(),
                    default_remote_model: "glm-test".to_string(),
                    local_enabled: false,
                    local_base_url: DEFAULT_LOCAL_BASE_URL.to_string(),
                    local_available_models: Vec::new(),
                    task_defaults: default_task_defaults("glm-test"),
                })
                .unwrap();

            let result = service.prepare_turn(ChatPromptRequest {
                workspace_id: "workspace-1".to_string(),
                session_id: None,
                prompt: "hello".to_string(),
                agent_profile: None,
                model_override: None,
                provider_override: None,
            });

            assert_eq!(result.unwrap_err(), AppError::AgentNotConfigured);
        });
    }

    fn sample_runtime() -> AgentRuntimeConfig {
        AgentRuntimeConfig {
            provider_mode: ProviderMode::Remote,
            base_url: "https://example.com".to_string(),
            model: "glm-5.1".to_string(),
            api_key: Some("key".to_string()),
            task_profile: TaskProfile::InteractiveChat,
        }
    }

    fn build_test_app() -> tauri::App<MockRuntime> {
        mock_builder()
            .plugin(tauri_plugin_store::Builder::new().build())
            .build(mock_context(noop_assets()))
            .unwrap()
    }

    fn unique_identifier(prefix: &str) -> String {
        let nanos = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_nanos();
        format!("com.mosaiciq.tests.{prefix}.{nanos}")
    }

    fn with_test_home<T>(prefix: &str, test: impl FnOnce() -> T) -> T {
        let _lock = env_lock().lock().unwrap();
        let home = env::temp_dir().join(unique_identifier(prefix));
        fs::create_dir_all(&home).unwrap();

        let original_home = env::var_os("HOME");
        env::set_var("HOME", &home);

        let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(test));

        match original_home {
            Some(value) => env::set_var("HOME", value),
            None => env::remove_var("HOME"),
        }

        cleanup_test_data_dir(home);

        match result {
            Ok(value) => value,
            Err(payload) => std::panic::resume_unwind(payload),
        }
    }

    fn env_lock() -> &'static Mutex<()> {
        static LOCK: OnceLock<Mutex<()>> = OnceLock::new();
        LOCK.get_or_init(|| Mutex::new(()))
    }

    fn cleanup_test_data_dir(path: PathBuf) {
        let _ = fs::remove_dir_all(path);
    }
}
