"""
Streamlit UI for the Code Explainer + Bug Fixer.
 
Run:
    pip install streamlit
    streamlit run streamlit_app.py
"""
 
import streamlit as st
import torch
from transformers import AutoTokenizer, T5ForConditionalGeneration
from line_explainer import explain_line_by_line
from bug_checker import detect_bugs, check_style
 
st.set_page_config(page_title="Code Explainer & Bug Fixer", page_icon="🛠️", layout="wide")
 
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
EXPLAIN_MODEL_NAME = "Salesforce/codet5-base-multi-sum"
 
 
@st.cache_resource(show_spinner="Loading explanation model... (first run only, this takes a minute)")
def load_explain_model():
    explain_tokenizer = AutoTokenizer.from_pretrained(EXPLAIN_MODEL_NAME, use_fast=False)
    explain_model = T5ForConditionalGeneration.from_pretrained(EXPLAIN_MODEL_NAME).to(DEVICE)
    return explain_tokenizer, explain_model
 
 
explain_tokenizer, explain_model = load_explain_model()
 
 
def explain_code(code: str) -> str:
    inputs = explain_tokenizer(code, return_tensors="pt", truncation=True, max_length=512).to(DEVICE)
    output_ids = explain_model.generate(**inputs, max_length=64, num_beams=5, early_stopping=True)
    return explain_tokenizer.decode(output_ids[0], skip_special_tokens=True)
 
 
# --------------------------- Language config ---------------------------
# codet5-base-multi-sum was trained on 6 languages for the EXPLANATION step.
# line_explainer.py and bug_checker.py are Python-only (built on Python's ast
# module), so those two features are gated to Python only below.
LANGUAGE_OPTIONS = {
    "Python": {"extensions": ["py"], "st_lang": "python", "supports_static_analysis": True},
    "Java": {"extensions": ["java"], "st_lang": "java", "supports_static_analysis": False},
    "JavaScript": {"extensions": ["js"], "st_lang": "javascript", "supports_static_analysis": False},
    "PHP": {"extensions": ["php"], "st_lang": "php", "supports_static_analysis": False},
    "Ruby": {"extensions": ["rb"], "st_lang": "ruby", "supports_static_analysis": False},
    "Go": {"extensions": ["go"], "st_lang": "go", "supports_static_analysis": False},
}
 
 
# --------------------------- UI ---------------------------
 
st.title("🛠️ Code Explainer & Bug Checker")
st.caption("Paste code or upload a file. It will be explained in plain English and checked for bugs/style issues.")
 
col_lang, col_upload = st.columns([1, 2])
with col_lang:
    language = st.selectbox("Language:", list(LANGUAGE_OPTIONS.keys()))
with col_upload:
    all_extensions = [ext for cfg in LANGUAGE_OPTIONS.values() for ext in cfg["extensions"]] + ["txt"]
    uploaded_file = st.file_uploader("Or upload a code file:", type=all_extensions)
 
lang_cfg = LANGUAGE_OPTIONS[language]
 
if not lang_cfg["supports_static_analysis"]:
    st.info(
        f"Line-by-line breakdown and bug/style checking currently support Python only. "
        f"For {language}, you'll get the AI explanation only."
    )
 
default_code = """def add_numbers(a, b)
    return a + b"""
 
if uploaded_file is not None:
    file_bytes = uploaded_file.read()
    try:
        starter_code = file_bytes.decode("utf-8")
    except UnicodeDecodeError:
        st.error("Couldn't decode the uploaded file as text (UTF-8). Please upload a plain text code file.")
        starter_code = default_code
else:
    starter_code = default_code
 
code_input = st.text_area("Your code:", value=starter_code, height=250, placeholder="Paste your code here...")
 
col1, col2 = st.columns(2)
with col1:
    run_explain = st.button("🔍 Analyze Code", type="primary", use_container_width=True)
with col2:
    clear = st.button("🗑️ Clear", use_container_width=True)
 
if clear:
    st.rerun()
 
if run_explain:
    if not code_input.strip():
        st.warning("Please paste some code first.")
    else:
        report_sections = []  # collects (title, text) for the downloadable report
 
        with st.spinner("Analyzing..."):
            explanation = explain_code(code_input)
 
        st.subheader("📝 Explanation")
        st.info(explanation)
        report_sections.append(("Explanation", explanation))
 
        if lang_cfg["supports_static_analysis"]:
            st.subheader("📋 Line-by-line breakdown")
            ok, breakdown = explain_line_by_line(code_input)
            if ok:
                breakdown_text = "\n".join(breakdown)
                st.code(breakdown_text, language="text")
                report_sections.append(("Line-by-line breakdown", breakdown_text))
            else:
                st.warning(breakdown)
                report_sections.append(("Line-by-line breakdown", breakdown))
 
            bug_result = detect_bugs(code_input)
            style_issues = check_style(code_input) if bug_result["syntax_valid"] else []
 
            st.subheader("🐛 Bug Check (static analysis)")
            st.caption("Checks for real errors: syntax mistakes, undefined variables, unused imports. Based on Python's own parser + pyflakes — deterministic, not a guess.")
            if not bug_result["syntax_valid"]:
                st.error(bug_result["syntax_error"])
                bug_report_text = bug_result["syntax_error"]
            elif bug_result["has_bug"]:
                st.error(bug_result["summary"])
                for issue in bug_result["static_issues"]:
                    st.write(f"- {issue}")
                bug_report_text = bug_result["summary"] + "\n" + "\n".join(f"- {i}" for i in bug_result["static_issues"])
            else:
                st.success(bug_result["summary"])
                bug_report_text = bug_result["summary"]
            report_sections.append(("Bug Check", bug_report_text))
 
            st.subheader("🎨 Style Check (PEP8)")
            st.caption("Checks formatting conventions (line length, spacing, naming) — not correctness. Code with style warnings can still run perfectly fine; this just flags deviations from Python's standard style guide.")
            if style_issues:
                st.warning(f"{len(style_issues)} style issue(s) found.")
                for issue in style_issues:
                    st.write(f"- {issue}")
                style_report_text = "\n".join(f"- {i}" for i in style_issues)
            else:
                st.success("No PEP8 style issues found.")
                style_report_text = "No PEP8 style issues found."
            report_sections.append(("Style Check (PEP8)", style_report_text))
 
        # --------------------------- Build downloadable report ---------------------------
        report_lines = [f"# Code Analysis Report\n", f"**Language:** {language}\n", "## Original Code\n", f"```{lang_cfg['st_lang']}\n{code_input}\n```\n"]
        for title, text in report_sections:
            report_lines.append(f"## {title}\n")
            report_lines.append(f"{text}\n")
        report_text = "\n".join(report_lines)
 
        st.divider()
        st.download_button(
            label="⬇️ Download Full Report (.md)",
            data=report_text,
            file_name="code_analysis_report.md",
            mime="text/markdown",
        )
 
st.divider()
st.caption(
    "Bug detection uses real static analysis (Python's ast parser + pyflakes), "
    "so results are accurate and deterministic — not a model's guess."
)