use async_trait::async_trait;
use kernelx_core::{models, prelude::*, Error};
use serde_json::Value;

#[derive(Clone, Debug)]
pub struct MockProvider {
    models: Vec<ModelInfo>,
    config: ModelConfig,
}

impl MockProvider {
    pub fn builder() -> ProviderBuilder<Self> {
        ProviderBuilder::new()
    }

    pub fn with_capabilities(capabilities: Vec<Capability>) -> Result<Self> {
        Self::builder()
            .api_key("test-key")
            .models(vec![ModelInfo {
                id: "mock-model".to_string(),
                capabilities,
            }])
            .build()
    }
}

impl Provider for MockProvider {
    fn models(&self) -> &[ModelInfo] {
        &self.models
    }

    fn from_builder(mut builder: ProviderBuilder<Self>) -> Result<Self> {
        if builder.get_api_key().is_none() {
            return Err(Error::Config("API key is required".into()));
        }

        Ok(Self {
            models: builder.take_models().unwrap_or_else(|| {
                models![
                    "mock-model" => [Capability::Chat, Capability::Complete, Capability::Structured]
                ]
            }),
            config: builder.take_config().unwrap_or_default(),
        })
    }
}

impl HasCapability for MockProvider {
    fn model_id(&self) -> &str {
        "mock"
    }

    fn config(&self) -> &ModelConfig {
        &self.config
    }
}

#[async_trait]
impl Chat for MockProvider {
    async fn chat_impl(
        &self,
        _model: &str,
        _messages: Vec<ChatMessage>,
        _config: &ModelConfig,
    ) -> Result<String> {
        Ok("Mock chat response".to_string())
    }
}

#[async_trait]
impl Complete for MockProvider {
    async fn complete_impl(
        &self,
        _model: &str,
        _prompt: &str,
        _config: &ModelConfig,
    ) -> Result<String> {
        Ok("Mock completion response".to_string())
    }
}

#[async_trait]
impl Structured for MockProvider {
    async fn structured_impl(
        &self,
        _model: &str,
        _prompt: &str,
        _schema: &Value,
        _config: &ModelConfig,
    ) -> Result<Value> {
        Ok(serde_json::json!({
            "message": "Mock structured response",
            "confidence": 0.95,
            "tags": ["test", "mock", "structured"],
            "setup": "Why did the mock cross the road?",
            "punchline": "To test the other side!",
            "genre": "Programming Humor"
        }))
    }
}

// Helper functions for tests
pub fn setup_mock_provider(capabilities: Vec<Capability>) -> Result<MockProvider> {
    MockProvider::with_capabilities(capabilities)
}

#[allow(dead_code)]
pub fn setup_mock_llm() -> Result<MockProvider> {
    setup_mock_provider(vec![Capability::Chat, Capability::Complete])
}

#[allow(dead_code)]
pub fn setup_mock_structured_llm() -> Result<MockProvider> {
    setup_mock_provider(vec![
        Capability::Chat,
        Capability::Complete,
        Capability::Structured,
    ])
}

#[allow(dead_code)]
pub fn setup_mock_provider_with_models(models: Vec<ModelInfo>) -> Result<MockProvider> {
    MockProvider::builder()
        .api_key("test-key")
        .models(models)
        .build()
}

#[allow(dead_code)]
pub fn create_model_info(id: &str, capabilities: Vec<Capability>) -> ModelInfo {
    ModelInfo {
        id: id.to_string(),
        capabilities,
    }
}