use crate::client::ClientManager;
use crate::errors::AppError;
use crate::logging::{RequestLog, RequestLogger};
use crate::providers::{Provider, ProviderStreamChunk};
use crate::utils::tokens::estimate_completion_tokens;
use futures::stream::Stream;
use sqlx::Row;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};

/// Configuration for creating an AggregatingStream.
pub struct StreamConfig {
    pub client_id: String,
    pub provider: Arc<dyn Provider>,
    pub model: String,
    pub prompt_tokens: u32,
    pub has_images: bool,
    pub logger: Arc<RequestLogger>,
    pub client_manager: Arc<ClientManager>,
    pub model_registry: Arc<crate::models::registry::ModelRegistry>,
    pub db_pool: crate::database::DbPool,
}

pub struct AggregatingStream<S> {
    inner: S,
    client_id: String,
    provider: Arc<dyn Provider>,
    model: String,
    prompt_tokens: u32,
    has_images: bool,
    accumulated_content: String,
    accumulated_reasoning: String,
    logger: Arc<RequestLogger>,
    client_manager: Arc<ClientManager>,
    model_registry: Arc<crate::models::registry::ModelRegistry>,
    db_pool: crate::database::DbPool,
    start_time: std::time::Instant,
    has_logged: bool,
}

impl<S> AggregatingStream<S>
where
    S: Stream<Item = Result<ProviderStreamChunk, AppError>> + Unpin,
{
    pub fn new(inner: S, config: StreamConfig) -> Self {
        Self {
            inner,
            client_id: config.client_id,
            provider: config.provider,
            model: config.model,
            prompt_tokens: config.prompt_tokens,
            has_images: config.has_images,
            accumulated_content: String::new(),
            accumulated_reasoning: String::new(),
            logger: config.logger,
            client_manager: config.client_manager,
            model_registry: config.model_registry,
            db_pool: config.db_pool,
            start_time: std::time::Instant::now(),
            has_logged: false,
        }
    }

    fn finalize(&mut self) {
        if self.has_logged {
            return;
        }
        self.has_logged = true;

        let duration = self.start_time.elapsed();
        let client_id = self.client_id.clone();
        let provider_name = self.provider.name().to_string();
        let model = self.model.clone();
        let logger = self.logger.clone();
        let client_manager = self.client_manager.clone();
        let provider = self.provider.clone();
        let prompt_tokens = self.prompt_tokens;
        let has_images = self.has_images;
        let registry = self.model_registry.clone();
        let pool = self.db_pool.clone();

        // Estimate completion tokens (including reasoning if present)
        let content_tokens = estimate_completion_tokens(&self.accumulated_content, &model);
        let reasoning_tokens = if !self.accumulated_reasoning.is_empty() {
            estimate_completion_tokens(&self.accumulated_reasoning, &model)
        } else {
            0
        };

        let completion_tokens = content_tokens + reasoning_tokens;
        let total_tokens = prompt_tokens + completion_tokens;

        // Spawn a background task to log the completion
        tokio::spawn(async move {
            // Check database for cost overrides
            let db_cost =
                sqlx::query("SELECT prompt_cost_per_m, completion_cost_per_m FROM model_configs WHERE id = ?")
                    .bind(&model)
                    .fetch_optional(&pool)
                    .await
                    .unwrap_or(None);

            let cost = if let Some(row) = db_cost {
                let prompt_rate = row.get::<Option<f64>, _>("prompt_cost_per_m");
                let completion_rate = row.get::<Option<f64>, _>("completion_cost_per_m");

                if let (Some(p), Some(c)) = (prompt_rate, completion_rate) {
                    (prompt_tokens as f64 * p / 1_000_000.0) + (completion_tokens as f64 * c / 1_000_000.0)
                } else {
                    provider.calculate_cost(&model, prompt_tokens, completion_tokens, &registry)
                }
            } else {
                provider.calculate_cost(&model, prompt_tokens, completion_tokens, &registry)
            };

            // Log to database
            logger.log_request(RequestLog {
                timestamp: chrono::Utc::now(),
                client_id: client_id.clone(),
                provider: provider_name,
                model,
                prompt_tokens,
                completion_tokens,
                total_tokens,
                cost,
                has_images,
                status: "success".to_string(),
                error_message: None,
                duration_ms: duration.as_millis() as u64,
            });

            // Update client usage
            let _ = client_manager
                .update_client_usage(&client_id, total_tokens as i64, cost)
                .await;
        });
    }
}

