use super::traits::{ChatMessage, ChatRequest, ChatResponse};
use super::Provider;
use crate::channels::media::ImageRejectionReason;
use async_trait::async_trait;
use std::collections::HashMap;

/// A single route: maps a task hint to a provider + model combo.
#[derive(Debug, Clone)]
pub struct Route {
    pub provider_name: String,
    pub model: String,
}

/// Multi-model router — routes requests to different provider+model combos
/// based on a task hint encoded in the model parameter.
///
/// The model parameter can be:
/// - A regular model name (e.g. "anthropic/claude-sonnet-4") → uses default provider
/// - A hint-prefixed string (e.g. "hint:reasoning") → resolves via route table
///
/// This wraps multiple pre-created providers and selects the right one per request.
pub struct RouterProvider {
    routes: HashMap<String, (usize, String)>, // hint → (provider_index, model)
    providers: Vec<(String, Box<dyn Provider>)>,
    /// Reverse map built at construction: provider_index → sorted hint names.
    /// Used during warmup to report which route hints are impacted when a
    /// provider fails to initialise.
    provider_hints: HashMap<usize, Vec<String>>,
    default_index: usize,
    default_model: String,
}

impl RouterProvider {
    /// Create a new router with a default provider and optional routes.
    ///
    /// `providers` is a list of (name, provider) pairs. The first one is the default.
    /// `routes` maps hint names to Route structs containing provider_name and model.
    pub fn new(
        providers: Vec<(String, Box<dyn Provider>)>,
        routes: Vec<(String, Route)>,
        default_model: String,
    ) -> Self {
        // Build provider name → index lookup
        let name_to_index: HashMap<&str, usize> = providers
            .iter()
            .enumerate()
            .map(|(i, (name, _))| (name.as_str(), i))
            .collect();

        // Resolve routes to provider indices
        let resolved_routes: HashMap<String, (usize, String)> = routes
            .into_iter()
            .filter_map(|(hint, route)| {
                let index = name_to_index.get(route.provider_name.as_str()).copied();
                match index {
                    Some(i) => Some((hint, (i, route.model))),
                    None => {
                        tracing::warn!(
                            hint = hint,
                            provider = route.provider_name,
                            "Route references unknown provider, skipping"
                        );
                        None
                    }
                }
            })
            .collect();

        // Build reverse map: provider_index → sorted list of hint names.
        let mut provider_hints: HashMap<usize, Vec<String>> = HashMap::new();
        for (hint, (provider_idx, _)) in &resolved_routes {
            provider_hints
                .entry(*provider_idx)
                .or_default()
                .push(hint.clone());
        }
        for hints in provider_hints.values_mut() {
            hints.sort_unstable();
        }

        Self {
            routes: resolved_routes,
            providers,
            provider_hints,
            default_index: 0,
            default_model,
        }
    }

    /// Resolve a model parameter to a (provider, actual_model) pair.
    ///
    /// If the model starts with "hint:", look up the hint in the route table.
    /// Otherwise, use the default provider with the given model name.
    /// Resolve a model parameter to a (provider_index, actual_model) pair.
    fn resolve(&self, model: &str) -> (usize, String) {
        if let Some(hint) = model.strip_prefix("hint:") {
            if let Some((idx, resolved_model)) = self.routes.get(hint) {
                return (*idx, resolved_model.clone());
            }
            tracing::warn!(
                hint = hint,
                fallback_model = model,
                "Unknown route hint, falling back to default provider with raw model string"
            );
        }

        // Not a hint or hint not found — use default provider with the model as-is
        (self.default_index, model.to_string())
    }
}

#[async_trait]
impl Provider for RouterProvider {
    fn capabilities(&self) -> super::traits::ProviderCapabilities {
        let mut merged = super::traits::ProviderCapabilities::default();
        for (_, provider) in &self.providers {
            let caps = provider.capabilities();
            merged.native_tool_calling = merged.native_tool_calling || caps.native_tool_calling;
            merged.image_input = merged.image_input || caps.image_input;
            for form in caps.image_transport_forms {
                if !merged.image_transport_forms.contains(&form) {
                    merged.image_transport_forms.push(form);
                }
            }
        }
        merged
    }

