Update app.py

app.py

@@ -1,21 +1,36 @@
 import streamlit as st
 import os
-import base64
-from google import genai
 import zipfile
 import io
 import re
+from google import genai
 
 # Initialize Gemini Client with API key from environment
 client = genai.Client(api_key=os.environ.get("GEMINI_API_KEY"))
 
-# Extended supported file types
+# Supported extensions and corresponding language names for syntax highlighting
 SUPPORTED_EXTENSIONS = [
-    "py", "c", "cpp", "java", "js", "jsx",
-    "ts", "go", "rs", "php", "rb", "swift", "kt",
-    "html", "css"
+    "py", "c", "cpp", "java", "js", "jsx", "ts", "go", "rs", "php", "rb", "swift", "kt", "html", "css"
 ]
 
+EXTENSION_LANGUAGE_MAP = {
+    "py": "python",
+    "c": "c",
+    "cpp": "cpp",
+    "java": "java",
+    "js": "javascript",
+    "jsx": "jsx",
+    "ts": "typescript",
+    "go": "go",
+    "rs": "rust",
+    "php": "php",
+    "rb": "ruby",
+    "swift": "swift",
+    "kt": "kotlin",
+    "html": "html",
+    "css": "css"
+}
+
 # Helper: Read and decode code file
 def read_code_file(file):
     content = file.read().decode("utf-8")
@@ -42,7 +57,7 @@ def analyze_code(filename, code_text):
 # Extract corrected code from response
 def extract_corrected_code(response_text):
     code_blocks = re.findall(r"```(?:[a-zA-Z]*\n)?(.*?)```", response_text, re.DOTALL)
-    return code_blocks[0] if code_blocks else ""
+    return code_blocks[0].strip() if code_blocks else ""
 
 # Main App
 def main():
@@ -61,16 +76,26 @@ def main():
 
         for file in uploaded_files:
             filename, code_text = read_code_file(file)
+            lang_ext = filename.split(".")[-1]
+            lang = EXTENSION_LANGUAGE_MAP.get(lang_ext, "")
+
             st.subheader(f"📄 Analysis: `{filename}`")
+            st.markdown("### 🧾 Original Code:")
+            st.code(code_text, language=lang)
 
             with st.spinner(f"Analyzing `{filename}`..."):
                 response_text = analyze_code(filename, code_text)
 
+            st.markdown("### 🧠 AI Analysis:")
             st.markdown(response_text)
 
             corrected_code = extract_corrected_code(response_text)
             if corrected_code:
+                st.markdown("### ✅ Fixed Code:")
+                st.code(corrected_code, language=lang)
+
                 fixed_files.append((filename, corrected_code))
+
                 st.download_button(
                     label=f"📥 Download fixed `{filename}`",
                     data=corrected_code,