use std::sync::Arc;
use std::time::Duration;

use mongodb::bson::doc;
use tokio::sync::broadcast;

use compliance_core::models::dast::DastTarget;
use compliance_core::models::pentest::*;
use compliance_core::traits::pentest_tool::PentestToolContext;
use compliance_dast::ToolRegistry;

use crate::database::Database;
use crate::llm::{
    ChatMessage, LlmClient, LlmResponse, ToolCallRequest, ToolCallRequestFunction, ToolDefinition,
};

/// Maximum duration for a single pentest session before timeout
const SESSION_TIMEOUT: Duration = Duration::from_secs(30 * 60); // 30 minutes

pub struct PentestOrchestrator {
    pub(crate) tool_registry: ToolRegistry,
    pub(crate) llm: Arc<LlmClient>,
    pub(crate) db: Database,
    pub(crate) event_tx: broadcast::Sender<PentestEvent>,
}

impl PentestOrchestrator {
    pub fn new(llm: Arc<LlmClient>, db: Database) -> Self {
        let (event_tx, _) = broadcast::channel(256);
        Self {
            tool_registry: ToolRegistry::new(),
            llm,
            db,
            event_tx,
        }
    }

    #[allow(dead_code)]
    pub fn subscribe(&self) -> broadcast::Receiver<PentestEvent> {
        self.event_tx.subscribe()
    }

    #[allow(dead_code)]
    pub fn event_sender(&self) -> broadcast::Sender<PentestEvent> {
        self.event_tx.clone()
    }

    /// Run a pentest session with timeout and automatic failure marking on errors.
    pub async fn run_session_guarded(
        &self,
        session: &PentestSession,
        target: &DastTarget,
        initial_message: &str,
    ) {
        let session_id = session.id;

        match tokio::time::timeout(
            SESSION_TIMEOUT,
            self.run_session(session, target, initial_message),
        )
        .await
        {
            Ok(Ok(())) => {
                tracing::info!(?session_id, "Pentest session completed successfully");
            }
            Ok(Err(e)) => {
                tracing::error!(?session_id, error = %e, "Pentest session failed");
                self.mark_session_failed(session_id, &format!("Error: {e}"))
                    .await;
                let _ = self.event_tx.send(PentestEvent::Error {
                    message: format!("Session failed: {e}"),
                });
            }
            Err(_) => {
                tracing::warn!(?session_id, "Pentest session timed out after 30 minutes");
                self.mark_session_failed(session_id, "Session timed out after 30 minutes")
                    .await;
                let _ = self.event_tx.send(PentestEvent::Error {
                    message: "Session timed out after 30 minutes".to_string(),
                });
            }
        }
    }

    async fn mark_session_failed(
        &self,
        session_id: Option<mongodb::bson::oid::ObjectId>,
        reason: &str,
    ) {
        if let Some(sid) = session_id {
            let _ = self
                .db
                .pentest_sessions()
                .update_one(
                    doc! { "_id": sid },
                    doc! { "$set": {
                        "status": "failed",
                        "completed_at": mongodb::bson::DateTime::now(),
                        "error_message": reason,
                    }},
                )
                .await;
        }
    }

