pub use knowdit_kg_model::audit_finding::FindingSeverity;
pub use knowdit_kg_model::finding_category::VulnerabilityCategory;

#[derive(Debug, Clone, Copy)]
pub struct FindingTaxonomyEntry {
    pub category: VulnerabilityCategory,
    pub subcategory: &'static str,
    pub description: &'static str,
}

pub const FINDING_TAXONOMY: &[FindingTaxonomyEntry] = &[
    FindingTaxonomyEntry {
        category: VulnerabilityCategory::AccessControl,
        subcategory: "Arbitrary `from` in transferFrom() without msg.sender Check",
        description: "Detect when transferFrom uses an arbitrary from address without validating msg.sender.",
    },
    FindingTaxonomyEntry {
        category: VulnerabilityCategory::AccessControl,
        subcategory: "Call to Arbitrary Addresses with Unchecked Calldata",
        description: "Calling arbitrary addresses with attacker-controlled calldata can trigger unexpected behavior.",
    },
    FindingTaxonomyEntry {
        category: VulnerabilityCategory::AccessControl,
        subcategory: "Caller Not Checked",
        description: "Authorization assumes caller properties that can be bypassed, such as constructor-time extcodesize checks.",
    },
    FindingTaxonomyEntry {
        category: VulnerabilityCategory::AccessControl,
        subcategory: "Contract Could be Destructed",
        description: "A privileged or unguarded path can destruct the contract or logic contract.",
    },
    FindingTaxonomyEntry {
        category: VulnerabilityCategory::AccessControl,
        subcategory: "Dangerous Immediate Initialization of State Variables",
        description: "State is initialized through unsafe non-constant calls or mutable state dependencies.",
    },
    FindingTaxonomyEntry {
        category: VulnerabilityCategory::AccessControl,
        subcategory: "Dangerous Usage of `tx.origin`",
        description: "Authorization depends on tx.origin and can be abused through call relays.",
    },
    FindingTaxonomyEntry {
        category: VulnerabilityCategory::AccessControl,
        subcategory: "Default Function Visibility",
        description: "Missing explicit visibility exposes state-changing functions publicly.",
    },
    FindingTaxonomyEntry {
        category: VulnerabilityCategory::AccessControl,
        subcategory: "Initializing Method without Permission Check",
        description: "Initialization can be triggered without an adequate permission check.",
    },
    FindingTaxonomyEntry {
        category: VulnerabilityCategory::AccessControl,
        subcategory: "Method permit() Used for Arbitrary `from` in transferFrom()",
        description: "permit-based flows allow arbitrary transferFrom source selection without validating the true owner.",
    },
    FindingTaxonomyEntry {
        category: VulnerabilityCategory::AccessControl,
        subcategory: "Missing `msg.sender` Check for transferFrom()",
        description: "transferFrom uses an attacker-controlled from address instead of binding to msg.sender or equivalent authorization.",
    },
    FindingTaxonomyEntry {
        category: VulnerabilityCategory::AccessControl,
        subcategory: "Missing Input Validation",
        description: "Critical input values such as addresses, ids, chain selectors, or enum values are not validated.",
    },
    FindingTaxonomyEntry {
        category: VulnerabilityCategory::AccessControl,
        subcategory: "Sending Ether to Arbitrary Destinations",
        description: "Ether can be sent to attacker-controlled destinations due to insufficient validation.",
    },
    FindingTaxonomyEntry {
        category: VulnerabilityCategory::AccessControl,
        subcategory: "Unprotected Contract Destruction",
        description: "selfdestruct or equivalent destructive behavior is reachable without adequate protection.",
    },
    FindingTaxonomyEntry {
        category: VulnerabilityCategory::AccessControl,
        subcategory: "Unprotected Ether Withdrawal",
        description: "Ether withdrawal paths are exposed without adequate access control or accounting.",
    },
    FindingTaxonomyEntry {
        category: VulnerabilityCategory::AccessControl,
        subcategory: "Unsafe Delegatecall",
        description: "delegatecall is performed against unsafe or attacker-controlled targets.",
    },
    FindingTaxonomyEntry {
        category: VulnerabilityCategory::AccessControl,
        subcategory: "Unused Return Value",
        description: "A critical external call return value is ignored, hiding failure or authorization problems.",
    },
    FindingTaxonomyEntry {
        category: VulnerabilityCategory::AccessControl,
        subcategory: "Usage of public mint or burn",
        description: "Mint or burn functionality is exposed publicly or without sufficient authorization.",
    },
    FindingTaxonomyEntry {
        category: VulnerabilityCategory::AccessControl,
        subcategory: "Write to Arbitrary Storage Location",
        description: "An attacker can influence writes to arbitrary storage slots or equivalent protected state.",
    },
    FindingTaxonomyEntry {
        category: VulnerabilityCategory::Arithmetic,
        subcategory: "Inappropriate Integer Division before Multiplication",
        description: "Integer division before multiplication truncates precision and causes incorrect accounting.",
    },
    FindingTaxonomyEntry {
        category: VulnerabilityCategory::Arithmetic,
        subcategory: "Integer Overflow",
        description: "Arithmetic can overflow and wrap or revert in a way that breaks intended invariants.",
    },
    FindingTaxonomyEntry {
        category: VulnerabilityCategory::Arithmetic,
        subcategory: "Integer Underflow",
        description: "Arithmetic can underflow and wrap or revert in a way that breaks intended invariants.",
    },
    FindingTaxonomyEntry {
        category: VulnerabilityCategory::Arithmetic,
        subcategory: "Unsafe Array Length Assignment",
        description: "User-controlled inputs can influence array length assignments or equivalent bounds changes.",
    },
    FindingTaxonomyEntry {
        category: VulnerabilityCategory::BlockManipulation,
        subcategory: "Dangerous Usage of `block.timestamp`",
        description: "Security-critical logic depends on block.timestamp or similarly manipulable block metadata.",
    },
    FindingTaxonomyEntry {
        category: VulnerabilityCategory::BlockManipulation,
        subcategory: "Transaction Order Dependency",
        description: "Correctness or fairness depends on transaction ordering and can be manipulated by frontrunning or reordering.",
    },
    FindingTaxonomyEntry {
        category: VulnerabilityCategory::BlockManipulation,
        subcategory: "Weak PRNG (Pseudorandom Number Generator)",
        description: "Randomness depends on predictable or miner-influenced blockchain values.",
    },
    FindingTaxonomyEntry {
        category: VulnerabilityCategory::Cryptographic,
        subcategory: "Lack of Proper Signature Verification",
        description: "Signed messages are accepted without sufficient signer, domain, nonce, or intent validation.",
    },
    FindingTaxonomyEntry {
        category: VulnerabilityCategory::Cryptographic,
        subcategory: "Signature Malleability",
        description: "Signatures can be altered into alternative valid encodings and replayed or mis-accounted.",
    },
    FindingTaxonomyEntry {
        category: VulnerabilityCategory::DenialOfServices,
        subcategory: "`transfer()` and `send()` with Hardcoded Gas Amount",
        description: "Value transfer relies on fixed gas forwarding assumptions that can break protocol liveness.",
    },
    FindingTaxonomyEntry {
        category: VulnerabilityCategory::DenialOfServices,
        subcategory: "Contract Could Lock Ether",
        description: "Funds can become permanently locked because the contract has no viable withdrawal or escape path.",
    },
    FindingTaxonomyEntry {
        category: VulnerabilityCategory::DenialOfServices,
        subcategory: "DoS with Block Gas Limit",
        description: "Unbounded work or data growth can push execution beyond the block gas limit.",
    },
    FindingTaxonomyEntry {
        category: VulnerabilityCategory::DenialOfServices,
        subcategory: "DoS With Failed Call",
        description: "A failed external call can block progress or cause protocol-wide denial of service.",
    },
    FindingTaxonomyEntry {
        category: VulnerabilityCategory::DenialOfServices,
        subcategory: "Force Sending Ether with this.balance check in require() or assert()",
        description: "Force-sent ether can break logic that relies on exact balance checks.",
    },
    FindingTaxonomyEntry {
        category: VulnerabilityCategory::DenialOfServices,
        subcategory: "Unsafe send() in the require() Condition",
        description: "send inside a require/assert style condition can be weaponized to revert protocol progress.",
    },
    FindingTaxonomyEntry {
        category: VulnerabilityCategory::Reentrancy,
        subcategory: "Reentrancy Vulnerability with Negative Events",
        description: "External control flow can reenter before state is safely updated, leading to invalid negative-state effects.",
    },
    FindingTaxonomyEntry {
        category: VulnerabilityCategory::Reentrancy,
        subcategory: "Reentrancy Vulnerability with Transfer",
        description: "An external transfer opens a reentrant path before invariant-preserving state updates complete.",
    },
    FindingTaxonomyEntry {
        category: VulnerabilityCategory::Reentrancy,
        subcategory: "Reentrancy Vulnerability with Same Effect",
        description: "Reentrancy repeats the same state transition or accounting effect multiple times.",
    },
    FindingTaxonomyEntry {
        category: VulnerabilityCategory::Reentrancy,
        subcategory: "Reentrancy Vulnerability with ETH Transfer",
        description: "ETH transfer creates a callback path that can reenter before state is finalized.",
    },
    FindingTaxonomyEntry {
        category: VulnerabilityCategory::Reentrancy,
        subcategory: "Reentrancy Vulnerability without ETH Transfer",
        description: "An external call without direct ETH transfer still enables reentrant state corruption.",
    },
    FindingTaxonomyEntry {
        category: VulnerabilityCategory::StorageAndMemory,
        subcategory: "Arbitrary Function Jump via Inline Assembly",
        description: "Inline assembly can redirect execution or corrupt function pointers or equivalent dispatch state.",
    },
    FindingTaxonomyEntry {
        category: VulnerabilityCategory::StorageAndMemory,
        subcategory: "Bytes Variables Risk",
        description: "Improper ABI-decoding or bytes handling causes ambiguous values, truncation, or collision-like behavior.",
    },
    FindingTaxonomyEntry {
        category: VulnerabilityCategory::StorageAndMemory,
        subcategory: "Dangerous Usage of `msg.value` inside a Loop",
        description: "msg.value is reused inside loops and causes repeated accounting or payout errors.",
    },
    FindingTaxonomyEntry {
        category: VulnerabilityCategory::StorageAndMemory,
        subcategory: "Error-prone Assembly Usage",
        description: "Low-level assembly is used in a way that can corrupt state or bypass invariants.",
    },
    FindingTaxonomyEntry {
        category: VulnerabilityCategory::StorageAndMemory,
        subcategory: "Memory Manipulation",
        description: "Unsafe memory access or mutation causes incorrect behavior, corruption, or exploitable state transitions.",
    },
    FindingTaxonomyEntry {
        category: VulnerabilityCategory::StorageAndMemory,
        subcategory: "Modifying storage array by value",
        description: "Storage arrays are handled with the wrong copy/reference semantics and cause unintended mutation.",
    },
    FindingTaxonomyEntry {
        category: VulnerabilityCategory::StorageAndMemory,
        subcategory: "Payable Functions using `delegatecall` inside a Loop",
        description: "delegatecall in a payable loop can multiply accounting effects or corrupt balance tracking.",
    },
];

