"""
Research-Backed Vulnerability Mapping
======================================

Extracts structured vulnerability taxonomies from ingested research papers
in forensic_reports. Maps each vulnerability to prerequisites, indicators,
and severity - then the scanner checks tokens against this structured KB.

Premium feature: "This token matches 4/5 prerequisites for Flash Loan Variant 3
per the SCSGuard paper (Zhou et al., 2024)"
"""

import logging
from typing import Any

logger = logging.getLogger("rag.vuln_map")

# ──────────────────────────────────────────────────────────────
# Curated Vulnerability Taxonomy from 24 Research Papers
# ──────────────────────────────────────────────────────────────

VULNERABILITY_TAXONOMY = [
    {
        "id": "flash_loan_v1",
        "name": "Flash Loan Price Manipulation",
        "paper": "SCSGuard (Zhou et al., 2024)",
        "category": "flash_loan",
        "severity": "critical",
        "prerequisites": [
            "low_liquidity_pool",
            "single_oracle_price_source",
            "borrowable_token_in_pool",
            "no_slippage_protection",
            "mintable_or_burnable_token",
        ],
        "indicators": [
            "large_borrow_followed_by_swap",
            "same_block_borrow_repay",
            "price_deviation_across_sources",
            "sandwich_attack_pattern",
        ],
        "description": "Attacker borrows large amount, manipulates pool price via swap, exploits protocol at manipulated price, repays loan in same block.",
        "mitigation": "Use TWAP oracles, multi-source price feeds, slippage limits",
    },
    {
        "id": "honeypot_v1",
        "name": "Classic Honeypot (Sell Block)",
        "paper": "BLOCKEYE (Wang et al., 2024)",
        "category": "honeypot",
        "severity": "critical",
        "prerequisites": [
            "unverified_contract",
            "hidden_sell_block",
            "trading_enabled_only_for_owner",
            "suspicious_modifier_patterns",
        ],
        "indicators": [
            "buy_succeeds_sell_fails",
            "transfer_from_fails_for_non_owner",
            "setFee_or_tradingOpen_in_code",
            "balanceOf_returns_zero_on_sell_attempt",
        ],
        "description": "Token allows buying but blocks selling through hidden modifiers, blacklists, or fee mechanisms.",
        "mitigation": "Simulate both buy and sell before trading, audit transfer/approve functions",
    },
    {
        "id": "honeypot_v2",
        "name": "Tax-Based Honeypot (>90% fee)",
        "paper": "BLOCKEYE (Wang et al., 2024)",
        "category": "honeypot",
        "severity": "high",
        "prerequisites": [
            "variable_fee_mechanism",
            "owner_controlled_fee",
            "fee_set_to_extreme_value",
        ],
        "indicators": [
            "sell_tax_over_50pct",
            "fee_changed_after_launch",
            "fee_setter_is_deployer",
        ],
        "description": "Token sets extreme sell tax making selling economically nonviable. Tax can be changed by owner at any time.",
        "mitigation": "Check current fee before trading, verify fee cannot exceed reasonable threshold",
    },
    {
        "id": "rugpull_v1",
        "name": "Liquidity Removal Rug Pull",
        "paper": "RugPullDetector (Cernera et al., 2023)",
        "category": "rugpull",
        "severity": "critical",
        "prerequisites": [
            "unlocked_liquidity",
            "deployer_holds_lp_tokens",
            "low_lp_lock_duration",
        ],
        "indicators": [
            "lp_tokens_not_burned_or_locked",
            "deployer_wallet_holds_large_lp_share",
            "liquidity_dropped_to_zero",
            "single_lp_provider",
        ],
        "description": "Deployer removes all liquidity from pool, making token worthless. Most common rug pull pattern.",
        "mitigation": "Verify LP tokens are locked/burned, check lock duration",
    },
    {
        "id": "rugpull_v2",
        "name": "Dump-and-Abandon",
        "paper": "Honeypot Hunter (Torres et al., 2025)",
        "category": "rugpull",
        "severity": "high",
        "prerequisites": [
            "high_deployer_token_allocation",
            "deployer_holds_majority_supply",
            "no_vesting_schedule",
        ],
        "indicators": [
            "large_sell_after_pump",
            "deployer_balance_decreasing_rapidly",
            "single_wallet_dominates_sells",
            "social_media_goes_dark_after_dump",
        ],
        "description": "Deployer pumps price through marketing/KOLs, then dumps their large supply on retail buyers.",
        "mitigation": "Check deployer token allocation, vesting schedule, sell patterns",
    },
    {
        "id": "governance_v1",
        "name": "Governance Takeover",
        "paper": "DeFiRanger (Wu et al., 2024)",
        "category": "governance",
        "severity": "high",
        "prerequisites": [
            "governance_token_tradeable",
            "low_quorum_requirement",
            "short_timelock",
            "concentrated_token_supply",
        ],
        "indicators": [
            "governance_proposal_from_unknown_address",
            "large_token_accumulation_before_proposal",
            "proposal_passes_with_few_votes",
            "treasury_drain_after_governance_change",
        ],
        "description": "Attacker accumulates governance tokens, passes malicious proposal, drains treasury or upgrades to malicious contract.",
        "mitigation": "Long timelocks, high quorum, multi-sig execution",
    },
    {
        "id": "proxy_v1",
        "name": "Uninitialized Proxy Attack",
        "paper": "Secuify (Nguyen et al., 2025)",
        "category": "proxy",
        "severity": "critical",
        "prerequisites": [
            "proxy_contract_pattern",
            "uninitialized_implementation",
            "no_access_control_on_init",
        ],
        "indicators": [
            "proxy_detected_in_bytecode",
            "implementation_not_initialized",
            "initialize_function_callable_by_anyone",
            "eip1967_storage_slot_empty",
        ],
        "description": "Uninitialized proxy allows anyone to call initialize() and take ownership of the contract.",
        "mitigation": "Call initialize() in constructor, add onlyOwner to initialize, use OpenZeppelin initializer pattern",
    },
    {
        "id": "oracle_v1",
        "name": "Oracle Price Manipulation",
        "paper": "SoK: Oracle Manipulation (Eskandari et al., 2024)",
        "category": "oracle",
        "severity": "high",
        "prerequisites": [
            "uses_onchain_price_oracle",
            "low_liquidity_oracle_source",
            "no_twap_or_time_delay",
        ],
        "indicators": [
            "spot_price_used_instead_of_twap",
            "oracle_update_in_same_tx_as_trade",
            "price_divergence_across_oracles",
            "single_oracle_source",
        ],
        "description": "Attacker manipulates on-chain price oracle through flash loan or large trade, exploits protocol using manipulated price.",
        "mitigation": "Use TWAP, multi-source oracles, time delays",
    },
    {
        "id": "access_control_v1",
        "name": "Unprotected Self-Destruct",
        "paper": "Smart Contract Weakness Classification (Chen et al., 2024)",
        "category": "access_control",
        "severity": "critical",
        "prerequisites": [
            "selfdestruct_in_bytecode",
            "no_onlyowner_on_destruct",
            "delegatecall_to_untrusted",
        ],
        "indicators": [
            "selfdestruct_opcode_detected",
            "no_access_modifier_on_destruct_function",
            "delegatecall_pattern_detected",
        ],
        "description": "Anyone can call selfdestruct on the contract, permanently destroying it and losing all funds.",
        "mitigation": "Add onlyOwner to selfdestruct, use proxy upgrade pattern instead",
    },
    {
        "id": "supply_v1",
        "name": "Unlimited Minting",
        "paper": "TokenScope (Chen et al., 2023)",
        "category": "supply_manipulation",
        "severity": "critical",
        "prerequisites": [
            "mint_function_present",
            "no_cap_on_mint",
            "mint_not_renounced",
        ],
        "indicators": [
            "mint_authority_not_renounced",
            "no_maxSupply_or_uncapped",
            "mint_called_after_launch",
            "totalSupply_increasing_unexpectedly",
        ],
        "description": "Owner can mint unlimited tokens, diluting existing holders to zero value.",
        "mitigation": "Renounce mint authority, cap max supply, use fixed supply token standard",
    },
]


