app.py

"""
CodeMentor PRO - Main Gradio App
"""
import os
import difflib
import tempfile

import gradio as gr

from config import SUPPORTED_LANGUAGES, EXPLANATION_LEVELS
from llm import call_groq_chat, call_groq_vision_ocr, LLMNotConfiguredError
from agent import (
    generate_breakdown, generate_plan, cleanup_code,
    suggest_optimal_solution, search_documentation, find_resources,
    extract_code_topic,
)
from error_detector import detect_errors
from sandbox import benchmark_two_versions, verify_fix_equivalence
from store import get_progress_snapshot
from prompts import FOLLOWUP_SYSTEM_PROMPT_TEMPLATE


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _validate_inputs(code: str) -> bool:
    if not code or not code.strip():
        gr.Warning("Please paste a code snippet before running an analysis.")
        return False
    return True

def _greeting_message(language=None, level=None):
    if language and level:
        text = (
            f"Got it! I've analyzed your **{language}** code at **{level}** level. "
            "Ask me anything about it — bugs, alternate approaches, or just \"explain line 5 again\"."
        )
    else:
        text = (
            "👋 Hi, I'm CodeMentor! Head over to the **Analyze Code** tab first and run an analysis — "
            "once that's done, come back here and I'll remember the context for follow-up questions."
        )
    return [{"role": "assistant", "content": text}]


def _safe_call(label, fn, *args, **kwargs):
    """Runs fn and never lets an exception escape to the UI as a raw traceback."""
    try:
        return fn(*args, **kwargs), None
    except LLMNotConfiguredError as e:
        return None, f"{label} unavailable — {e}"
    except Exception as e:
        return None, f"{label} failed — {e}"
    
def _citation_cards_html(citations: list) -> str:
    if not citations:
        return "<p style='opacity:0.7'>No citations found.</p>"
    cards = []
    for i, c in enumerate(citations, start=1):
        thumb = (
            f"<img src='{c['thumbnail']}' style='width:120px;height:68px;object-fit:cover;border-radius:6px;margin-right:12px;'/>"
            if c.get("thumbnail") else
            "<div style='width:120px;height:68px;background:#2a2a2a;border-radius:6px;margin-right:12px;"
            "display:flex;align-items:center;justify-content:center;font-size:22px;'>🔗</div>"
        )
        cards.append(f"""
        <div style='display:flex;align-items:center;padding:10px;border-bottom:1px solid #333;'>
          <div style='min-width:24px;font-weight:bold;opacity:0.6;'>{i}.</div>
          {thumb}
          <div>
            <a href='{c['url']}' target='_blank' style='font-weight:600;text-decoration:none;'>{c['title']}</a>
            <div style='font-size:12px;opacity:0.7;margin-top:2px;'>{c.get('snippet','')}</div>
          </div>
        </div>""")
    return "<div style='max-height:280px;overflow-y:auto;border:1px solid #333;border-radius:8px;'>" + "".join(cards) + "</div>"



def build_session_report(code, language, level, breakdown_md, analogy_md,
                          errors_md, cleanup_md, optimal_md, docs_md, resources_md) -> str:
    report = (
        "# CodeMentor PRO — Session Report\n\n"
        f"**Language:** {language}  |  **Explanation level:** {level}\n\n"
        "## Original Code\n"
        f"```{language}\n{code}\n```\n\n"
        "## Explanation\n"
        f"{breakdown_md}\n\n"
        "## Analogy\n"
        f"{analogy_md}\n\n"
        "## Errors & Weakness Review\n"
        f"{errors_md}\n\n"
        "## Cleanup\n"
        f"{cleanup_md}\n\n"
        "## Optimal Solution\n"
        f"{optimal_md}\n\n"
        "## Documentation Findings\n"
        f"{docs_md}\n\n"
        "## Further Resources\n"
        f"{resources_md}\n"
    )
    tmp_dir = tempfile.mkdtemp(prefix="codementor_report_")
    path = os.path.join(tmp_dir, "codementor_session_report.md")
    with open(path, "w", encoding="utf-8") as f:
        f.write(report)
    return path