pub fn taxonomy_prompt() -> String {
    let mut out = String::from(
        "Choose exactly one vulnerability category and one subcategory from the taxonomy below.\n\n",
    );
    let mut current_category = None;

    for entry in FINDING_TAXONOMY {
        if current_category != Some(entry.category) {
            if current_category.is_some() {
                out.push('\n');
            }
            out.push_str(&format!("### {}\n", entry.category));
            current_category = Some(entry.category);
        }

        out.push_str(&format!("- {}: {}\n", entry.subcategory, entry.description));
    }

    out
}

pub fn normalize_taxonomy_key(text: &str) -> String {
    let mut out = String::new();
    let mut prev_sep = false;

    for ch in text.chars() {
        if ch.is_ascii_alphanumeric() {
            out.push(ch.to_ascii_lowercase());
            prev_sep = false;
        } else if !prev_sep {
            out.push(' ');
            prev_sep = true;
        }
    }

    out.split_whitespace().collect::<Vec<_>>().join(" ")
}

pub fn resolve_taxonomy_entry(
    category: VulnerabilityCategory,
    subcategory: &str,
) -> Option<&'static FindingTaxonomyEntry> {
    let normalized = normalize_taxonomy_key(subcategory);

    FINDING_TAXONOMY.iter().find(|entry| {
        entry.category == category && normalize_taxonomy_key(entry.subcategory) == normalized
    })
}

