import streamlit as st
from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification
import torch

# Load Pre-trained Model (Replace with your fine-tuned model)
@st.cache_resource
def load_model():
    model_name = "microsoft/codebert-base"  # Replace with a fine-tuned model for vulnerability detection
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForSequenceClassification.from_pretrained(model_name)
    return tokenizer, model

tokenizer, model = load_model()

# Vulnerability Explanation Function
def analyze_code(code_snippet):
    # Tokenize Input
    inputs = tokenizer(code_snippet, return_tensors="pt", truncation=True, max_length=512)
    outputs = model(**inputs)
    predictions = torch.softmax(outputs.logits, dim=1)
    vulnerability_score = predictions[0][1].item()  # Assuming index 1 corresponds to "vulnerable"
    
    # Generate Explanation
    if vulnerability_score > 0.6:
        explanation = (
            f"The code has a high likelihood of being vulnerable. The model detected patterns "
            f"indicative of potential security issues."
        )
    elif vulnerability_score > 0.3:
        explanation = (
            f"The code has some potential vulnerabilities. Review the logic carefully, especially in "
            f"sensitive operations like input validation or file handling."
        )
    else:
        explanation = (
            f"The code appears to be safe based on the analysis. However, manual review is always recommended."
        )
    return vulnerability_score, explanation

# Streamlit UI
st.title("AI-Enhanced Code Vulnerability Scanner")
st.markdown("""
This tool uses AI to detect vulnerabilities in Python code and provides explanations for potential issues.
""")

# Input Section
code_snippet = st.text_area("Paste your Python code here:", height=200)
analyze_button = st.button("Analyze Code")

if analyze_button and code_snippet.strip():
    with st.spinner("Analyzing your code..."):
        score, explanation = analyze_code(code_snippet)
    
    # Display Results
    st.subheader("Analysis Results")
    st.write(f"**Vulnerability Score:** {score:.2f}")
    st.write(f"**Explanation:** {explanation}")
else:
    st.info("Please paste Python code and click 'Analyze Code'.")