use std::collections::HashMap;

use bamboo_agent_core::tools::{
    normalize_tool_name, parse_tool_args, parse_tool_args_best_effort, ToolCall, ToolResult,
};
use bamboo_agent_core::{Role, Session};

const MAX_TOOL_CALLS_PER_ROUND: usize = 80;
const MAX_CONSECUTIVE_FAILURES_PER_TOOL: usize = 3;
const RESET_POLICY_TOOL_NAME: &str = "conclusion_with_options";
const COPILOT_CONCLUSION_WITH_OPTIONS_ENHANCEMENT_METADATA_KEY: &str =
    "copilot_conclusion_with_options_enhancement_enabled";

const STRICT_ARGUMENT_TOOL_NAMES: [&str; 11] = [
    "Write",
    "Edit",
    "NotebookEdit",
    "apply_patch",
    "Bash",
    "Task",
    "SubAgent",
    "scheduler",
    "sub_session_manager",
    "session_note",
    "memory_note",
];

fn normalize_tool_for_policy(raw_tool_name: &str) -> String {
    bamboo_tools::normalize_tool_ref(raw_tool_name)
        .unwrap_or_else(|| normalize_tool_name(raw_tool_name).trim().to_string())
}

fn is_copilot_conclusion_with_options_enhancement_enabled(session: &Session) -> bool {
    session
        .metadata
        .get(COPILOT_CONCLUSION_WITH_OPTIONS_ENHANCEMENT_METADATA_KEY)
        .is_some_and(|value| value.trim().eq_ignore_ascii_case("true"))
}

fn find_assistant_message_for_tool_call<'a>(
    session: &'a Session,
    tool_call_id: &str,
) -> Option<&'a bamboo_agent_core::Message> {
    session.messages.iter().rev().find(|message| {
        matches!(message.role, Role::Assistant)
            && message
                .tool_calls
                .as_ref()
                .map(|calls| calls.iter().any(|call| call.id == tool_call_id))
                .unwrap_or(false)
    })
}

fn conclusion_with_options_arguments_include_conclusion(tool_call: &ToolCall) -> bool {
    let (parsed, _) = parse_tool_args_best_effort(&tool_call.function.arguments);
    let Some(args) = parsed.as_object() else {
        return false;
    };

    let Some(conclusion) = args.get("conclusion").and_then(|value| value.as_object()) else {
        return false;
    };
    let has_summary = conclusion
        .get("summary")
        .and_then(|value| value.as_str())
        .is_some_and(|value| !value.trim().is_empty());
    let has_mermaid_graph = conclusion
        .get("mermaid")
        .and_then(|value| value.as_object())
        .and_then(|value| value.get("graph"))
        .and_then(|value| value.as_str())
        .is_some_and(|value| !value.trim().is_empty());

    has_summary && has_mermaid_graph
}

pub(super) fn validate_tool_call_arguments(tool_call: &ToolCall) -> Result<(), String> {
    let normalized_tool_name = normalize_tool_for_policy(&tool_call.function.name);
    if !STRICT_ARGUMENT_TOOL_NAMES
        .iter()
        .any(|name| name.eq_ignore_ascii_case(&normalized_tool_name))
    {
        return Ok(());
    }

    if normalized_tool_name.eq_ignore_ascii_case("session_note")
        || normalized_tool_name.eq_ignore_ascii_case("memory_note")
    {
        let (parsed, parse_warning) = parse_tool_args_best_effort(&tool_call.function.arguments);
        if parse_warning.is_some()
            && parsed
                .as_object()
                .map(|map| map.is_empty())
                .unwrap_or(false)
        {
            return Err(format!(
                "Tool policy blocked '{}' due to invalid JSON arguments: unable to recover arguments. Rewrite the {} call with a valid JSON object, for example {{\"action\":\"read\",\"topic\":\"default\"}}.",
                normalized_tool_name, normalized_tool_name,
            ));
        }
        return Ok(());
    }

    parse_tool_args(&tool_call.function.arguments).map_err(|error| {
        format!(
            "Tool policy blocked '{}' due to invalid JSON arguments: {}",
            normalized_tool_name, error
        )
    })?;

    Ok(())
}

