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

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

use compliance_core::models::dast::DastTarget;
use compliance_core::models::finding::{Finding, FindingStatus, Severity};
use compliance_core::models::pentest::*;
use compliance_core::models::sbom::SbomEntry;
use compliance_core::traits::pentest_tool::PentestToolContext;
use compliance_dast::ToolRegistry;

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

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

pub struct PentestOrchestrator {
    tool_registry: ToolRegistry,
    llm: Arc<LlmClient>,
    db: Database,
    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,
        }
    }

    pub fn subscribe(&self) -> broadcast::Receiver<PentestEvent> {
        self.event_tx.subscribe()
    }

    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;

        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(tool_calls) => {
                    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: None,
                        tool_calls: Some(tc_requests),
                        tool_call_id: None,
                    });

                    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(),
                            String::new(),
                        );
                        node.status = AttackNodeStatus::Running;
                        node.started_at = Some(chrono::Utc::now());
                        let _ = self.db.attack_chain_nodes().insert_one(&node).await;

                        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;

                                    for mut finding in result.findings {
                                        finding.scan_run_id = session_id.clone();
                                        finding.session_id = Some(session_id.clone());
                                        let _ =
                                            self.db.dast_findings().insert_one(&finding).await;
                                        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(),
                                            });
                                    }

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

                                    let _ = self
                                        .db
                                        .attack_chain_nodes()
                                        .update_one(
                                            doc! {
                                                "session_id": &session_id,
                                                "node_id": &node_id,
                                            },
                                            doc! { "$set": {
                                                "status": "completed",
                                                "tool_output": mongodb::bson::to_bson(&result.data)
                                                    .unwrap_or(mongodb::bson::Bson::Null),
                                                "completed_at": mongodb::bson::DateTime::now(),
                                            }},
                                        )
                                        .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()),
                        });
                    }

                    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(())
    }

    // ── Code-Awareness: Gather context from linked repo ─────────

    /// Fetch SAST findings, SBOM entries (with CVEs), and code graph entry points
    /// for the repo linked to this DAST target.
    async fn gather_repo_context(
        &self,
        target: &DastTarget,
    ) -> (Vec<Finding>, Vec<SbomEntry>, Vec<CodeContextHint>) {
        let Some(repo_id) = &target.repo_id else {
            return (Vec::new(), Vec::new(), Vec::new());
        };

        let sast_findings = self.fetch_sast_findings(repo_id).await;
        let sbom_entries = self.fetch_vulnerable_sbom(repo_id).await;
        let code_context = self.fetch_code_context(repo_id, &sast_findings).await;

        tracing::info!(
            repo_id,
            sast_findings = sast_findings.len(),
            vulnerable_deps = sbom_entries.len(),
            code_hints = code_context.len(),
            "Gathered code-awareness context for pentest"
        );

        (sast_findings, sbom_entries, code_context)
    }

    /// Fetch open/triaged SAST findings for the repo (not false positives or resolved)
    async fn fetch_sast_findings(&self, repo_id: &str) -> Vec<Finding> {
        let cursor = self
            .db
            .findings()
            .find(doc! {
                "repo_id": repo_id,
                "status": { "$in": ["open", "triaged"] },
            })
            .sort(doc! { "severity": -1 })
            .limit(100)
            .await;

        match cursor {
            Ok(mut c) => {
                let mut results = Vec::new();
                while let Some(Ok(f)) = c.next().await {
                    results.push(f);
                }
                results
            }
            Err(e) => {
                tracing::warn!("Failed to fetch SAST findings for pentest: {e}");
                Vec::new()
            }
        }
    }

    /// Fetch SBOM entries that have known vulnerabilities
    async fn fetch_vulnerable_sbom(&self, repo_id: &str) -> Vec<SbomEntry> {
        let cursor = self
            .db
            .sbom_entries()
            .find(doc! {
                "repo_id": repo_id,
                "known_vulnerabilities": { "$exists": true, "$ne": [] },
            })
            .limit(50)
            .await;

        match cursor {
            Ok(mut c) => {
                let mut results = Vec::new();
                while let Some(Ok(e)) = c.next().await {
                    results.push(e);
                }
                results
            }
            Err(e) => {
                tracing::warn!("Failed to fetch vulnerable SBOM entries: {e}");
                Vec::new()
            }
        }
    }

    /// Build CodeContextHint objects from the code knowledge graph.
    /// Maps entry points to their source files and links SAST findings.
    async fn fetch_code_context(
        &self,
        repo_id: &str,
        sast_findings: &[Finding],
    ) -> Vec<CodeContextHint> {
        // Get entry point nodes from the code graph
        let cursor = self
            .db
            .graph_nodes()
            .find(doc! {
                "repo_id": repo_id,
                "is_entry_point": true,
            })
            .limit(50)
            .await;

        let nodes = match cursor {
            Ok(mut c) => {
                let mut results = Vec::new();
                while let Some(Ok(n)) = c.next().await {
                    results.push(n);
                }
                results
            }
            Err(_) => return Vec::new(),
        };

        // Build hints by matching graph nodes to SAST findings by file path
        nodes
            .into_iter()
            .map(|node| {
                // Find SAST findings in the same file
                let linked_vulns: Vec<String> = sast_findings
                    .iter()
                    .filter(|f| {
                        f.file_path.as_deref() == Some(&node.file_path)
                    })
                    .map(|f| {
                        format!(
                            "[{}] {}: {} (line {})",
                            f.severity,
                            f.scanner,
                            f.title,
                            f.line_number.unwrap_or(0)
                        )
                    })
                    .collect();

                CodeContextHint {
                    endpoint_pattern: node.qualified_name.clone(),
                    handler_function: node.name.clone(),
                    file_path: node.file_path.clone(),
                    code_snippet: String::new(), // Could fetch from embeddings
                    known_vulnerabilities: linked_vulns,
                }
            })
            .collect()
    }

    // ── System Prompt Builder ───────────────────────────────────

    async fn build_system_prompt(
        &self,
        session: &PentestSession,
        target: &DastTarget,
        sast_findings: &[Finding],
        sbom_entries: &[SbomEntry],
        code_context: &[CodeContextHint],
    ) -> String {
        let tool_names = self.tool_registry.list_names().join(", ");
        let strategy_guidance = match session.strategy {
            PentestStrategy::Quick => {
                "Focus on the most common and impactful vulnerabilities. Run a quick recon, then target the highest-risk areas."
            }
            PentestStrategy::Comprehensive => {
                "Perform a thorough assessment covering all vulnerability types. Start with recon, then systematically test each attack surface."
            }
            PentestStrategy::Targeted => {
                "Focus specifically on areas highlighted by SAST findings and known CVEs. Prioritize exploiting known weaknesses."
            }
            PentestStrategy::Aggressive => {
                "Use all available tools aggressively. Test with maximum payloads and attempt full exploitation."
            }
            PentestStrategy::Stealth => {
                "Minimize noise. Use fewer requests, avoid aggressive payloads. Focus on passive analysis and targeted probes."
            }
        };

        // Build SAST findings section
        let sast_section = if sast_findings.is_empty() {
            String::from("No SAST findings available for this target.")
        } else {
            let critical = sast_findings
                .iter()
                .filter(|f| f.severity == Severity::Critical)
                .count();
            let high = sast_findings
                .iter()
                .filter(|f| f.severity == Severity::High)
                .count();

            let mut section = format!(
                "{} open findings ({} critical, {} high):\n",
                sast_findings.len(),
                critical,
                high
            );

            // List the most important findings (critical/high first, up to 20)
            for f in sast_findings.iter().take(20) {
                let file_info = f
                    .file_path
                    .as_ref()
                    .map(|p| {
                        format!(
                            " in {}:{}",
                            p,
                            f.line_number.unwrap_or(0)
                        )
                    })
                    .unwrap_or_default();
                let status_note = match f.status {
                    FindingStatus::Triaged => " [TRIAGED]",
                    _ => "",
                };
                section.push_str(&format!(
                    "- [{sev}] {title}{file}{status}\n",
                    sev = f.severity,
                    title = f.title,
                    file = file_info,
                    status = status_note,
                ));
                if let Some(cwe) = &f.cwe {
                    section.push_str(&format!("  CWE: {cwe}\n"));
                }
            }
            if sast_findings.len() > 20 {
                section.push_str(&format!(
                    "... and {} more findings\n",
                    sast_findings.len() - 20
                ));
            }
            section
        };

        // Build SBOM/CVE section
        let sbom_section = if sbom_entries.is_empty() {
            String::from("No vulnerable dependencies identified.")
        } else {
            let mut section = format!(
                "{} dependencies with known vulnerabilities:\n",
                sbom_entries.len()
            );
            for entry in sbom_entries.iter().take(15) {
                let cve_ids: Vec<&str> = entry
                    .known_vulnerabilities
                    .iter()
                    .map(|v| v.id.as_str())
                    .collect();
                section.push_str(&format!(
                    "- {} {} ({}): {}\n",
                    entry.name,
                    entry.version,
                    entry.package_manager,
                    cve_ids.join(", ")
                ));
            }
            if sbom_entries.len() > 15 {
                section.push_str(&format!(
                    "... and {} more vulnerable dependencies\n",
                    sbom_entries.len() - 15
                ));
            }
            section
        };

        // Build code context section
        let code_section = if code_context.is_empty() {
            String::from("No code knowledge graph available for this target.")
        } else {
            let with_vulns = code_context
                .iter()
                .filter(|c| !c.known_vulnerabilities.is_empty())
                .count();

            let mut section = format!(
                "{} entry points identified ({} with linked SAST findings):\n",
                code_context.len(),
                with_vulns
            );

            for hint in code_context.iter().take(20) {
                section.push_str(&format!(
                    "- {} ({})\n",
                    hint.endpoint_pattern, hint.file_path
                ));
                for vuln in &hint.known_vulnerabilities {
                    section.push_str(&format!("  SAST: {vuln}\n"));
                }
            }
            section
        };

        format!(
            r#"You are an expert penetration tester conducting an authorized security assessment.

## Target
- **Name**: {target_name}
- **URL**: {base_url}
- **Type**: {target_type}
- **Rate Limit**: {rate_limit} req/s
- **Destructive Tests Allowed**: {allow_destructive}
- **Linked Repository**: {repo_linked}

## Strategy
{strategy_guidance}

## SAST Findings (Static Analysis)
{sast_section}

## Vulnerable Dependencies (SBOM)
{sbom_section}

## Code Entry Points (Knowledge Graph)
{code_section}

## Available Tools
{tool_names}

## Instructions
1. Start by running reconnaissance (recon tool) to fingerprint the target and discover technologies.
2. Run the OpenAPI parser to discover API endpoints from specs.
3. Check infrastructure: DNS, DMARC, TLS, security headers, cookies, CSP, CORS.
4. Based on SAST findings, prioritize testing endpoints where vulnerabilities were found in code.
5. For each vulnerability type found in SAST, use the corresponding DAST tool to verify exploitability.
6. If vulnerable dependencies are listed, try to trigger known CVE conditions against the running application.
7. Test rate limiting on critical endpoints (login, API).
8. Check for console.log leakage in frontend JavaScript.
9. Analyze tool results and chain findings — if one vulnerability enables others, explore the chain.
10. When testing is complete, provide a structured summary with severity and remediation.
11. Always explain your reasoning before invoking each tool.
12. When done, say "Testing complete" followed by a final summary.

## Important
- This is an authorized penetration test. All testing is permitted within the target scope.
- Respect the rate limit of {rate_limit} requests per second.
- Only use destructive tests if explicitly allowed ({allow_destructive}).
- Use SAST findings to guide your testing — they tell you WHERE in the code vulnerabilities exist.
- Use SBOM data to understand what technologies and versions the target runs.
"#,
            target_name = target.name,
            base_url = target.base_url,
            target_type = target.target_type,
            rate_limit = target.rate_limit,
            allow_destructive = target.allow_destructive,
            repo_linked = target.repo_id.as_deref().unwrap_or("None"),
        )
    }
}
