import gradio as gr
from huggingface_hub import InferenceClient

# Qwen2.5-Coder is currently the best open-source coding model
try:
    client = InferenceClient("Qwen/Qwen2.5-Coder-32B-Instruct")
except:
    # Fallback if that specific model isn't on free tier right now
    client = InferenceClient("meta-llama/Llama-3.2-3B-Instruct")

def analyze_code(code, language, progress=gr.Progress()):
    progress(0.2, desc="Initializing Security Audit...")
    system_prompt = f"""You are a Senior Application Security Engineer and Expert Code Reviewer.
Analyze the following {language} code. 
1. Identify any security vulnerabilities (OWASP Top 10, Injection, etc.).
2. Point out performance bottlenecks or bad engineering practices.
3. Provide a secure, refactored version of the code.

Structure your response in Markdown with clear headings for 'Vulnerabilities', 'Best Practices', and 'Refactored Secure Code'.
"""
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": f"```{language}\n{code}\n```"}
    ]
    try:
        progress(0.4, desc="Analyzing codebase and generating report (This takes a few seconds)...")
        response = client.chat_completion(messages, max_tokens=1500)
        progress(1.0, desc="Audit Complete!")
        return response.choices[0].message.content
    except Exception as e:
        return f"⚠️ **Error connecting to Analysis Engine**: {str(e)}"

# A sleek Gradio interface
with gr.Blocks(theme=gr.themes.Base()) as demo:
    gr.Markdown("# 🔍 AI Smart Code Auditor")
    gr.Markdown("Secure your application. Paste your code and have an AI Security Engineer audit it for zero-days, vulnerabilities, and bad practices.")
    
    with gr.Row():
        with gr.Column(scale=1):
            lang = gr.Dropdown(choices=["Python", "JavaScript/TypeScript", "C/C++", "Java", "Go", "Rust", "PHP"], value="Python", label="Programming Language")
            code_input = gr.Code(label="Source Code", language="python", lines=15)
            btn = gr.Button("Analyze Code 🚀", variant="primary")
            
            example_code = '''import sqlite3
from flask import Flask, request

app = Flask(__name__)

@app.route('/user')
def get_user():
    username = request.args.get('username')
    conn = sqlite3.connect('users.db')
    cursor = conn.cursor()
    # Vulnerable to SQL Injection
    cursor.execute(f"SELECT * FROM users WHERE username = '{username}'")
    user = cursor.fetchone()
    return str(user)
'''
            gr.Markdown("### Try an example:")
            gr.Examples(examples=[[example_code, "Python"]], inputs=[code_input, lang])
            
        with gr.Column(scale=1):
            output = gr.Markdown(label="Audit Report")
            
    btn.click(analyze_code, inputs=[code_input, lang], outputs=output)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0")