impl<S> Stream for AggregatingStream<S>
where
    S: Stream<Item = Result<ProviderStreamChunk, AppError>> + Unpin,
{
    type Item = Result<ProviderStreamChunk, AppError>;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let result = Pin::new(&mut self.inner).poll_next(cx);

        match &result {
            Poll::Ready(Some(Ok(chunk))) => {
                self.accumulated_content.push_str(&chunk.content);
                if let Some(reasoning) = &chunk.reasoning_content {
                    self.accumulated_reasoning.push_str(reasoning);
                }
            }
            Poll::Ready(Some(Err(_))) => {
                // If there's an error, we might still want to log what we got so far?
                // For now, just finalize if we have content
                if !self.accumulated_content.is_empty() {
                    self.finalize();
                }
            }
            Poll::Ready(None) => {
                self.finalize();
            }
            Poll::Pending => {}
        }

        result
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use anyhow::Result;
    use futures::stream::{self, StreamExt};

    // Simple mock provider for testing
    struct MockProvider;
    #[async_trait::async_trait]
    impl Provider for MockProvider {
        fn name(&self) -> &str {
            "mock"
        }
        fn supports_model(&self, _model: &str) -> bool {
            true
        }
        fn supports_multimodal(&self) -> bool {
            false
        }
        async fn chat_completion(
            &self,
            _req: crate::models::UnifiedRequest,
        ) -> Result<crate::providers::ProviderResponse, AppError> {
            unimplemented!()
        }
        async fn chat_completion_stream(
            &self,
            _req: crate::models::UnifiedRequest,
        ) -> Result<futures::stream::BoxStream<'static, Result<ProviderStreamChunk, AppError>>, AppError> {
            unimplemented!()
        }
        fn estimate_tokens(&self, _req: &crate::models::UnifiedRequest) -> Result<u32> {
            Ok(10)
        }
        fn calculate_cost(&self, _model: &str, _p: u32, _c: u32, _r: &crate::models::registry::ModelRegistry) -> f64 {
            0.05
        }
    }

    #[tokio::test]
    async fn test_aggregating_stream() {
        let chunks = vec![
            Ok(ProviderStreamChunk {
                content: "Hello".to_string(),
                reasoning_content: None,
                finish_reason: None,
                model: "test".to_string(),
            }),
            Ok(ProviderStreamChunk {
                content: " World".to_string(),
                reasoning_content: None,
                finish_reason: Some("stop".to_string()),
                model: "test".to_string(),
            }),
        ];
        let inner_stream = stream::iter(chunks);

        let pool = sqlx::SqlitePool::connect("sqlite::memory:").await.unwrap();
        let (dashboard_tx, _) = tokio::sync::broadcast::channel(16);
        let logger = Arc::new(RequestLogger::new(pool.clone(), dashboard_tx));
        let client_manager = Arc::new(ClientManager::new(pool.clone()));
        let registry = Arc::new(crate::models::registry::ModelRegistry {
            providers: std::collections::HashMap::new(),
        });

        let mut agg_stream = AggregatingStream::new(
            inner_stream,
            StreamConfig {
                client_id: "client_1".to_string(),
                provider: Arc::new(MockProvider),
                model: "test".to_string(),
                prompt_tokens: 10,
                has_images: false,
                logger,
                client_manager,
                model_registry: registry,
                db_pool: pool.clone(),
            },
        );

        while let Some(item) = agg_stream.next().await {
            assert!(item.is_ok());
        }

        assert_eq!(agg_stream.accumulated_content, "Hello World");
        assert!(agg_stream.has_logged);
    }
}