# ---------------------------------------------------------------------------
# Tab 1 callbacks — Analyze Code
# ---------------------------------------------------------------------------

def on_generate_plan(code, language):
    if not _validate_inputs(code):
        return gr.update(value="", visible=False), gr.update(visible=False)

    plan, err = _safe_call("Plan generation", generate_plan, code, language)
    if err:
        gr.Warning(err)
        return gr.update(value="", visible=False), gr.update(visible=False)

    steps = (plan or {}).get("steps", [])
    plan_md = "**Here's what I'll run:**\n\n" + "\n".join(f"- {s}" for s in steps)
    if (plan or {}).get("estimated_focus"):
        plan_md += f"\n\n*Focus: {plan['estimated_focus']}*"

    return gr.update(value=plan_md, visible=True), gr.update(visible=True)

def on_extract_from_image(image_path):
    if not image_path:
        gr.Warning("Please upload an image first.")
        return gr.update()
    extracted, err = _safe_call("Vision OCR", call_groq_vision_ocr, image_path)
    if err:
        gr.Warning(err)
        return gr.update()
    if not extracted:
        gr.Warning("Couldn't extract any code from that image — try a clearer photo.")
        return gr.update()
    gr.Info("Code extracted — please review it before running analysis, OCR can misread characters.")
    return gr.update(value=extracted)


def on_cancel_plan():
    return gr.update(value="", visible=False), gr.update(visible=False)