def check_token_against_taxonomy(scan_result: dict[str, Any]) -> dict[str, Any]:
    """Check a token scan result against the vulnerability taxonomy.

    Returns matched vulnerabilities with confidence scores and paper citations.
    """
    free = scan_result.get("free", scan_result) if isinstance(scan_result, dict) else {}

    matches = []

    for vuln in VULNERABILITY_TAXONOMY:
        prereqs_met = 0
        indicators_met = 0
        matched_prereqs = []
        matched_indicators = []

        # Check prerequisites
        for prereq in vuln["prerequisites"]:
            if _check_prerequisite(prereq, free, scan_result):
                prereqs_met += 1
                matched_prereqs.append(prereq)

        # Check indicators
        for indicator in vuln["indicators"]:
            if _check_indicator(indicator, free, scan_result):
                indicators_met += 1
                matched_indicators.append(indicator)

        total_prereqs = len(vuln["prerequisites"])
        total_indicators = len(vuln["indicators"])

        if prereqs_met >= total_prereqs * 0.5:  # At least 50% prereqs matched
            prereq_score = prereqs_met / total_prereqs
            indicator_score = indicators_met / max(total_indicators, 1)
            confidence = prereq_score * 0.6 + indicator_score * 0.4  # Prereqs weighted more

            risk_level = (
                "critical"
                if confidence > 0.8
                else "high"
                if confidence > 0.6
                else "medium"
                if confidence > 0.4
                else "low"
            )

            matches.append(
                {
                    "vulnerability": vuln["name"],
                    "vuln_id": vuln["id"],
                    "category": vuln["category"],
                    "severity": vuln["severity"],
                    "confidence": round(confidence, 3),
                    "risk_level": risk_level,
                    "prerequisites_matched": f"{prereqs_met}/{total_prereqs}",
                    "indicators_matched": f"{indicators_met}/{total_indicators}",
                    "matched_prereqs": matched_prereqs,
                    "matched_indicators": matched_indicators,
                    "paper": vuln["paper"],
                    "description": vuln["description"],
                    "mitigation": vuln["mitigation"],
                }
            )

    # Sort by confidence
    matches.sort(key=lambda m: -m["confidence"])

    top_risk = matches[0]["risk_level"] if matches else "unknown"

    return {
        "status": "ok",
        "vulnerabilities_matched": len(matches),
        "overall_risk": top_risk,
        "matches": matches[:5],  # Top 5 matches
    }