/// Validate a `(category, subcategory)` pair against the taxonomy and
/// return a self-healing error message for the LLM if it is invalid. The
/// message tells the agent the most likely fix:
///
/// 1. If the subcategory exists under a different category, point the
///    agent at the correct category.
/// 2. Otherwise list the valid subcategories under the agent's requested
///    category so it can pick one.
///
/// Returns `None` when the pair resolves cleanly.
pub fn validate_taxonomy_pair(
    category: VulnerabilityCategory,
    subcategory: &str,
) -> Option<String> {
    if resolve_taxonomy_entry(category, subcategory).is_some() {
        return None;
    }

    let normalized = normalize_taxonomy_key(subcategory);
    let owning_categories: Vec<VulnerabilityCategory> = FINDING_TAXONOMY
        .iter()
        .filter(|e| normalize_taxonomy_key(e.subcategory) == normalized)
        .map(|e| e.category)
        .collect();

    if !owning_categories.is_empty() {
        let owners = owning_categories
            .iter()
            .map(|c| format!("`{}`", c))
            .collect::<Vec<_>>()
            .join(", ");
        return Some(format!(
            "subcategory `{}` is not under category `{}`. It actually belongs to {}. Re-emit the finding with the correct `category` value.",
            subcategory, category, owners,
        ));
    }

    let siblings: Vec<&'static str> = FINDING_TAXONOMY
        .iter()
        .filter(|e| e.category == category)
        .map(|e| e.subcategory)
        .collect();
    let listed = if siblings.is_empty() {
        "(no subcategories defined)".to_string()
    } else {
        siblings
            .iter()
            .map(|s| format!("`{}`", s))
            .collect::<Vec<_>>()
            .join(", ")
    };
    Some(format!(
        "subcategory `{}` is not in the taxonomy under any category. Valid subcategories under `{}` are: {}. Re-emit the finding with one of these (or correct the `category` if the subcategory belongs elsewhere).",
        subcategory, category, listed,
    ))
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn normalizes_subcategory_names() {
        let resolved = resolve_taxonomy_entry(
            VulnerabilityCategory::DenialOfServices,
            "DoS With Block Gas Limit",
        )
        .unwrap();

        assert_eq!(resolved.subcategory, "DoS with Block Gas Limit");
    }
}