def on_approve_run(code, language, level, progress=gr.Progress()):
    empty_return = ("", "", "", "", "", "", "", None, gr.update(visible=False), gr.update(visible=False))
    if not _validate_inputs(code):
        return empty_return

    # --- Explanation & breakdown ---
    progress(0.1, desc="Explaining code...")
    breakdown_result, err = _safe_call("Explanation", generate_breakdown, code, language, level)
    if err:
        gr.Warning(err)
    breakdown_items = (breakdown_result or {}).get("breakdown", [])
    breakdown_md = "\n\n".join(
        f"**{item.get('line_or_block', '')}** — {item.get('explanation', '')}" for item in breakdown_items
    ) or "_No breakdown available._"
    analogy_md = (breakdown_result or {}).get("analogy") or "_No analogy returned._"
    bugs = (breakdown_result or {}).get("bugs", [])

    # --- Multi-type error detection & weakness review ---
    progress(0.3, desc="Detecting errors & weaknesses...")
    err_result, err = _safe_call("Error detection", detect_errors, code, language)
    if err:
        gr.Warning(err)

    content_lines = []
    if err_result:
        if err_result["syntax"]["found"]:
            content_lines.append(f"🔴 **Syntax error**: {err_result['syntax']['message']}")
        if err_result["runtime"]["found"]:
            content_lines.append(f"🔴 **Runtime error**: {err_result['runtime']['message']}")
        if err_result["tle"]["found"]:
            content_lines.append(f"🟠 **TLE**: {err_result['tle']['message']}")
        for issue in err_result["logical"]["issues"]:
            content_lines.append(
                f"🟡 **Logical issue**: {issue.get('description', '')}\n   ↳ Fix: {issue.get('suggested_fix', '')}"
            )
        for bug in bugs:
            content_lines.append(f"🟡 {bug.get('description', '')}\n   ↳ Fix: {bug.get('concrete_fix', '')}")
        if not content_lines:
            content_lines.append("✅ No errors detected.")
        if err_result.get("weakness_note"):
            content_lines.append(f"\n📊 **Weakness note**: {err_result['weakness_note']}")
        if not err_result.get("verified_by_execution"):
            content_lines.append("\n_Note: this language is checked via LLM static estimate, not sandboxed execution._")
    elif err:
        content_lines.append(f"⚠️ {err}")
    else:
        content_lines.append("No error data available.")
    errors_md = "\n\n".join(content_lines)

    # --- Cleanup pass ---
    progress(0.5, desc="Cleaning up code...")
    clean_res, err = _safe_call("Cleanup", cleanup_code, code, language)
    if err:
        gr.Warning(err)
    clean_res = clean_res or {}
    cleaned_code = clean_res.get("cleaned_code", "") or ""

    cleanup_lines = []
    if clean_res.get("needs_cleanup") and cleaned_code.strip():
        diff = difflib.unified_diff(
            code.splitlines(), cleaned_code.splitlines(),
            fromfile="original", tofile="cleaned", lineterm="",
        )
        diff_text = "\n".join(diff)
        cleanup_lines.append(f"```diff\n{diff_text}\n```" if diff_text else "No line-level differences.")

        changes = clean_res.get("changes_made", [])
        if changes:
            cleanup_lines.append("**Changes made:**\n" + "\n".join(f"- {c}" for c in changes))

        if language == "python":
            equivalence, eq_err = _safe_call("Fix verification", verify_fix_equivalence, code, cleaned_code)
            if eq_err:
                cleanup_lines.append(f"⚠️ Could not verify equivalence — {eq_err}")
            elif equivalence:
                if equivalence["outputs_match"]:
                    cleanup_lines.append("✅ **Verified by execution**: cleaned code produces identical output to the original.")
                else:
                    cleanup_lines.append(
                        "⚠️ **Could not confirm equivalence** — outputs differ. Please review the cleaned version before trusting it."
                    )
    elif err:
        cleanup_lines.append(f"⚠️ {err}")
    else:
        cleanup_lines.append("Code already looks clean — no restructuring needed.")
    cleanup_md = "\n\n".join(cleanup_lines)

    # --- Optimal solution pass ---
    progress(0.7, desc="Checking for a better approach...")
    opt_res, err = _safe_call("Optimization", suggest_optimal_solution, code, language)
    if err:
        gr.Warning(err)
    opt_res = opt_res or {}

    optimal_lines = []
    cur = opt_res.get("current_complexity", {}) or {}
    structures = opt_res.get("data_structures_used", []) or []
    pattern = opt_res.get("algorithm_pattern", "")

    if opt_res:
        optimal_lines.append(f"**Time complexity:** {cur.get('time', '?')}  |  **Space complexity:** {cur.get('space', '?')}")
        if structures:
            optimal_lines.append(f"**Data structures used:** {', '.join(structures)}")
        if pattern:
            optimal_lines.append(f"**Pattern:** {pattern}")

        if opt_res.get("has_better_approach"):
            sug = opt_res.get("suggested_complexity", {}) or {}
            optimal_lines.append(f"**Suggested complexity:** time {sug.get('time', '?')}, space {sug.get('space', '?')}")
            if opt_res.get("explanation"):
                optimal_lines.append(f"**Why it's better:** {opt_res['explanation']}")
            suggested_code = opt_res.get("suggested_code", "") or ""
            if suggested_code.strip():
                optimal_lines.append(f"```{language}\n{suggested_code}\n```")
                if language == "python":
                    bench, bench_err = _safe_call("Benchmark", benchmark_two_versions, code, suggested_code)
                    if bench_err:
                        optimal_lines.append(f"⚠️ Could not measure runtime — {bench_err}")
                    elif bench:
                        optimal_lines.append(
                            f"**Measured runtime** — original: {bench['original_runtime']}s "
                            f"({'ok' if bench['original_ok'] else 'failed/TLE'}), "
                            f"suggested: {bench['optimized_runtime']}s "
                            f"({'ok' if bench['optimized_ok'] else 'failed/TLE'})"
                        )
                        if bench.get("speedup"):
                            optimal_lines.append(f"🚀 Speedup: {bench['speedup']}x")
        else:
            optimal_lines.append(f"**Why this is already optimal:** {opt_res.get('explanation', 'No further improvement possible.')}")
    elif err:
        optimal_lines.append(f"⚠️ {err}")
    else:
        optimal_lines.append("Optimization analysis unavailable.")
    optimal_md = "\n\n".join(optimal_lines)


    # --- Documentation search ---
    progress(0.85, desc="Searching documentation...")
    doc_findings, err = _safe_call("Documentation search", search_documentation, code)
    if err:
        docs_html = f"<p>⚠️ {err}</p>"
    elif doc_findings:
        blocks = []
        for d in doc_findings:
            blocks.append(f"<h4>{d['library']}</h4><p>{d['summary']}</p>")
            blocks.append(_citation_cards_html(d.get("citations", [])))
        docs_html = "".join(blocks)
    else:
        docs_html = "<p>No unfamiliar libraries detected — nothing to look up.</p>"

    # --- Resource finder (topic-based, not just language-based) ---
    progress(0.9, desc="Identifying core topic...")
    topic, err = _safe_call("Topic extraction", extract_code_topic, code, language)
    topic = topic or f"{language} programming"

    progress(0.95, desc=f"Finding resources on {topic}...")
    resources, err = _safe_call("Resource search", find_resources, language, topic)
    if err:
        resources_html = f"<p>⚠️ {err}</p>"
    elif resources and resources.get("ok"):
        topic_header = (
            f"<p style='opacity:0.75;font-size:13px;margin-bottom:8px;'>"
            f"📌 Topic identified: <strong>{topic}</strong></p>"
        )
        resources_html = topic_header + _citation_cards_html(resources.get("citations", []))
    else:
        resources_html = f"<p>No resources found for <strong>{topic}</strong> this time — try again in a moment.</p>"


    # --- Downloadable report ---
    report_path = None
    try:
        report_path = build_session_report(
            code, language, level, breakdown_md, analogy_md,
            errors_md, cleanup_md, optimal_md, docs_html, resources_html,
        )
    except Exception as e:
        gr.Warning(f"Could not build downloadable report: {e}")

    progress(1.0, desc="Done!")
    return (
        breakdown_md, analogy_md, errors_md, cleanup_md, optimal_md,
        docs_html, resources_html, report_path,
        gr.update(visible=False), gr.update(visible=False),
    )
    