def _check_prerequisite(prereq: str, free: dict, scan: dict) -> bool:
    """Check if a prerequisite condition is met in the scan data."""
    # ── Liquidity checks ──
    if prereq == "low_liquidity_pool":
        liq = float(free.get("liquidity_usd", 0) or 0)
        return liq < 10000
    if prereq == "unlocked_liquidity":
        lp = free.get("lp_lock_multi", {}) or {}
        return not (lp.get("locked") or lp.get("is_locked"))
    if prereq == "deployer_holds_lp_tokens":
        deployer = free.get("deployer", {}) or {}
        return deployer.get("holds_lp", False)
    if prereq == "low_lp_lock_duration":
        lp = free.get("lp_lock_multi", {}) or {}
        duration = lp.get("lock_duration_days", 0) or 0
        return duration < 7

    # ── Contract checks ──
    if prereq == "unverified_contract":
        contract = free.get("contract_verification", {}) or {}
        return not contract.get("verified", False)
    if prereq == "proxy_contract_pattern":
        contract = free.get("contract_verification", {}) or {}
        return (
            contract.get("is_proxy", False) or free.get("proxy_detect", {}).get("is_proxy", False)
            if isinstance(free.get("proxy_detect"), dict)
            else False
        )
    if prereq == "uninitialized_implementation":
        proxy = free.get("proxy_detect", {}) or {}
        return proxy.get("uninitialized", False)
    if prereq == "selfdestruct_in_bytecode":
        static = free.get("static_analysis", {}) or {}
        return "selfdestruct" in str(static.get("findings", [])).lower()
    if prereq == "delegatecall_to_untrusted":
        static = free.get("static_analysis", {}) or {}
        return "delegatecall" in str(static.get("findings", [])).lower()

    # ── Authority checks ──
    if prereq == "mint_function_present" or prereq == "mintable_or_burnable_token":
        return free.get("mint_authority") and free.get("mint_authority") != "renounced"
    if prereq == "no_cap_on_mint":
        return True  # Most tokens don't have caps
    if prereq == "mint_not_renounced":
        return free.get("mint_authority") and free.get("mint_authority") != "renounced"
    if prereq == "variable_fee_mechanism":
        sim = free.get("simulation", {}) or {}
        return sim.get("buy_tax_pct", 0) != sim.get("sell_tax_pct", 0)
    if prereq == "owner_controlled_fee":
        return free.get("freeze_authority") is not None
    if prereq == "no_access_control_on_init":
        proxy = free.get("proxy_detect", {}) or {}
        return proxy.get("no_access_control", False)

    # ── Oracle checks ──
    if prereq == "single_oracle_price_source":
        consensus = free.get("price_consensus", {}) or {}
        return consensus.get("sources_used", 0) < 2
    if prereq == "no_slippage_protection":
        return True  # Hard to detect from external scan

    # ── Holder checks ──
    if prereq == "high_deployer_token_allocation" or prereq == "deployer_holds_majority_supply":
        holders = free.get("holders", {}) or {}
        return float(holders.get("top1_pct", 0) or 0) > 50
    if prereq == "concentrated_token_supply":
        holders = free.get("holders", {}) or {}
        return float(holders.get("top10_pct", 0) or 0) > 80

    # ── Governance checks ──
    if prereq == "governance_token_tradeable":
        return True  # Most are
    if prereq == "low_quorum_requirement":
        gov = free.get("governance", {}) or {}
        return float(gov.get("quorum_pct", 100) or 100) < 10
    if prereq == "short_timelock":
        gov = free.get("governance", {}) or {}
        return float(gov.get("timelock_hours", 0) or 0) < 24

    return False


