import streamlit as st
import plotly.graph_objects as go
from combined_detector import CombinedCodeDetector
from rule_detector import RuleBasedCodeDetector
from fix_generator import FixSuggestionGenerator
import pandas as pd
import time
import requests
import json
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

# 1. Page Config
st.set_page_config(
    page_title="AI Code Security Scanner",
    page_icon="🔒",
    layout="wide"
)

# 2. Optimized Component Loading for Deployment
@st.cache_resource
def load_tools():
    # Hugging Face Model ID
    model_path = "mubi-613/ai-code-security-scanner"
    
    with st.spinner("🚀 Loading AI Models from Hugging Face... Please wait."):
        with torch.inference_mode():
            # Load Tokenizer and Model from HF Hub
            tokenizer = AutoTokenizer.from_pretrained(model_path)
            model = AutoModelForSequenceClassification.from_pretrained(model_path)
            
            # Initialize detectors (these internal classes should use the loaded model)
            detector = CombinedCodeDetector()
            fix_gen = FixSuggestionGenerator()
            rules = RuleBasedCodeDetector()
            
            # Inject the HF model into the detector if it expects a local one
            if hasattr(detector, 'model'):
                detector.model = model
                detector.model.eval()
            if hasattr(detector, 'tokenizer'):
                detector.tokenizer = tokenizer
                
            if hasattr(fix_gen, 'model'):
                fix_gen.model.eval()
                
            return {
                "detector": detector,
                "fix_gen": fix_gen,
                "rules": rules,
                "model": model,
                "tokenizer": tokenizer
            }

tools = load_tools()

# --- 3. CUSTOM STYLING ---
st.markdown("""
<style>
    .main-header {
        text-align: center;
        padding: 1.5rem;
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        color: white;
        border-radius: 10px;
        margin-bottom: 2rem;
    }
</style>
""", unsafe_allow_html=True)

# --- 4. HEADER ---
st.markdown('<div class="main-header"><h1>🔒 AI Code Security Scanner</h1><p>Real-time vulnerability detection and AI-powered fixes</p></div>', unsafe_allow_html=True)

# 5. Sidebar
with st.sidebar:
    st.header("⚙️ Configuration")
    analysis_mode = st.radio(
        "Analysis Mode",
        ["Fast (Rules Only)", "Deep (Rules + AI)", "API Mode"],
        help="Deep Mode uses the Hugging Face AI Model"
    )
    show_fixes = st.checkbox("Show Fix Suggestions", value=True)
    api_url = st.text_input("API URL", "http://localhost:8000")
    
    st.markdown("---")
    st.info("""
    **Features:**
    - Real-time code analysis
    - AI-powered fix suggestions
    - Batch processing
    - REST API integration
    """)
    

# 4. Main Tabs
tab1, tab2, tab3, tab4 = st.tabs(["🔍 Analyzer", "📁 Batch", "🌐 API", "📊 Dashboard"])

# --- TAB 1: ANALYZER ---
with tab1:
    col1, col2 = st.columns([2, 1])
    
    with col1:
        st.subheader("Code Analysis")
        default_code = """def get_user(user_id):
    # VULNERABLE: SQL Injection
    query = f"SELECT * FROM users WHERE id = {user_id}"
    api_key = "sk_live_1234567890abcdef"
    return execute_query(query)"""
        
        if 'code_input' not in st.session_state:
            st.session_state.code_input = default_code

        code_input = st.text_area(
            "Paste Python code:",
            value=st.session_state.code_input,
            height=300,
            key="main_editor",
            label_visibility="collapsed"
        )
    
    with col2:
        st.subheader("Templates")
        templates = {
            "SQL Injection": 'query = f"SELECT * FROM users WHERE id = {user_input}"',
            "Hardcoded Secret": 'api_key = "sk_test_1234567890"',
            "XSS Vulnerability": 'return f"<div>{user_input}</div>"',
            "Safe Example": 'query = "SELECT * FROM users WHERE id = %s"\ncursor.execute(query, (user_id,))',
        }
        
        for name, template in templates.items():
            if st.button(f"📋 {name}", use_container_width=True):
                st.session_state.code_input = template
                st.rerun()
    
    # Action Buttons
    btn_col1, btn_col2 = st.columns(2)
    with btn_col1:
        analyze_clicked = st.button("🔍 Analyze Code", type="primary", use_container_width=True)
    with btn_col2:
        if st.button("🗑️ Clear Analysis", use_container_width=True):
            if 'result' in st.session_state:
                del st.session_state.result
            st.rerun()

    if analyze_clicked:
        if not code_input.strip():
            st.warning("Please enter some code!")
        else:
            with st.spinner("Analyzing with AI..."):
                start_time = time.time()
                try:
                    with torch.no_grad():
                        if analysis_mode == "API Mode":
                            response = requests.post(f"{api_url}/analyze", json={"code": code_input, "detailed": True}, timeout=10)
                            result = response.json()
                        elif analysis_mode == "Fast (Rules Only)":
                            result = tools["rules"].analyze(code_input)
                        else:
                            result = tools["detector"].combined_analysis(code_input)
                        
                        result["analysis_time"] = time.time() - start_time
                        st.session_state.result = result
                except Exception as e:
                    st.error(f"Analysis failed: {str(e)}")

    # Display results
    if 'result' in st.session_state:
        result = st.session_state.result
        st.markdown("---")
        
        score = result.get("security_score", 0)
        # --- BALLOON LOGIC ---
        if score == 100:
            st.balloons()
            st.markdown("### **:green[✓ The code is safe! You have a 100 security score.]**")
        elif score >= 80:
            st.markdown("**:orange[⚠ Code is mostly safe, but minor issues were found.]**")
        else:
            st.markdown("**:red[✖ Critical vulnerabilities detected. Action required.]**")
        
        col_m1, col_m2, col_m3, col_m4 = st.columns(4)
        summary = result.get("summary", {})
        critical_count = summary.get("critical", 0) if isinstance(summary, dict) else 0
        
        with col_m1:
            fig = go.Figure(go.Indicator(
                mode="gauge+number",
                value=score,
                title={"text": "Security Score"},
                gauge={
                    'axis': {'range': [0, 100]}, 
                    'bar': {'color': "darkblue"},
                    'steps': [
                        {'range': [0, 50], 'color': "red"}, 
                        {'range': [50, 80], 'color': "orange"}, 
                        {'range': [80, 100], 'color': "green"}
                    ]
                }
            ))
            fig.update_layout(height=200, margin=dict(t=50, b=0, l=10, r=10))
            st.plotly_chart(fig, use_container_width=True)

        col_m2.metric("Issues", result.get("issue_count", 0))
        col_m3.metric("Critical", critical_count)
        col_m4.metric("Analysis Time", f"{result.get('analysis_time', 0):.2f}s")
        
        # --- AI FIX SUGGESTIONS ---
        issues = result.get("issues", [])
        if issues:
            st.subheader("🚨 Detected Issues")
            for issue in issues:
                with st.expander(f"[{issue.get('severity', 'UNKNOWN')}] {issue.get('type', 'issue').replace('_', ' ').title()} - Line {issue.get('line', '??')}"):
                    st.write(f"**Message:** {issue.get('message')}")
                    if show_fixes:
                        st.write("**💡 Fix Suggestions:**")
                        st.write("**💡 AI Recommended Fixes:**")
                        fixes = tools["fix_gen"].get_fixes(issue.get('message', ''), issue.get('type', ''))
                        # Numbered Display
                        for idx, fix in enumerate(fixes, 1):
                            st.markdown(f"**Suggestion {idx}:**")
                            st.code(fix, language="python")

