import streamlit as st
import os
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"))

# Supported extensions and corresponding language names for syntax highlighting
SUPPORTED_EXTENSIONS = [
    "py", "c", "cpp", "java", "js", "jsx", "go", "rs", "php", "rb", "swift", "kt", "html", "css"
]

EXTENSION_LANGUAGE_MAP = {
    "py": "python",
    "c": "c",
    "cpp": "cpp",
    "java": "java",
    "js": "javascript",
    "jsx": "jsx",
    "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")
    filename = file.name
    return filename, content

# Analyze a single file and return result
def analyze_code(filename, code_text):
    prompt = (
        f"You are a software code analysis assistant. Given the code below, perform the following tasks:\n\n"
        f"1. Provide a short README-style summary of what the code is intended to do.\n"
        f"2. Identify any bugs or issues in the code.\n"
        f"3. Provide a corrected, clean version of the entire code.\n"
        f"4. List the changes you made in bullet points.\n\n"
        f"--- FILE: {filename} ---\n{code_text}\n--- END ---"
    )

    response = client.models.generate_content(
        model="gemini-2.5-flash",
        contents=[{"text": prompt}]
    )
    return response.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].strip() if code_blocks else ""

# Main App
def main():
    st.set_page_config(page_title="Multi-Language Code Analyzer & Fixer", layout="centered")
    st.title("Multi-Language Code Analyzer & Bug Fixer")
    st.write("Upload one or more code files and get a full analysis with downloadable fixed versions.")

    uploaded_files = st.file_uploader(
        "Upload your code file(s)",
        type=SUPPORTED_EXTENSIONS,
        accept_multiple_files=True
    )

    if uploaded_files and st.button("🔍 Analyze and Fix All"):
        fixed_files = []

        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,
                    file_name=f"fixed_{filename}",
                    mime="text/plain"
                )
            else:
                st.warning(f"❌ Could not extract corrected code for `{filename}`.")

        if fixed_files:
            zip_buffer = io.BytesIO()
            with zipfile.ZipFile(zip_buffer, "w") as zip_file:
                for filename, fixed_code in fixed_files:
                    zip_file.writestr(f"fixed_{filename}", fixed_code)
            zip_buffer.seek(0)

            st.download_button(
                label="📦 Download All Fixed Files as ZIP",
                data=zip_buffer,
                file_name="fixed_code_files.zip",
                mime="application/zip"
            )

if __name__ == "__main__":
    main()