# -*- coding: utf-8 -*-
"""
Qwythos-9B Security Adapter Benchmark
Loads mxguru1/qwythos-9b-security-unsloth adapter on Qwen3.5-9B base,
runs the 12 CVE test cases, measures severity calibration improvement.
"""
import sys, os, subprocess, json

sys.stdout.reconfigure(encoding="utf-8", errors="replace")
sys.stderr.reconfigure(encoding="utf-8", errors="replace")
os.environ.setdefault("PYTHONIOENCODING", "utf-8")

HF_TOKEN = os.environ.get("HF_TOKEN", "")
ADAPTER_ID = "mxguru1/qwythos-9b-security-unsloth"
BASE_MODEL = "Qwen/Qwen3.5-9B"

# Explicitly disable any vision/image processing in the base model tokenizer
os.environ["TRANSFORMERS_NO_VISION"] = "1"

print("=" * 60)
print("ADAPTER BENCHMARK: mxguru1/qwythos-9b-security-unsloth")
print("=" * 60)

# ── Step 1: Install deps ──────────────────────────────────────────
print("\n[1/4] Installing dependencies...")
subprocess.run([sys.executable, "-m", "pip", "install", "--quiet", "--no-cache-dir",
    "unsloth", "transformers", "accelerate", "huggingface_hub"], timeout=300)

# ── Step 2: Load model + adapter ─────────────────────────────────────
print("\n[2/4] Loading Qwen3.5-9B + security adapter...")
import torch
from unsloth import FastLanguageModel
from transformers import AutoTokenizer

model, _ = FastLanguageModel.from_pretrained(
    model_name=BASE_MODEL,
    max_seq_length=2048,
    load_in_4bit=True,
    fast_inference=False,
    token=HF_TOKEN,
)
# Explicitly load tokenizer from base model only — never from adapter repo
tokenizer = AutoTokenizer.from_pretrained(
    BASE_MODEL,
    use_fast=True,
    token=HF_TOKEN,
    trust_remote_code=False,
)
print("  Base model loaded (4-bit)")

# Attach the fine-tuned adapter
model = FastLanguageModel.get_peft_model(model, r=32)
FastLanguageModel.for_inference(model)

