use std::collections::HashMap;
use axum::{
    extract::{State, Query},
    http::StatusCode,
    response::{Json, sse::{Event, Sse}},
    routing::{get, post},
    Router,
};
use common_contracts::observability::{HealthStatus, ServiceStatus, TaskProgress};
use serde::Deserialize;
use uuid::Uuid;
use crate::state::AppState;
use crate::llm_client::LlmClient;
use futures::Stream; // Ensure futures is available

pub fn create_router(app_state: AppState) -> Router {
    Router::new()
        .route("/health", get(health_check))
        .route("/tasks", get(get_current_tasks))
        .route("/test-llm", post(test_llm_connection))
        .route("/analysis-results/stream", get(stream_analysis_results))
        .with_state(app_state)
}

/// [GET /health]
/// Provides the current health status of the module.
async fn health_check(State(_state): State<AppState>) -> Json<HealthStatus> {
    let mut details = HashMap::new();
    // In a real scenario, we would check connections to the message bus, etc.
    details.insert("message_bus_connection".to_string(), "ok".to_string());
    
    let status = HealthStatus {
        module_id: "report-generator-service".to_string(),
        status: ServiceStatus::Ok,
        version: env!("CARGO_PKG_VERSION").to_string(),
        details,
    };
    Json(status)
}

/// [GET /tasks]
/// Reports all currently processing tasks and their progress.
async fn get_current_tasks(State(state): State<AppState>) -> Json<Vec<TaskProgress>> {
    let tasks: Vec<TaskProgress> = state
        .tasks
        .iter()
        .map(|entry| entry.value().clone())
        .collect();
    Json(tasks)
}

#[derive(Deserialize)]
pub struct TestLlmRequest {
    pub api_base_url: String,
    pub api_key: String,
    pub model_id: String,
}

async fn test_llm_connection(
    State(_state): State<AppState>,
    Json(payload): Json<TestLlmRequest>,
) -> Result<Json<String>, (StatusCode, String)> {
    let client = LlmClient::new(
        payload.api_base_url,
        payload.api_key,
        payload.model_id,
        None,
    );

    let response = client.generate_text("Hello".to_string()).await
        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;

    Ok(Json(response))
}

#[derive(Deserialize)]
pub struct StreamQuery {
    pub request_id: Uuid,
}

/// [GET /analysis-results/stream]
/// SSE endpoint for streaming analysis results.
async fn stream_analysis_results(
    State(state): State<AppState>,
    Query(query): Query<StreamQuery>,
) -> Sse<impl Stream<Item = Result<Event, axum::Error>>> {
    
    let mut rx = state.streams.get_or_create_sender(query.request_id).subscribe();
    
    let stream = async_stream::stream! {
        loop {
             match rx.recv().await {
                 Ok(msg) => {
                     yield Ok(Event::default().data(msg));
                 },
                 Err(tokio::sync::broadcast::error::RecvError::Closed) => {
                     break;
                 },
                 Err(tokio::sync::broadcast::error::RecvError::Lagged(cnt)) => {
                     // Lagged, maybe log it. 
                     tracing::warn!("Stream lagged by {} messages", cnt);
                 }
             }
        }
    };

    Sse::new(stream).keep_alive(axum::response::sse::KeepAlive::default())
}