# ---------------------------------------------------------------------------
# Tab 2 callbacks — Chat / Follow-up
# ---------------------------------------------------------------------------

def on_chat_send(message, chat_history, last_code, last_language, last_level):
    chat_history = chat_history or []
    if not message or not message.strip():
        return chat_history, ""
    if not last_code:
        gr.Warning("Analyze a code snippet in the 'Analyze Code' tab first so I have context for follow-ups.")
        chat_history = chat_history + [
            {"role": "user", "content": message},
            {"role": "assistant", "content": "⚠️ I don't have any code to discuss yet — please run an analysis in the **Analyze Code** tab first."},
        ]
        return chat_history, ""

    system_prompt = FOLLOWUP_SYSTEM_PROMPT_TEMPLATE.format(
        level=last_level or "Intermediate",
        language=last_language or "python",
        code=last_code,
    )
    api_history = chat_history

    try:
        reply = call_groq_chat(system_prompt, api_history, message)
    except LLMNotConfiguredError as e:
        reply = f"⚠️ {e}"
    except Exception as e:
        reply = f"⚠️ Something went wrong answering that: {e}"

    chat_history = chat_history + [
        {"role": "user", "content": message},
        {"role": "assistant", "content": reply},
    ]
    return chat_history, ""
# ---------------------------------------------------------------------------
# Tab 3 callbacks — My Progress
# ---------------------------------------------------------------------------

def on_refresh_progress():
    try:
        return get_progress_snapshot()
    except Exception as e:
        gr.Warning(f"Could not refresh progress: {e}")
        return gr.update(), gr.update(), gr.update()

# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------