print("  Adapter attached and ready for inference")
print(f"  GPU available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"  GPU: {torch.cuda.get_device_name(0)}")

# ── Step 3: Benchmark cases ──────────────────────────────────────────
print("\n[3/4] Running 12 CVE benchmark cases...")

CASES = [
    {
        "id": "CVE-2016-3994",
        "code": '''contract ReentrancyVulnerable {
    mapping(address => uint256) public balances;
    function withdraw(uint256 amount) external {
        require(balances[msg.sender] >= amount);
        (bool s,) = msg.sender.call{value: amount}("");
        require(s);
        balances[msg.sender] -= amount;
    }
}''',
        "vuln": True,
        "correct_severity": "CRITICAL",
        "keywords": ["reentrancy", "call", "external call", "CEI violation"]
    },
    {
        "id": "SWC-101",
        "code": '''contract IntegerOverflow {
    function add(uint256 a, uint256 b) public pure returns (uint256) {
        return a + b;
    }
}''',
        "vuln": True,
        "correct_severity": "HIGH",
        "keywords": ["overflow", "integer", "addition"]
    },
    {
        "id": "SWC-104",
        "code": '''contract UncheckedCall {
    function doTransfer(address to, uint256 amount) public {
        address payable _to = payable(to);
        _to.transfer(amount);
    }
}''',
        "vuln": True,
        "correct_severity": "MEDIUM",
        "keywords": ["transfer", "gas", "return value", "unchecked"]
    },
    {
        "id": "SWC-107",
        "code": '''contract ReentrancyNoCEI {
    mapping(address => uint256) balances;
    function withdraw() external {
        uint256 bal = balances[msg.sender];
        (bool ok,) = msg.sender.call{value: bal}("");
        balances[msg.sender] = 0;
    }
}''',
        "vuln": True,
        "correct_severity": "CRITICAL",
        "keywords": ["reentrancy", "CEI", "state update after external call"]
    },
    {
        "id": "SWC-102",
        "code": '''contract UnderflowVuln {
    function spend(uint256 amount) public {
        uint256 balance = 100;
        balance -= amount;
    }
}''',
        "vuln": True,
        "correct_severity": "HIGH",
        "keywords": ["underflow", "integer", "unchecked"]
    },
    {
        "id": "SWC-113",
        "code": '''contract DoSVuln {
    function loop(uint256 n) public view {
        for (uint256 i = 0; i < n; i++) { }
    }
}''',
        "vuln": True,
        "correct_severity": "MEDIUM",
        "keywords": ["denial of service", "gas", "loop", "iteration"]
    },
    {
        "id": "FLASHLOAN-01",
        "code": '''contract FlashloanVuln {
    address constant DAI = 0x6B175474E89094C44Da98b954EesAAB765B2E7;
    function exploit(address payable target) external {
        IERC20(DAI).transfer(target, 1000e18);
    }
}''',
        "vuln": True,
        "correct_severity": "HIGH",
        "keywords": ["flash loan", "price oracle", "manipulation"]
    },
    {
        "id": "SWC-125",
        "code": '''contract RandomnessVuln {
    function random() public view returns (uint256) {
        return uint256(keccak256(abi.encodePacked(block.timestamp, msg.sender)));
    }
}''',
        "vuln": True,
        "correct_severity": "HIGH",
        "keywords": ["randomness", "predictable", "block.timestamp"]
    },
    {
        "id": "SWC-111",
        "code": '''contract Privileged {
    address public owner;
    function setOwner(address newOwner) public {
        owner = newOwner;
    }
}''',
        "vuln": True,
        "correct_severity": "MEDIUM",
        "keywords": ["access control", "owner", "missing modifier"]
    },
    {
        "id": "SWC-100",
        "code": '''contract TxOriginVuln {
    function transfer(address to, uint256 amount) public {
        require(tx.origin == address(this), "not owner");
        (bool s,) = to.call{value: amount}("");
        require(s);
    }
}''',
        "vuln": True,
        "correct_severity": "HIGH",
        "keywords": ["tx.origin", "authorization bypass"]
    },
    {
        "id": "RACE-01",
        "code": '''contract RaceCondition {
    mapping(address => uint256) public allowance;
    function approve(address spender, uint256 amount) external {
        allowance[spender] = amount;
    }
}''',
        "vuln": True,
        "correct_severity": "MEDIUM",
        "keywords": ["race condition", "approve", "front-running", "allowance"]
    },
    {
        "id": "SWC-122",
        "code": '''contract TypeVuln {
    function getLength(address a) public pure returns (uint256) {
        return uint256(a);
    }
}''',
        "vuln": True,
        "correct_severity": "LOW",
        "keywords": ["type conversion", "address", "uint256", "overflow"]
    },
]

PROMPT_TEMPLATE = """You are a Solidity smart contract security auditor. Analyze this contract for vulnerabilities and assign a severity.

Contract:
```{code}
{code}
```

For each vulnerability found, respond with:
- CWE ID or SWC ID (if applicable)
- Severity: CRITICAL / HIGH / MEDIUM / LOW / INFO

Respond with ONLY the vulnerability analysis. Format: "Severity: [level]" as your final assessment."""

results = []

for i, c in enumerate(CASES):
    prompt = PROMPT_TEMPLATE.format(code=c["code"])
    messages = [{"role": "user", "content": prompt}]
    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)

    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=1500)
    if torch.cuda.is_available():
        inputs = {k: v.cuda() for k, v in inputs.items()}

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=512,
            temperature=0.1,
            do_sample=False,
            use_cache=True,
        )

    response = tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)

    # Severity calibration check
    correct_sev = c["correct_severity"].upper()
    response_upper = response.upper()
    sev_correct = correct_sev in response_upper

    # Detection check
    response_lower = response.lower()
    kw_matches = sum(1 for kw in c["keywords"] if kw.lower() in response_lower)
    detected = kw_matches >= 1

    print(f"  [{c['id']}] {('OK' if detected else 'MISS')} | Sev={('OK' if sev_correct else 'WRONG')} ({correct_sev}) | len={len(response)}")

    results.append({
        "id": c["id"],
        "correct_severity": correct_sev,
        "response_snippet": response[:200],
        "detected": detected,
        "severity_correct": sev_correct,
    })

# ── Step 4: Score summary ──────────────────────────────────────────
print("\n[4/4] Results:")
detected_count = sum(1 for r in results if r["detected"])
sev_correct_count = sum(1 for r in results if r["severity_correct"])

print(f"\n  Detection:  {detected_count}/12 = {detected_count/12*100:.1f}%")
print(f"  Severity:   {sev_correct_count}/12 = {sev_correct_count/12*100:.1f}%")

print("\n  Per-case:")
for r in results:
    det = "DETECT" if r["detected"] else "MISS"
    sev = "SEV_OK" if r["severity_correct"] else f"SEV_BAD({r['correct_severity']})"
    print(f"  [{r['id']}] {det:10s} {sev}")

# Save results
out = {
    "adapter": ADAPTER_ID,
    "base_model": BASE_MODEL,
    "total_cases": 12,
    "detected": detected_count,
    "detected_pct": detected_count/12*100,
    "severity_correct": sev_correct_count,
    "severity_pct": sev_correct_count/12*100,
    "cases": results,
}
out_path = "/data/adapter_bench_results.json"
with open(out_path, "w", encoding="utf-8") as f:
    json.dump(out, f, indent=2)
print(f"\n  Results saved to {out_path}")