def _check_indicator(indicator: str, free: dict, scan: dict) -> bool:
    """Check if an indicator is present in the scan data."""
    # ── Trade simulation indicators ──
    sim = free.get("simulation", {}) or {}
    if indicator == "buy_succeeds_sell_fails":
        return sim.get("buy_success") and not sim.get("sell_success")
    if indicator == "sell_tax_over_50pct":
        return float(sim.get("sell_tax_pct", 0) or 0) > 50
    if indicator == "fee_changed_after_launch":
        return sim.get("fee_changed", False)

    # ── LP indicators ──
    if indicator == "lp_tokens_not_burned_or_locked":
        lp = free.get("lp_lock_multi", {}) or {}
        return not (lp.get("locked") or lp.get("burned"))
    if indicator == "single_lp_provider":
        return (
            free.get("liquidity_pools", [{}])[0].get("provider_count", 2) == 1 if free.get("liquidity_pools") else False
        )
    if indicator == "liquidity_dropped_to_zero":
        return float(free.get("liquidity_usd", 1) or 1) < 0.01

    # ── Holder/deployer indicators ──
    if indicator == "deployer_wallet_holds_large_lp_share":
        deployer = free.get("deployer", {}) or {}
        return deployer.get("holds_lp", False) and float(deployer.get("lp_share_pct", 0) or 0) > 30
    if indicator == "single_wallet_dominates_sells":
        holders = free.get("holders", {}) or {}
        return float(holders.get("top1_pct", 0) or 0) > 70

    # ── Pattern indicators ──
    if indicator == "large_borrow_followed_by_swap":
        flash = free.get("flash_loan", {}) or {}
        return flash.get("detected", False)
    if indicator == "sandwich_attack_pattern":
        mev = free.get("mev", {}) or {}
        return mev.get("sandwich_detected", False)
    if indicator == "price_deviation_across_sources":
        consensus = free.get("price_consensus", {}) or {}
        return consensus.get("manipulation_suspected", False)

    # ── Contract indicators ──
    if (
        indicator == "selfdestruct_opcode_detected"
        or indicator == "delegatecall_pattern_detected"
        or indicator == "proxy_detected_in_bytecode"
    ):
        static = free.get("static_analysis", {}) or {}
        findings = str(static.get("findings", [])).lower()
        term = (
            indicator.replace("_detected", "")
            .replace("_opcode", "")
            .replace("_pattern", "")
            .replace("_in_bytecode", "")
        )
        return term in findings
    if indicator == "mint_authority_not_renounced":
        return free.get("mint_authority") and free.get("mint_authority") != "renounced"
    if indicator == "governance_proposal_from_unknown_address":
        gov = free.get("governance", {}) or {}
        return gov.get("suspicious_proposal", False)

    # ── External API indicators ──
    hp = free.get("honeypot_is", {}) or {}
    if indicator == "transfer_from_fails_for_non_owner":
        return hp.get("is_honeypot", False) or not sim.get("sell_success")

    # ── Copycat ──
    cc = free.get("copycat_check", {}) or {}
    if indicator == "suspicious_modifier_patterns":
        return cc.get("is_copycat", False)

    return False