with gr.Blocks(title="CodeMentor PRO") as demo:
    gr.Markdown("# 🧑‍💻 CodeMentor PRO\nYour personal, execution-verified coding coach.")

    state_last_code = gr.State("")
    _initial_table, _initial_chart, _initial_summary = get_progress_snapshot()
    state_last_language = gr.State(SUPPORTED_LANGUAGES[0])
    state_last_level = gr.State(EXPLANATION_LEVELS[0])

    with gr.Tab("Analyze Code"):
        with gr.Row():
            with gr.Column(scale=2):
                code_input = gr.Textbox(label="Paste your code", lines=16, placeholder="Paste your code here, or upload a screenshot below…")
                with gr.Accordion("📸 Or upload a code screenshot", open=False):
                    image_input = gr.Image(label="Code screenshot", type="filepath", height=240)
                    extract_btn = gr.Button("🔍 Extract Code from Image", variant="secondary")
            with gr.Column(scale=1):
                language_dd = gr.Dropdown(SUPPORTED_LANGUAGES, value=SUPPORTED_LANGUAGES[0], label="Language")
                level_dd = gr.Dropdown(EXPLANATION_LEVELS, value=EXPLANATION_LEVELS[0], label="Explanation level")
                plan_btn = gr.Button("1. Preview Plan", variant="primary")

        plan_output = gr.Markdown(visible=False)
        with gr.Row(visible=False) as plan_actions:
            approve_btn = gr.Button("✅ Approve & Run Analysis", variant="primary")
            cancel_btn = gr.Button("❌ Cancel")

        with gr.Accordion("📖 Breakdown", open=True):
            breakdown_out = gr.Markdown()
        with gr.Accordion("💡 Analogy", open=False):
            analogy_out = gr.Markdown()
        with gr.Accordion("🐞 Errors & Weakness Review", open=True):
            errors_out = gr.Markdown()
        with gr.Accordion("🧹 Cleanup Diff", open=False):
            cleanup_out = gr.Markdown()
        with gr.Accordion("⚡ Optimal Solution", open=False):
            optimal_out = gr.Markdown()
        with gr.Accordion("📚 Documentation Findings", open=False):
            docs_out = gr.HTML()
        with gr.Accordion("🔗 Further Resources", open=False):
            resources_out = gr.HTML()

        report_file = gr.File(label="Download session report")

    with gr.Tab("Chat / Follow-up"):
        gr.Markdown("Ask follow-up questions about the code you analyzed — I'll remember the context.")
        chatbot = gr.Chatbot(label="CodeMentor", height=400, value=_greeting_message())
        chat_input = gr.Textbox(label="Your question", placeholder="e.g. why did you suggest that fix?")
        chat_send_btn = gr.Button("Send", variant="primary")

    with gr.Tab("My Progress"):
        gr.Markdown("Your language-wise revision history and skill breakdown.")
        refresh_btn = gr.Button("🔄 Refresh")
        revision_table = gr.Dataframe(
            headers=["Language", "Submissions", "Syntax Errors", "Logical Errors", "TLE Count", "Last Seen"],
            value=_initial_table,
            interactive=False,
        )
        skill_pie = gr.Plot(value=_initial_chart)
        skill_summary = gr.Markdown(value=_initial_summary)

    # ------------------------------ Event wiring ------------------------------
     
    plan_btn.click(
        on_generate_plan,
        inputs=[code_input, language_dd],
        outputs=[plan_output, plan_actions],
    )

    cancel_btn.click(
        on_cancel_plan,
        outputs=[plan_output, plan_actions],
    )
     
    extract_btn.click(
        on_extract_from_image,
        inputs=[image_input],
        outputs=[code_input],
    )
    
    approve_btn.click(
        on_approve_run,
        inputs=[code_input, language_dd, level_dd],
        outputs=[
            breakdown_out, analogy_out, errors_out, cleanup_out,
            optimal_out, docs_out, resources_out, report_file,
            plan_output, plan_actions,
        ],
    ).then(
    lambda c, l, lv: (c, l, lv, _greeting_message(language=l, level=lv)),
    inputs=[code_input, language_dd, level_dd],
    outputs=[state_last_code, state_last_language, state_last_level, chatbot],
     )

    chat_send_btn.click(
        on_chat_send,
        inputs=[chat_input, chatbot, state_last_code, state_last_language, state_last_level],
        outputs=[chatbot, chat_input],
    )
    chat_input.submit(
        on_chat_send,
        inputs=[chat_input, chatbot, state_last_code, state_last_language, state_last_level],
        outputs=[chatbot, chat_input],
    )

    refresh_btn.click(on_refresh_progress, outputs=[revision_table, skill_pie, skill_summary])

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