    async fn chat_with_system(
        &self,
        system_prompt: Option<&str>,
        message: &str,
        model: &str,
        temperature: f64,
    ) -> anyhow::Result<String> {
        let (provider_idx, resolved_model) = self.resolve(model);

        let (provider_name, provider) = &self.providers[provider_idx];
        tracing::info!(
            provider = provider_name.as_str(),
            model = resolved_model.as_str(),
            "Router dispatching request"
        );

        provider
            .chat_with_system(system_prompt, message, &resolved_model, temperature)
            .await
    }

    async fn chat_with_history(
        &self,
        messages: &[ChatMessage],
        model: &str,
        temperature: f64,
    ) -> anyhow::Result<String> {
        let (provider_idx, resolved_model) = self.resolve(model);
        let (_, provider) = &self.providers[provider_idx];
        provider
            .chat_with_history(messages, &resolved_model, temperature)
            .await
    }

    async fn chat(
        &self,
        request: ChatRequest<'_>,
        model: &str,
        temperature: f64,
    ) -> anyhow::Result<ChatResponse> {
        let (provider_idx, resolved_model) = self.resolve(model);
        let (provider_name, provider) = &self.providers[provider_idx];

        // Fail-closed: reject image turns to text-only providers.
        if !request.images.is_empty() && !provider.capabilities().supports_image_input() {
            tracing::warn!(
                provider = provider_name.as_str(),
                model = resolved_model.as_str(),
                "Rejected image turn for non-image-capable route"
            );
            return Err(ImageRejectionReason::RouteNotImageCapable.into());
        }

        provider.chat(request, &resolved_model, temperature).await
    }

    async fn chat_with_tools(
        &self,
        messages: &[ChatMessage],
        tools: &[serde_json::Value],
        model: &str,
        temperature: f64,
    ) -> anyhow::Result<ChatResponse> {
        let (provider_idx, resolved_model) = self.resolve(model);
        let (_, provider) = &self.providers[provider_idx];
        provider
            .chat_with_tools(messages, tools, &resolved_model, temperature)
            .await
    }

    fn supports_native_tools(&self) -> bool {
        self.providers
            .get(self.default_index)
            .map(|(_, p)| p.supports_native_tools())
            .unwrap_or(false)
    }