# --- TAB 2: BATCH ---
with tab2:
    st.subheader("Batch Analysis")
    uploaded_files = st.file_uploader("Upload multiple .py files", type=['py'], accept_multiple_files=True)
    
    batch_btn1, batch_btn2 = st.columns(2)
    with batch_btn1:
        run_batch = st.button("🔍 Analyze All Files", type="primary", use_container_width=True)
    with batch_btn2:
        if st.button("🗑️ Clear Batch Results", use_container_width=True):
            if 'batch_df' in st.session_state:
                del st.session_state.batch_df
            st.rerun()

    if uploaded_files and run_batch:
        with st.spinner(f"Analyzing {len(uploaded_files)} files..."):
            results = []
            for file in uploaded_files:
                code = file.getvalue().decode("utf-8")
                res = tools["detector"].combined_analysis(code)
                results.append({
                    "file": file.name,
                    "score": res.get("security_score", 0),
                    "issues": res.get("issue_count", 0),
                    "critical": res.get("summary", {}).get("critical", 0)
                })
            st.session_state.batch_df = pd.DataFrame(results)

    if 'batch_df' in st.session_state:
        df = st.session_state.batch_df
        st.dataframe(df, use_container_width=True)
        fig_b = go.Figure(data=[go.Bar(name='Score', x=df['file'], y=df['score']), go.Bar(name='Issues', x=df['file'], y=df['issues'])])
        st.plotly_chart(fig_b, use_container_width=True)

# --- TAB 3: API ---
with tab3:
    st.subheader("API Integration")
    col_api1, col_api2 = st.columns(2)
    with col_api1:
        if st.button("Test Health Check"):
            try:
                response = requests.get(f"{api_url}/health")
                st.success(f"✅ API Healthy: {response.json()['status']}")
            except: st.error("❌ Connection Failed")
    with col_api2:
        st.markdown(f"- [OpenAPI Docs]({api_url}/docs)\n- [ReDoc]({api_url}/redoc)")

# --- TAB 4: DASHBOARD (Updated per your request) ---
with tab4:
    st.subheader("Project Dashboard")
    
    col_d1, col_d2, col_d3 = st.columns(3)
    
    with col_d1:
        st.metric("Vulnerability Types", "10+")
        st.metric("Detection Accuracy", "92%")
    
    with col_d2:
        st.metric("Supported Languages", "Python")
        st.metric("Response Time", "< 2s")
    
    with col_d3:
        st.metric("Integration Points", "4")
        st.metric("Lines Analyzed", "1000+")
    
    st.markdown("---")
    
    # 2. Centered Architecture Diagram
    st.write("**System Architecture:**")
    
    # Create 3 columns to center the image in the middle one
    # The [1, 2, 1] ratio makes the middle column 50% of the page width
    buf1, main_col, buf2 = st.columns([1, 2, 1])
    
    with main_col:
        import os
        # Use the filename of the image you saved
        img_path = os.path.join(os.getcwd(), "architecture.jpg") 
        
        if os.path.exists(img_path):
            
            st.image(img_path, width=600, caption="Project Workflow Diagram")
        else:
            st.info("💡 Please save your diagram as 'architecture.jpg' in the project folder.")

st.markdown("---")
st.caption("AI Code Security Scanner")