use std::sync::Arc;

use axum::extract::{Extension, Path};
use axum::http::StatusCode;
use axum::response::sse::{Event, Sse};
use futures_util::stream;
use mongodb::bson::doc;

use compliance_core::models::pentest::*;

use crate::agent::ComplianceAgent;

use super::super::dto::collect_cursor_async;

type AgentExt = Extension<Arc<ComplianceAgent>>;

/// GET /api/v1/pentest/sessions/:id/stream — SSE endpoint for real-time events
///
/// Returns recent messages as SSE events (polling approach).
/// True real-time streaming with broadcast channels will be added in a future iteration.
#[tracing::instrument(skip_all, fields(session_id = %id))]
pub async fn session_stream(
    Extension(agent): AgentExt,
    Path(id): Path<String>,
) -> Result<
    Sse<impl futures_util::Stream<Item = Result<Event, std::convert::Infallible>>>,
    StatusCode,
> {
    let oid = mongodb::bson::oid::ObjectId::parse_str(&id).map_err(|_| StatusCode::BAD_REQUEST)?;

    // Verify session exists
    let _session = agent
        .db
        .pentest_sessions()
        .find_one(doc! { "_id": oid })
        .await
        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?
        .ok_or(StatusCode::NOT_FOUND)?;

    // Fetch recent messages for this session
    let messages: Vec<PentestMessage> = match agent
        .db
        .pentest_messages()
        .find(doc! { "session_id": &id })
        .sort(doc! { "created_at": 1 })
        .limit(100)
        .await
    {
        Ok(cursor) => collect_cursor_async(cursor).await,
        Err(_) => Vec::new(),
    };

    // Fetch recent attack chain nodes
    let nodes: Vec<AttackChainNode> = match agent
        .db
        .attack_chain_nodes()
        .find(doc! { "session_id": &id })
        .sort(doc! { "started_at": 1 })
        .limit(100)
        .await
    {
        Ok(cursor) => collect_cursor_async(cursor).await,
        Err(_) => Vec::new(),
    };

    // Build SSE events from stored data
    let mut events: Vec<Result<Event, std::convert::Infallible>> = Vec::new();

    for msg in &messages {
        let event_data = serde_json::json!({
            "type": "message",
            "role": msg.role,
            "content": msg.content,
            "created_at": msg.created_at.to_rfc3339(),
        });
        if let Ok(data) = serde_json::to_string(&event_data) {
            events.push(Ok(Event::default().event("message").data(data)));
        }
    }

    for node in &nodes {
        let event_data = serde_json::json!({
            "type": "tool_execution",
            "node_id": node.node_id,
            "tool_name": node.tool_name,
            "status": node.status,
            "findings_produced": node.findings_produced,
        });
        if let Ok(data) = serde_json::to_string(&event_data) {
            events.push(Ok(Event::default().event("tool").data(data)));
        }
    }

    // Add session status event
    let session = agent
        .db
        .pentest_sessions()
        .find_one(doc! { "_id": oid })
        .await
        .ok()
        .flatten();

    if let Some(s) = session {
        let status_data = serde_json::json!({
            "type": "status",
            "status": s.status,
            "findings_count": s.findings_count,
            "tool_invocations": s.tool_invocations,
        });
        if let Ok(data) = serde_json::to_string(&status_data) {
            events.push(Ok(Event::default().event("status").data(data)));
        }
    }

    Ok(Sse::new(stream::iter(events)))
}