    async fn warmup(&self) -> anyhow::Result<()> {
        for (i, (name, provider)) in self.providers.iter().enumerate() {
            tracing::info!(provider = name, "Warming up routed provider");
            if let Err(e) = provider.warmup().await {
                let affected_hints = match self.provider_hints.get(&i) {
                    Some(hints) if !hints.is_empty() => hints.join(", "),
                    _ => "none".to_string(),
                };
                tracing::warn!(
                    provider = name,
                    affected_hints = affected_hints,
                    "Warmup failed; routes using this provider may not be available: {e}"
                );
            }
        }
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::test_support::tracing_capture::capture_tracing_events;
    use std::sync::atomic::{AtomicUsize, Ordering};
    use std::sync::Arc;

    struct MockProvider {
        calls: Arc<AtomicUsize>,
        chat_calls: Arc<AtomicUsize>,
        response: &'static str,
        last_model: parking_lot::Mutex<String>,
    }

    struct MockProviderHandle(Arc<MockProvider>);

    impl MockProvider {
        fn new(response: &'static str) -> Self {
            Self {
                calls: Arc::new(AtomicUsize::new(0)),
                chat_calls: Arc::new(AtomicUsize::new(0)),
                response,
                last_model: parking_lot::Mutex::new(String::new()),
            }
        }

        fn call_count(&self) -> usize {
            self.calls.load(Ordering::SeqCst)
        }

        fn chat_call_count(&self) -> usize {
            self.chat_calls.load(Ordering::SeqCst)
        }

        fn last_model(&self) -> String {
            self.last_model.lock().clone()
        }
    }

    #[async_trait]
    impl Provider for MockProvider {
        async fn chat_with_system(
            &self,
            _system_prompt: Option<&str>,
            _message: &str,
            model: &str,
            _temperature: f64,
        ) -> anyhow::Result<String> {
            self.calls.fetch_add(1, Ordering::SeqCst);
            *self.last_model.lock() = model.to_string();
            Ok(self.response.to_string())
        }

        async fn chat(
            &self,
            request: ChatRequest<'_>,
            model: &str,
            temperature: f64,
        ) -> anyhow::Result<ChatResponse> {
            self.chat_calls.fetch_add(1, Ordering::SeqCst);
            let text = self
                .chat_with_history(request.messages, model, temperature)
                .await?;
            Ok(ChatResponse {
                text: Some(text),
                tool_calls: Vec::new(),
            })
        }
    }

    fn make_router(
        providers: Vec<(&'static str, &'static str)>,
        routes: Vec<(&str, &str, &str)>,
    ) -> (RouterProvider, Vec<Arc<MockProvider>>) {
        let mocks: Vec<Arc<MockProvider>> = providers
            .iter()
            .map(|(_, response)| Arc::new(MockProvider::new(response)))
            .collect();

        let provider_list: Vec<(String, Box<dyn Provider>)> = providers
            .iter()
            .zip(mocks.iter())
            .map(|((name, _), mock)| {
                (
                    name.to_string(),
                    Box::new(MockProviderHandle(Arc::clone(mock))) as Box<dyn Provider>,
                )
            })
            .collect();

        let route_list: Vec<(String, Route)> = routes
            .iter()
            .map(|(hint, provider_name, model)| {
                (
                    hint.to_string(),
                    Route {
                        provider_name: provider_name.to_string(),
                        model: model.to_string(),
                    },
                )
            })
            .collect();

        let router = RouterProvider::new(provider_list, route_list, "default-model".to_string());

        (router, mocks)
    }

    #[async_trait]
    impl Provider for MockProviderHandle {
        async fn chat_with_system(
            &self,
            system_prompt: Option<&str>,
            message: &str,
            model: &str,
            temperature: f64,
        ) -> anyhow::Result<String> {
            self.0
                .as_ref()
                .chat_with_system(system_prompt, message, model, temperature)
                .await
        }
    }

    #[tokio::test]
    async fn routes_hint_to_correct_provider() {
        let (router, mocks) = make_router(
            vec![("fast", "fast-response"), ("smart", "smart-response")],
            vec![
                ("fast", "fast", "llama-3-70b"),
                ("reasoning", "smart", "claude-opus"),
            ],
        );

        let result = router
            .simple_chat("hello", "hint:reasoning", 0.5)
            .await
            .unwrap();
        assert_eq!(result, "smart-response");
        assert_eq!(mocks[1].call_count(), 1);
        assert_eq!(mocks[1].last_model(), "claude-opus");
        assert_eq!(mocks[0].call_count(), 0);
    }

    #[tokio::test]
    async fn routes_fast_hint() {
        let (router, mocks) = make_router(
            vec![("fast", "fast-response"), ("smart", "smart-response")],
            vec![("fast", "fast", "llama-3-70b")],
        );

        let result = router.simple_chat("hello", "hint:fast", 0.5).await.unwrap();
        assert_eq!(result, "fast-response");
        assert_eq!(mocks[0].call_count(), 1);
        assert_eq!(mocks[0].last_model(), "llama-3-70b");
    }

    #[tokio::test]
    async fn unknown_hint_falls_back_to_default() {
        let (router, mocks) = make_router(
            vec![("default", "default-response"), ("other", "other-response")],
            vec![],
        );

        let result = router
            .simple_chat("hello", "hint:nonexistent", 0.5)
            .await
            .unwrap();
        assert_eq!(result, "default-response");
        assert_eq!(mocks[0].call_count(), 1);
        // Falls back to default with the hint as model name
        assert_eq!(mocks[0].last_model(), "hint:nonexistent");
    }

    #[tokio::test]
    async fn non_hint_model_uses_default_provider() {
        let (router, mocks) = make_router(
            vec![
                ("primary", "primary-response"),
                ("secondary", "secondary-response"),
            ],
            vec![("code", "secondary", "codellama")],
        );

        let result = router
            .simple_chat("hello", "anthropic/claude-sonnet-4-20250514", 0.5)
            .await
            .unwrap();
        assert_eq!(result, "primary-response");
        assert_eq!(mocks[0].call_count(), 1);
        assert_eq!(mocks[0].last_model(), "anthropic/claude-sonnet-4-20250514");
    }

    #[test]
    fn resolve_preserves_model_for_non_hints() {
        let (router, _) = make_router(vec![("default", "ok")], vec![]);

        let (idx, model) = router.resolve("gpt-4o");
        assert_eq!(idx, 0);
        assert_eq!(model, "gpt-4o");
    }

    #[test]
    fn resolve_strips_hint_prefix() {
        let (router, _) = make_router(
            vec![("fast", "ok"), ("smart", "ok")],
            vec![("reasoning", "smart", "claude-opus")],
        );

        let (idx, model) = router.resolve("hint:reasoning");
        assert_eq!(idx, 1);
        assert_eq!(model, "claude-opus");
    }

    #[test]
    fn skips_routes_with_unknown_provider() {
        let (router, _) = make_router(
            vec![("default", "ok")],
            vec![("broken", "nonexistent", "model")],
        );

        // Route should not exist
        assert!(!router.routes.contains_key("broken"));
    }

    #[test]
    fn unknown_hint_warning_includes_raw_fallback_model() {
        let (router, _) = make_router(vec![("default", "ok")], vec![]);

        let ((_idx, model), events) = capture_tracing_events(|| router.resolve("hint:nonexistent"));

        assert_eq!(model, "hint:nonexistent");
        let warning = events.iter().find(|event| {
            event
                .field("message")
                .is_some_and(|message| message.contains("Unknown route hint"))
        });
        assert!(warning.is_some());
        assert_eq!(warning.unwrap().field("hint"), Some("nonexistent"));
        assert_eq!(
            warning.unwrap().field("fallback_model"),
            Some("hint:nonexistent")
        );
    }

    #[test]
    fn known_hint_does_not_trigger_warning() {
        let (router, _) = make_router(
            vec![("default", "ok"), ("smart", "ok")],
            vec![("reasoning", "smart", "claude-opus")],
        );

        let ((_idx, model), events) = capture_tracing_events(|| router.resolve("hint:reasoning"));

        assert_eq!(model, "claude-opus");
        assert!(events.iter().all(|event| {
            !event
                .field("message")
                .is_some_and(|message| message.contains("Unknown route hint"))
        }));
    }

    #[test]
    fn non_hint_selector_does_not_trigger_warning() {
        let (router, _) = make_router(vec![("default", "ok")], vec![]);

        let ((_idx, model), events) = capture_tracing_events(|| router.resolve("gpt-4o"));

        assert_eq!(model, "gpt-4o");
        assert!(events.iter().all(|event| {
            !event
                .field("message")
                .is_some_and(|message| message.contains("Unknown route hint"))
        }));
    }

    #[tokio::test]
    async fn warmup_calls_all_providers() {
        let (router, _) = make_router(vec![("a", "ok"), ("b", "ok")], vec![]);

        // Warmup should not error
        assert!(router.warmup().await.is_ok());
    }

    struct FailingWarmupProvider;

    #[async_trait]
    impl Provider for FailingWarmupProvider {
        async fn chat_with_system(
            &self,
            _system_prompt: Option<&str>,
            _message: &str,
            _model: &str,
            _temperature: f64,
        ) -> anyhow::Result<String> {
            Ok("ok".into())
        }

        async fn warmup(&self) -> anyhow::Result<()> {
            Err(anyhow::anyhow!("simulated warmup failure"))
        }
    }

    #[tokio::test]
    async fn warmup_logs_affected_route_hints_on_failure() {
        use crate::test_support::tracing_capture::CaptureLayer;
        use tracing_subscriber::prelude::*;

        let layer = CaptureLayer::default();
        let subscriber = tracing_subscriber::registry().with(layer.clone());
        let _guard = tracing::subscriber::set_default(subscriber);

        let router = RouterProvider::new(
            vec![
                (
                    "default".into(),
                    Box::new(MockProvider::new("ok")) as Box<dyn Provider>,
                ),
                (
                    "failing".into(),
                    Box::new(FailingWarmupProvider) as Box<dyn Provider>,
                ),
            ],
            vec![
                (
                    "code".into(),
                    Route {
                        provider_name: "failing".into(),
                        model: "codellama".into(),
                    },
                ),
                (
                    "reasoning".into(),
                    Route {
                        provider_name: "failing".into(),
                        model: "o1-preview".into(),
                    },
                ),
            ],
            "default-model".into(),
        );

        assert!(router.warmup().await.is_ok());

        let events = layer.snapshot();
        let warning = events.iter().find(|event| {
            event
                .field("message")
                .is_some_and(|m| m.contains("Warmup failed"))
        });
        assert!(warning.is_some(), "expected warmup failure warning");
        let warning = warning.unwrap();
        assert_eq!(warning.field("provider"), Some("failing"));
        let hints = warning
            .field("affected_hints")
            .expect("expected affected_hints field");
        assert!(hints.contains("code"), "expected 'code' in affected hints");
        assert!(
            hints.contains("reasoning"),
            "expected 'reasoning' in affected hints"
        );
    }

    #[tokio::test]
    async fn warmup_reports_no_affected_hints_when_provider_has_no_routes() {
        use crate::test_support::tracing_capture::CaptureLayer;
        use tracing_subscriber::prelude::*;

        let layer = CaptureLayer::default();
        let subscriber = tracing_subscriber::registry().with(layer.clone());
        let _guard = tracing::subscriber::set_default(subscriber);

        // FailingWarmupProvider is registered but has no routes pointing to it
        let router = RouterProvider::new(
            vec![
                (
                    "default".into(),
                    Box::new(MockProvider::new("ok")) as Box<dyn Provider>,
                ),
                (
                    "failing".into(),
                    Box::new(FailingWarmupProvider) as Box<dyn Provider>,
                ),
            ],
            vec![],
            "default-model".into(),
        );

        assert!(router.warmup().await.is_ok());

        let events = layer.snapshot();
        let warning = events.iter().find(|event| {
            event
                .field("message")
                .is_some_and(|m| m.contains("Warmup failed"))
        });
        assert!(warning.is_some(), "expected warmup failure warning");
        assert_eq!(
            warning.unwrap().field("affected_hints"),
            Some("none"),
            "provider with no routes should report 'none' for affected_hints"
        );
    }

    #[tokio::test]
    async fn chat_with_system_passes_system_prompt() {
        let mock = Arc::new(MockProvider::new("response"));
        let router = RouterProvider::new(
            vec![(
                "default".into(),
                Box::new(MockProviderHandle(Arc::clone(&mock))) as Box<dyn Provider>,
            )],
            vec![],
            "model".into(),
        );

        let result = router
            .chat_with_system(Some("system"), "hello", "model", 0.5)
            .await
            .unwrap();
        assert_eq!(result, "response");
        assert_eq!(mock.call_count(), 1);
    }

    #[tokio::test]
    async fn chat_with_tools_delegates_to_resolved_provider() {
        let mock = Arc::new(MockProvider::new("tool-response"));
        let router = RouterProvider::new(
            vec![(
                "default".into(),
                Box::new(MockProviderHandle(Arc::clone(&mock))) as Box<dyn Provider>,
            )],
            vec![],
            "model".into(),
        );

        let messages = vec![ChatMessage {
            role: "user".to_string(),
            content: "use tools".to_string(),
            image_metadata: None,
            audio_metadata: None,
        }];
        let tools = vec![serde_json::json!({
            "type": "function",
            "function": {
                "name": "shell",
                "description": "Run shell command",
                "parameters": {}
            }
        })];

        // chat_with_tools should delegate through the router to the mock.
        // MockProvider's default chat_with_tools calls chat_with_history -> chat_with_system.
        let result = router
            .chat_with_tools(&messages, &tools, "model", 0.7)
            .await
            .unwrap();
        assert_eq!(result.text.as_deref(), Some("tool-response"));
        assert_eq!(mock.call_count(), 1);
        assert_eq!(mock.last_model(), "model");
    }

    #[tokio::test]
    async fn chat_with_tools_routes_hint_correctly() {
        let (router, mocks) = make_router(
            vec![("fast", "fast-tool"), ("smart", "smart-tool")],
            vec![("reasoning", "smart", "claude-opus")],
        );

        let messages = vec![ChatMessage {
            role: "user".to_string(),
            content: "reason about this".to_string(),
            image_metadata: None,
            audio_metadata: None,
        }];
        let tools = vec![serde_json::json!({"type": "function", "function": {"name": "test"}})];

        let result = router
            .chat_with_tools(&messages, &tools, "hint:reasoning", 0.5)
            .await
            .unwrap();
        assert_eq!(result.text.as_deref(), Some("smart-tool"));
        assert_eq!(mocks[1].call_count(), 1);
        assert_eq!(mocks[1].last_model(), "claude-opus");
        assert_eq!(mocks[0].call_count(), 0);
    }

    // ── §2.2 Image routing tests ─────────────────────────────

    #[tokio::test]
    async fn chat_rejects_image_turn_to_text_only_provider() {
        use crate::channels::media::{AllowedImageMime, ImageTransportForm, StagedImage};
        use std::path::PathBuf;

        let (router, mocks) = make_router(vec![("default", "ok")], vec![]);

        let staged = StagedImage {
            sha256: "abc".into(),
            mime_type: AllowedImageMime::Jpeg,
            byte_len: 100,
            temp_path: PathBuf::from("/tmp/fake.jpg"),
            transport_form: ImageTransportForm::InlineBytes,
            channel_origin: "test".into(),
        };

        let images = [staged];
        let request = super::super::traits::ChatRequest {
            messages: &[ChatMessage::user("describe this")],
            tools: None,
            images: &images,
        };

        let err = router
            .chat(request, "model", 0.5)
            .await
            .expect_err("should reject image turn");
        let rejection = err
            .downcast_ref::<ImageRejectionReason>()
            .expect("expected structured image rejection");
        assert_eq!(rejection, &ImageRejectionReason::RouteNotImageCapable);
        assert_eq!(mocks[0].chat_call_count(), 0);
    }

    #[tokio::test]
    async fn chat_text_only_passes_through_without_image_check() {
        let (router, mocks) = make_router(vec![("default", "text-ok")], vec![]);

        let request = super::super::traits::ChatRequest {
            messages: &[ChatMessage::user("hello")],
            tools: None,
            images: &[],
        };

        let result = router.chat(request, "model", 0.5).await.unwrap();
        assert_eq!(result.text.as_deref(), Some("text-ok"));
        assert_eq!(mocks[0].call_count(), 1);
    }

    #[test]
    fn capabilities_merges_all_providers() {
        let (router, _) = make_router(vec![("default", "ok")], vec![]);
        let caps = <RouterProvider as Provider>::capabilities(&router);
        // Default MockProvider has default capabilities (no image).
        assert!(!caps.supports_image_input());
    }
}