    async fn run_session(
        &self,
        session: &PentestSession,
        target: &DastTarget,
        initial_message: &str,
    ) -> Result<(), crate::error::AgentError> {
        let session_id = session.id.map(|oid| oid.to_hex()).unwrap_or_default();

        // Gather code-awareness context from linked repo
        let (sast_findings, sbom_entries, code_context) = self.gather_repo_context(target).await;

        // Build system prompt with code context
        let system_prompt = self
            .build_system_prompt(
                session,
                target,
                &sast_findings,
                &sbom_entries,
                &code_context,
            )
            .await;

        // Build tool definitions for LLM
        let tool_defs: Vec<ToolDefinition> = self
            .tool_registry
            .all_definitions()
            .into_iter()
            .map(|td| ToolDefinition {
                name: td.name,
                description: td.description,
                parameters: td.input_schema,
            })
            .collect();

        // Initialize messages
        let mut messages = vec![
            ChatMessage {
                role: "system".to_string(),
                content: Some(system_prompt),
                tool_calls: None,
                tool_call_id: None,
            },
            ChatMessage {
                role: "user".to_string(),
                content: Some(initial_message.to_string()),
                tool_calls: None,
                tool_call_id: None,
            },
        ];

        // Store user message
        let user_msg = PentestMessage::user(session_id.clone(), initial_message.to_string());
        let _ = self.db.pentest_messages().insert_one(&user_msg).await;

        // Build tool context with real data
        let tool_context = PentestToolContext {
            target: target.clone(),
            session_id: session_id.clone(),
            sast_findings,
            sbom_entries,
            code_context,
            rate_limit: target.rate_limit,
            allow_destructive: target.allow_destructive,
        };

        let max_iterations = 50;
        let mut total_findings = 0u32;
        let mut total_tool_calls = 0u32;
        let mut total_successes = 0u32;
        let mut prev_node_ids: Vec<String> = Vec::new();

        for _iteration in 0..max_iterations {
            let response = self
                .llm
                .chat_with_tools(messages.clone(), &tool_defs, Some(0.2), Some(8192))
                .await?;

            match response {
                LlmResponse::Content(content) => {
                    let msg = PentestMessage::assistant(session_id.clone(), content.clone());
                    let _ = self.db.pentest_messages().insert_one(&msg).await;
                    let _ = self.event_tx.send(PentestEvent::Message {
                        content: content.clone(),
                    });

                    messages.push(ChatMessage {
                        role: "assistant".to_string(),
                        content: Some(content.clone()),
                        tool_calls: None,
                        tool_call_id: None,
                    });

                    let done_indicators = [
                        "pentest complete",
                        "testing complete",
                        "scan complete",
                        "analysis complete",
                        "finished",
                        "that concludes",
                    ];
                    let content_lower = content.to_lowercase();
                    if done_indicators
                        .iter()
                        .any(|ind| content_lower.contains(ind))
                    {
                        break;
                    }
                    break;
                }
                LlmResponse::ToolCalls {
                    calls: tool_calls,
                    reasoning,
                } => {
                    let tc_requests: Vec<ToolCallRequest> = tool_calls
                        .iter()
                        .map(|tc| ToolCallRequest {
                            id: tc.id.clone(),
                            r#type: "function".to_string(),
                            function: ToolCallRequestFunction {
                                name: tc.name.clone(),
                                arguments: serde_json::to_string(&tc.arguments).unwrap_or_default(),
                            },
                        })
                        .collect();

                    messages.push(ChatMessage {
                        role: "assistant".to_string(),
                        content: if reasoning.is_empty() {
                            None
                        } else {
                            Some(reasoning.clone())
                        },
                        tool_calls: Some(tc_requests),
                        tool_call_id: None,
                    });

                    let mut current_batch_node_ids: Vec<String> = Vec::new();

                    for tc in &tool_calls {
                        total_tool_calls += 1;
                        let node_id = uuid::Uuid::new_v4().to_string();

                        let mut node = AttackChainNode::new(
                            session_id.clone(),
                            node_id.clone(),
                            tc.name.clone(),
                            tc.arguments.clone(),
                            reasoning.clone(),
                        );
                        // Link to previous iteration's nodes
                        node.parent_node_ids = prev_node_ids.clone();
                        node.status = AttackNodeStatus::Running;
                        node.started_at = Some(chrono::Utc::now());
                        let _ = self.db.attack_chain_nodes().insert_one(&node).await;
                        current_batch_node_ids.push(node_id.clone());

                        let _ = self.event_tx.send(PentestEvent::ToolStart {
                            node_id: node_id.clone(),
                            tool_name: tc.name.clone(),
                            input: tc.arguments.clone(),
                        });

                        let result = if let Some(tool) = self.tool_registry.get(&tc.name) {
                            match tool.execute(tc.arguments.clone(), &tool_context).await {
                                Ok(result) => {
                                    total_successes += 1;
                                    let findings_count = result.findings.len() as u32;
                                    total_findings += findings_count;

                                    let mut finding_ids: Vec<String> = Vec::new();
                                    for mut finding in result.findings {
                                        finding.scan_run_id = session_id.clone();
                                        finding.session_id = Some(session_id.clone());
                                        let insert_result =
                                            self.db.dast_findings().insert_one(&finding).await;
                                        if let Ok(res) = &insert_result {
                                            finding_ids.push(
                                                res.inserted_id
                                                    .as_object_id()
                                                    .map(|oid| oid.to_hex())
                                                    .unwrap_or_default(),
                                            );
                                        }
                                        let _ = self.event_tx.send(PentestEvent::Finding {
                                            finding_id: finding
                                                .id
                                                .map(|oid| oid.to_hex())
                                                .unwrap_or_default(),
                                            title: finding.title.clone(),
                                            severity: finding.severity.to_string(),
                                        });
                                    }

                                    // Compute risk score based on findings severity
                                    let risk_score: Option<u8> = if findings_count > 0 {
                                        Some(std::cmp::min(
                                            100,
                                            (findings_count as u8)
                                                .saturating_mul(15)
                                                .saturating_add(20),
                                        ))
                                    } else {
                                        None
                                    };

                                    let _ = self.event_tx.send(PentestEvent::ToolComplete {
                                        node_id: node_id.clone(),
                                        summary: result.summary.clone(),
                                        findings_count,
                                    });

                                    let finding_ids_bson: Vec<mongodb::bson::Bson> = finding_ids
                                        .iter()
                                        .map(|id| mongodb::bson::Bson::String(id.clone()))
                                        .collect();

                                    let mut update_doc = doc! {
                                        "status": "completed",
                                        "tool_output": mongodb::bson::to_bson(&result.data)
                                            .unwrap_or(mongodb::bson::Bson::Null),
                                        "completed_at": mongodb::bson::DateTime::now(),
                                        "findings_produced": finding_ids_bson,
                                    };
                                    if let Some(rs) = risk_score {
                                        update_doc.insert("risk_score", rs as i32);
                                    }

                                    let _ = self
                                        .db
                                        .attack_chain_nodes()
                                        .update_one(
                                            doc! {
                                                "session_id": &session_id,
                                                "node_id": &node_id,
                                            },
                                            doc! { "$set": update_doc },
                                        )
                                        .await;

                                    serde_json::json!({
                                        "summary": result.summary,
                                        "findings_count": findings_count,
                                        "data": result.data,
                                    })
                                    .to_string()
                                }
                                Err(e) => {
                                    let _ = self
                                        .db
                                        .attack_chain_nodes()
                                        .update_one(
                                            doc! {
                                                "session_id": &session_id,
                                                "node_id": &node_id,
                                            },
                                            doc! { "$set": {
                                                "status": "failed",
                                                "completed_at": mongodb::bson::DateTime::now(),
                                            }},
                                        )
                                        .await;
                                    format!("Tool execution failed: {e}")
                                }
                            }
                        } else {
                            format!("Unknown tool: {}", tc.name)
                        };

                        messages.push(ChatMessage {
                            role: "tool".to_string(),
                            content: Some(result),
                            tool_calls: None,
                            tool_call_id: Some(tc.id.clone()),
                        });
                    }

                    // Advance parent links so next iteration's nodes connect to this batch
                    prev_node_ids = current_batch_node_ids;

                    if let Some(sid) = session.id {
                        let _ = self
                            .db
                            .pentest_sessions()
                            .update_one(
                                doc! { "_id": sid },
                                doc! { "$set": {
                                    "tool_invocations": total_tool_calls as i64,
                                    "tool_successes": total_successes as i64,
                                    "findings_count": total_findings as i64,
                                }},
                            )
                            .await;
                    }
                }
            }
        }

        if let Some(sid) = session.id {
            let _ = self
                .db
                .pentest_sessions()
                .update_one(
                    doc! { "_id": sid },
                    doc! { "$set": {
                        "status": "completed",
                        "completed_at": mongodb::bson::DateTime::now(),
                        "tool_invocations": total_tool_calls as i64,
                        "tool_successes": total_successes as i64,
                        "findings_count": total_findings as i64,
                    }},
                )
                .await;
        }

        let _ = self.event_tx.send(PentestEvent::Complete {
            summary: format!(
                "Pentest complete. {} findings from {} tool invocations.",
                total_findings, total_tool_calls
            ),
        });

        Ok(())
    }
}