pub(super) fn validate_tool_call_context(
    tool_call: &ToolCall,
    session: &Session,
) -> Result<(), String> {
    let normalized_tool_name = normalize_tool_for_policy(&tool_call.function.name);
    let assistant_message = find_assistant_message_for_tool_call(session, &tool_call.id);
    let enhancement_enabled = is_copilot_conclusion_with_options_enhancement_enabled(session);

    if normalized_tool_name.eq_ignore_ascii_case("conclusion_with_options") {
        let has_narration = assistant_message
            .map(|message| !message.content.trim().is_empty())
            .unwrap_or(false);

        if !has_narration {
            return Err(
                "Tool policy blocked 'conclusion_with_options': add a brief assistant text summary before calling conclusion_with_options so the user can understand the conclusion. Then retry conclusion_with_options instead of repeating the same blocked call. For example, first send a short plain-text wrap-up sentence, then call conclusion_with_options in the next assistant action."
                    .to_string(),
            );
        }

        if enhancement_enabled && !conclusion_with_options_arguments_include_conclusion(tool_call) {
            return Err(
                "Tool policy blocked 'conclusion_with_options': when copilot conclusion-with-options enhancement is enabled, include `conclusion.summary` and `conclusion.mermaid.graph` in the conclusion_with_options arguments. Rewrite the tool arguments to include those required fields, then retry conclusion_with_options."
                    .to_string(),
            );
        }

        return Ok(());
    }

    Ok(())
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub(super) enum ToolPolicyPrecheckViolation {
    RoundToolLimit {
        limit: usize,
        tool_name: String,
    },
    ToolCircuitOpen {
        tool_name: String,
        consecutive_failures: usize,
        limit: usize,
    },
}

impl ToolPolicyPrecheckViolation {
    pub(super) fn should_stop_round(&self) -> bool {
        matches!(self, Self::RoundToolLimit { .. })
    }

    pub(super) fn into_message(self) -> String {
        match self {
            Self::RoundToolLimit { limit, tool_name } => format!(
                "Tool policy blocked '{}': per-round tool call limit ({limit}) reached",
                tool_name
            ),
            Self::ToolCircuitOpen {
                tool_name,
                consecutive_failures,
                limit,
            } => format!(
                "Tool policy blocked '{}': {} consecutive failures reached circuit limit ({}) in this round",
                tool_name, consecutive_failures, limit
            ),
        }
    }
}

#[derive(Debug, Clone)]
pub(super) struct ToolPolicyGuard {
    max_tool_calls_per_round: usize,
    max_consecutive_failures_per_tool: usize,
    executed_calls: usize,
    consecutive_failures: HashMap<String, usize>,
}

impl ToolPolicyGuard {
    pub(super) fn new(
        max_tool_calls_per_round: usize,
        max_consecutive_failures_per_tool: usize,
    ) -> Self {
        Self {
            max_tool_calls_per_round,
            max_consecutive_failures_per_tool,
            executed_calls: 0,
            consecutive_failures: HashMap::new(),
        }
    }

    pub(super) fn check_before_execution(
        &self,
        tool_call: &ToolCall,
        reserved_calls: usize,
    ) -> Result<(), ToolPolicyPrecheckViolation> {
        let normalized_tool_name = normalize_tool_for_policy(&tool_call.function.name);
        let projected_executions = self.executed_calls.saturating_add(reserved_calls);

        if projected_executions >= self.max_tool_calls_per_round {
            return Err(ToolPolicyPrecheckViolation::RoundToolLimit {
                limit: self.max_tool_calls_per_round,
                tool_name: normalized_tool_name,
            });
        }

        let failures = self
            .consecutive_failures
            .get(&normalized_tool_name)
            .copied()
            .unwrap_or(0);
        if failures >= self.max_consecutive_failures_per_tool {
            return Err(ToolPolicyPrecheckViolation::ToolCircuitOpen {
                tool_name: normalized_tool_name,
                consecutive_failures: failures,
                limit: self.max_consecutive_failures_per_tool,
            });
        }

        Ok(())
    }

    pub(super) fn observe_outcome(
        &mut self,
        tool_call: &ToolCall,
        outcome: &Result<ToolResult, String>,
    ) {
        self.executed_calls = self.executed_calls.saturating_add(1);

        let normalized_tool_name = normalize_tool_for_policy(&tool_call.function.name);

        if normalized_tool_name.eq_ignore_ascii_case(RESET_POLICY_TOOL_NAME) {
            self.reset();
            return;
        }

        let succeeded = matches!(outcome, Ok(result) if result.success);

        if succeeded {
            self.consecutive_failures.remove(&normalized_tool_name);
            return;
        }

        *self
            .consecutive_failures
            .entry(normalized_tool_name)
            .or_insert(0) += 1;
    }

    fn reset(&mut self) {
        self.executed_calls = 0;
        self.consecutive_failures.clear();
    }
}

impl Default for ToolPolicyGuard {
    fn default() -> Self {
        Self::new(MAX_TOOL_CALLS_PER_ROUND, MAX_CONSECUTIVE_FAILURES_PER_TOOL)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use bamboo_agent_core::tools::FunctionCall;
    use bamboo_agent_core::{Message, Session};

    fn tool_call(name: &str, arguments: &str) -> ToolCall {
        tool_call_with_id("call-1", name, arguments)
    }

    fn tool_call_with_id(id: &str, name: &str, arguments: &str) -> ToolCall {
        ToolCall {
            id: id.to_string(),
            tool_type: "function".to_string(),
            function: FunctionCall {
                name: name.to_string(),
                arguments: arguments.to_string(),
            },
        }
    }

    #[test]
    fn strict_tools_reject_invalid_json_arguments() {
        let invalid = tool_call("Write", "{invalid");
        let err = validate_tool_call_arguments(&invalid).expect_err("expected strict rejection");
        assert!(err.contains("Tool policy blocked 'Write'"));
    }

    #[test]
    fn non_strict_tools_allow_invalid_json_for_best_effort_path() {
        let call = tool_call("Read", "{invalid");
        assert!(validate_tool_call_arguments(&call).is_ok());
    }

    #[test]
    fn conclusion_with_options_context_validation_requires_matching_assistant_narration() {
        let call = tool_call("conclusion_with_options", "{}");
        let session = Session::new("session-1", "model");
        let err =
            validate_tool_call_context(&call, &session).expect_err("expected context rejection");
        assert!(err.contains("before calling conclusion_with_options"));
        assert!(err.contains("Then retry conclusion_with_options"));
        assert!(err.contains("short plain-text wrap-up sentence"));
    }

    #[test]
    fn conclusion_with_options_context_validation_rejects_empty_assistant_narration() {
        let call = tool_call("conclusion_with_options", "{}");
        let mut session = Session::new("session-1", "model");
        session.add_message(Message::assistant("   ", Some(vec![call.clone()])));

        let err =
            validate_tool_call_context(&call, &session).expect_err("expected context rejection");
        assert!(err.contains("before calling conclusion_with_options"));
        assert!(err.contains("Then retry conclusion_with_options"));
    }

    #[test]
    fn conclusion_with_options_context_validation_accepts_non_empty_assistant_narration() {
        let call = tool_call("conclusion_with_options", "{}");
        let mut session = Session::new("session-1", "model");
        session.add_message(Message::assistant(
            "Summary: completed requested changes and ready for your confirmation.",
            Some(vec![call.clone()]),
        ));

        assert!(validate_tool_call_context(&call, &session).is_ok());
    }

    #[test]
    fn non_conclusion_with_options_context_validation_is_noop() {
        let call = tool_call("Read", "{}");
        let session = Session::new("session-1", "model");
        assert!(validate_tool_call_context(&call, &session).is_ok());
    }

    #[test]
    fn conclusion_with_options_requires_conclusion_fields_when_enhancement_enabled() {
        let conclusion_with_options_call = tool_call_with_id(
            "call-conclusion-with-options",
            "conclusion_with_options",
            r#"{"question":"Continue?"}"#,
        );
        let mut session = Session::new("session-1", "model");
        session.metadata.insert(
            COPILOT_CONCLUSION_WITH_OPTIONS_ENHANCEMENT_METADATA_KEY.to_string(),
            "true".to_string(),
        );
        session.add_message(Message::assistant(
            "Final confirmation request.",
            Some(vec![conclusion_with_options_call.clone()]),
        ));

        let err = validate_tool_call_context(&conclusion_with_options_call, &session)
            .expect_err("expected policy rejection");
        assert!(err.contains("conclusion.summary"));
        assert!(err.contains("conclusion.mermaid.graph"));
        assert!(err.contains("then retry conclusion_with_options"));
    }

    #[test]
    fn conclusion_with_options_with_conclusion_fields_passes_when_enhancement_enabled() {
        let conclusion_with_options_call = tool_call_with_id(
            "call-conclusion-with-options",
            "conclusion_with_options",
            r#"{"question":"Continue?","conclusion":{"summary":"All done.","mermaid":{"graph":"graph TD\nA-->B"}}}"#,
        );
        let mut session = Session::new("session-1", "model");
        session.metadata.insert(
            COPILOT_CONCLUSION_WITH_OPTIONS_ENHANCEMENT_METADATA_KEY.to_string(),
            "true".to_string(),
        );
        session.add_message(Message::assistant(
            "Final summary and confirmation.",
            Some(vec![conclusion_with_options_call.clone()]),
        ));

        assert!(validate_tool_call_context(&conclusion_with_options_call, &session).is_ok());
    }

    #[test]
    fn session_note_accepts_repairable_json_arguments() {
        let call = tool_call("session_note", r#"{"action":"read""#);
        assert!(validate_tool_call_arguments(&call).is_ok());
    }

    #[test]
    fn session_note_rejects_unrecoverable_json_with_retry_guidance() {
        let call = tool_call("session_note", "not-json");
        let err = validate_tool_call_arguments(&call).expect_err("expected strict rejection");
        assert!(err.contains("Rewrite the session_note call with a valid JSON object"));
    }

    #[test]
    fn precheck_blocks_when_round_limit_is_reached() {
        let mut guard = ToolPolicyGuard::default();
        let call = tool_call("Read", "{}");

        for _ in 0..MAX_TOOL_CALLS_PER_ROUND {
            guard.observe_outcome(
                &call,
                &Ok(ToolResult {
                    success: true,
                    result: "ok".to_string(),
                    display_preference: None,
                }),
            );
        }

        let violation = guard
            .check_before_execution(&call, 0)
            .expect_err("expected round limit violation");
        assert!(matches!(
            violation,
            ToolPolicyPrecheckViolation::RoundToolLimit { .. }
        ));
    }

    #[test]
    fn precheck_blocks_after_consecutive_failures() {
        let mut guard = ToolPolicyGuard::default();
        let call = tool_call("Bash", "{}");

        for _ in 0..MAX_CONSECUTIVE_FAILURES_PER_TOOL {
            guard.observe_outcome(&call, &Err("boom".to_string()));
        }

        let violation = guard
            .check_before_execution(&call, 0)
            .expect_err("expected circuit violation");
        assert!(matches!(
            violation,
            ToolPolicyPrecheckViolation::ToolCircuitOpen { .. }
        ));
    }

    #[test]
    fn successful_outcome_resets_failure_streak() {
        let mut guard = ToolPolicyGuard::default();
        let call = tool_call("Task", "{}");

        guard.observe_outcome(&call, &Err("boom".to_string()));
        guard.observe_outcome(
            &call,
            &Ok(ToolResult {
                success: true,
                result: "ok".to_string(),
                display_preference: None,
            }),
        );

        assert!(guard.check_before_execution(&call, 0).is_ok());
    }

    #[test]
    fn conclusion_with_options_resets_round_limit_counters() {
        let mut guard = ToolPolicyGuard::default();
        let read_call = tool_call("Read", "{}");
        let conclusion_with_options_call = tool_call("conclusion_with_options", "{}");

        for _ in 0..MAX_TOOL_CALLS_PER_ROUND {
            guard.observe_outcome(
                &read_call,
                &Ok(ToolResult {
                    success: true,
                    result: "ok".to_string(),
                    display_preference: None,
                }),
            );
        }

        assert!(guard.check_before_execution(&read_call, 0).is_err());

        guard.observe_outcome(
            &conclusion_with_options_call,
            &Ok(ToolResult {
                success: true,
                result: "ask".to_string(),
                display_preference: None,
            }),
        );

        assert!(guard.check_before_execution(&read_call, 0).is_ok());
    }

    #[test]
    fn conclusion_with_options_resets_failure_circuit() {
        let mut guard = ToolPolicyGuard::default();
        let bash_call = tool_call("Bash", "{}");
        let conclusion_with_options_call = tool_call("conclusion_with_options", "{}");

        for _ in 0..MAX_CONSECUTIVE_FAILURES_PER_TOOL {
            guard.observe_outcome(&bash_call, &Err("boom".to_string()));
        }
        assert!(guard.check_before_execution(&bash_call, 0).is_err());

        guard.observe_outcome(
            &conclusion_with_options_call,
            &Ok(ToolResult {
                success: true,
                result: "ask".to_string(),
                display_preference: None,
            }),
        );

        assert!(guard.check_before_execution(&bash_call, 0).is_ok());
    }
}
