import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import PeftModel

BASE_MODEL = "codellama/CodeLlama-7b-Instruct-hf"
ADAPTER_PATH = "/home/ai/codellama-security-finetuned"

def load_model():
    print("Loading base model...")
    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16,
        bnb_4bit_use_double_quant=True,
    )
    
    model = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL,
        quantization_config=bnb_config,
        device_map="auto",
        trust_remote_code=True,
        torch_dtype=torch.bfloat16,
    )
    
    print("Loading LoRA adapter...")
    model = PeftModel.from_pretrained(model, ADAPTER_PATH)
    model = model.merge_and_unload()
    
    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)
    tokenizer.pad_token = tokenizer.eos_token
    
    return model, tokenizer

def analyze_code(model, tokenizer, code, instruction="Identify all security vulnerabilities in this code"):
    prompt = f"""<|begin_of_text|><|start_header_id|>system<|end_header_id|>
You are a cybersecurity expert specializing in vulnerability assessment, penetration testing, and secure code review. Identify security flaws, explain the impact, and provide remediation.
<|eot_id|><|start_header_id|>user<|end_header_id|>
{instruction}

```python
{code}
```
<|eot_id|><|start_header_id|>assistant<|end_header_id|>"""
    
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=1024,
            temperature=0.3,
            top_p=0.9,
            do_sample=True,
            repetition_penalty=1.1,
            pad_token_id=tokenizer.eos_token_id,
        )
    
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response.split("<|start_header_id|>assistant<|end_header_id|>")[-1].strip()

def main():
    model, tokenizer = load_model()
    
    test_cases = [
        ("SQL Injection", """
import sqlite3
def get_user(username):
    conn = sqlite3.connect('app.db')
    cursor = conn.cursor()
    query = f"SELECT * FROM users WHERE name = '{username}'"
    return cursor.execute(query).fetchall()
"""),
        ("Command Injection", """
import subprocess
def ping_host(host):
    result = subprocess.run(f"ping -c 4 {host}", shell=True, capture_output=True)
    return result.stdout
"""),
        ("Path Traversal", """
from flask import request, send_file
@app.route('/download')
def download():
    filename = request.args.get('file')
    return send_file(f'/var/www/uploads/{filename}')
"""),
        ("Insecure Deserialization", """
import pickle
def load_session(data):
    return pickle.loads(data)
"""),
        ("JWT Issues", """
import jwt
def verify_token(token):
    return jwt.decode(token, 'secret123', algorithms=['HS256'])
"""),
    ]
    
    for name, code in test_cases:
        print(f"\n{'='*60}")
        print(f"TEST: {name}")
        print(f"{'='*60}")
        result = analyze_code(model, tokenizer, code)
        print(result)

if __name__ == "__main__":
    main()