app.py

import time
import uuid
from typing import Generator

import streamlit as st

from src.config import ConfigError, PLOTS_DIR, get_settings
from src.costs import read_costs
from src.graph.workflow import resume_task, stream_task, stream_task_hitl


def _token_stream(text: str) -> Generator[str, None, None]:
    """Word-by-word generator for st.write_stream — gives typewriter effect.

    Uses the real stream_answer_tokens() API shape so this is a drop-in
    swap for production token streaming when the LLM supports it.
    """
    for word in text.split():
        yield word + " "
        time.sleep(0.012)

def _handle_api_error(exc: Exception) -> None:
    """Show a friendly banner for known API errors instead of a raw traceback."""
    msg = str(exc)
    # Groq / OpenAI rate limit
    if "429" in msg or "rate_limit" in msg.lower() or "rate limit" in msg.lower():
        if "tokens per day" in msg.lower() or "tpd" in msg.lower():
            st.error(
                "**Daily token limit reached** — you've used your full Groq free-tier quota (100k tokens/day). "
                "It resets at midnight UTC. Options:\n"
                "- Wait until tomorrow\n"
                "- Set `GROQ_MODEL=llama-3.1-8b-instant` in `.env` (separate quota)\n"
                "- Upgrade at https://console.groq.com/settings/billing"
            )
        else:
            import re
            wait = re.search(r"try again in ([\d.]+[ms]+)", msg)
            wait_str = f" Try again in **{wait.group(1)}**." if wait else ""
            st.error(f"**Rate limit hit** — too many requests per minute.{wait_str}")
    # 402 insufficient balance (DeepSeek / paid APIs)
    elif "402" in msg or "insufficient balance" in msg.lower() or "insufficient_balance" in msg.lower():
        st.error(
            "**Insufficient API balance** — your account has no credits.\n\n"
            "Options:\n"
            "- **DeepSeek**: top up at https://platform.deepseek.com/top-up\n"
            "- **Switch to Groq** (free): set `LLM_PROVIDER=groq` in `.env`"
        )
    # Auth errors
    elif "401" in msg or "authentication" in msg.lower() or "api key" in msg.lower():
        st.error("**Authentication error** — check your API key in `.env`.")
    # 400 model_decommissioned
    elif "model_decommissioned" in msg or "decommissioned" in msg.lower():
        import re
        m = re.search(r"model `?([^\s`']+)`?", msg)
        name = m.group(1) if m else "this model"
        st.error(
            f"**Model decommissioned:** `{name}` is no longer available on Groq.\n\n"
            "Update `GROQ_MODEL` in `.env` to an active model:\n"
            "```\nGROQ_MODEL=llama-3.3-70b-versatile   # best quality\n"
            "GROQ_MODEL=llama-3.1-70b-versatile   # fallback quota\n```\n"
            "Check active models: https://console.groq.com/docs/models"
        )
    # 400 tool_use_failed — model too small for tool calling
    elif "400" in msg and ("tool_use_failed" in msg or "failed_generation" in msg):
        st.error(
            "**Model tool-calling failure** — the selected model is too small to reliably call tools.\n\n"
            "Fix: switch to a larger model in `.env`:\n"
            "```\nGROQ_MODEL=qwen/qwen3-32b          # 32B, reliable tool use\n"
            "GROQ_MODEL=llama-3.3-70b-versatile  # best quality\n```"
        )
    # Network / timeout
    elif "timeout" in msg.lower() or "connection" in msg.lower():
        st.error("**Connection error** — Groq API unreachable. Check your internet connection.")
    else:
        st.error(f"**Agent error:** {exc}")


def _show_recent_plots(since: float, max_age_s: int = 120) -> None:
    """Display any image files written to PLOTS_DIR after `since` timestamp."""
    if not PLOTS_DIR.exists():
        return
    imgs = sorted(
        (p for p in PLOTS_DIR.glob("*") if p.suffix.lower() in {".png", ".jpg", ".jpeg"}),
        key=lambda p: p.stat().st_mtime,
    )
    recent = [p for p in imgs if p.stat().st_mtime >= since - 1]
    if not recent:
        return
    st.divider()
    st.subheader(f"Generated plot{'s' if len(recent) > 1 else ''}")
    for img_path in recent:
        if img_path.stat().st_size == 0:
            st.warning(f"{img_path.name} is empty — plot may not have saved correctly.")
            continue
        try:
            import base64 as _b64
            import io as _io
            from PIL import Image as _PIL

            # Flatten RGBA transparency onto white — avoids invisible-on-white-bg issue
            img = _PIL.open(img_path)
            if img.mode in ("RGBA", "LA", "P"):
                bg = _PIL.new("RGB", img.size, (255, 255, 255))
                bg.paste(img, mask=img.convert("RGBA").split()[3])
                img = bg
            else:
                img = img.convert("RGB")

            buf = _io.BytesIO()
            img.save(buf, format="PNG")
            b64 = _b64.b64encode(buf.getvalue()).decode()
            st.markdown(
                f'<img src="data:image/png;base64,{b64}" '
                f'style="max-width:100%;border-radius:4px;display:block">',
                unsafe_allow_html=True,
            )
            st.caption(img_path.name)
        except Exception as _e:
            st.warning(f"Could not render {img_path.name}: {_e}")


st.set_page_config(page_title="Coding Assistant Agent", layout="wide")

st.title("Coding Assistant Agent")
st.caption("LangGraph · Groq · planner → executor → critic · streaming · HITL · conversation memory")

# ── Session state init ────────────────────────────────────────────────────
for key, default in [
    ("thread_id", str(uuid.uuid4())),
    ("history_display", []),
    ("hitl_pending", False),       # waiting for human approve/revise
    ("hitl_snapshot", {}),         # last state snapshot before interrupt
    ("hitl_thread", None),         # thread_id of interrupted task
]:
    if key not in st.session_state:
        st.session_state[key] = default

# ── Sidebar ───────────────────────────────────────────────────────────────
with st.sidebar:
    st.header("Status")
    try:
        settings = get_settings()
        st.success("Configuration loaded")
        model_name = {
            "groq": settings.groq_model,
            "openai": settings.openai_model,
            "anthropic": settings.anthropic_model,
            "deepseek": settings.deepseek_model,
            "gemini": settings.gemini_model,
        }.get(settings.llm_provider, settings.groq_model)
        st.caption(f"Model: {model_name}")
        st.caption(f"Provider: {settings.llm_provider}")
        st.caption(f"Max fix attempts: {settings.max_fix_attempts}")
        st.caption(f"Session: `{st.session_state.thread_id[:8]}…`")
    except ConfigError as exc:
        st.error(str(exc))
        st.stop()

    hitl_mode = st.toggle("Human-in-the-loop (review before critic)", value=False)

    if st.button("New session (clear memory)", use_container_width=True):
        st.session_state.thread_id = str(uuid.uuid4())
        st.session_state.history_display = []
        st.session_state.hitl_pending = False
        st.rerun()

    # ── Cost display ──────────────────────────────────────────────────────
    costs = read_costs()
    if costs:
        st.divider()
        st.header("Token usage (this session)")
        total_in = sum(c.get("prompt_tokens", 0) for c in costs)
        total_out = sum(c.get("completion_tokens", 0) for c in costs)
        st.metric("Prompt tokens", f"{total_in:,}")
        st.metric("Completion tokens", f"{total_out:,}")

    st.divider()
    st.header("Try these")
    st.markdown(
        "- Write a function to check if a number is prime\n"
        "- Read a file called data.txt\n"
        "- Generate 100 random numbers and plot a histogram\n"
        "- Implement binary search and test it\n"
        "- Fix this code: `def divide(a,b): return a/b` then `print(divide(10,0))`\n"
        "- *(follow-up)* Now make it handle floats too"
    )

    if st.session_state.history_display:
        st.divider()
        st.header("Session history")
        for i, (t, _) in enumerate(st.session_state.history_display, 1):
            st.caption(f"{i}. {t[:60]}{'…' if len(t) > 60 else ''}")


# ── HITL review UI (shown when graph interrupted) ─────────────────────────
if st.session_state.hitl_pending:
    snap = st.session_state.hitl_snapshot
    hitl_run_start = st.session_state.get("hitl_run_start", time.time() - 30)
    st.info("**Human review required** — inspect the executor's output, then approve or request a revision.")

    col_left, col_right = st.columns([1, 2])
    with col_left:
        st.subheader("Executor output")
        code_runs = snap.get("code_runs", [])
        st.markdown(f"**{len(code_runs)} code run(s)**")
        for i, run in enumerate(code_runs, 1):
            with st.expander(f"Run {i}", expanded=(i == len(code_runs))):
                st.code(run["code"], language="python")
                st.text("Output:")
                st.code(run["output"])
    with col_right:
        st.subheader("Draft answer")
        st.markdown(snap.get("final_answer", "_No answer yet._"))

    # Show plots outside columns so layout doesn't constrain the image
    _show_recent_plots(hitl_run_start)

    st.divider()
    col_a, col_r = st.columns(2)
    with col_a:
        if st.button("Approve", type="primary", use_container_width=True):
            with st.spinner("Resuming graph…"):
                final = resume_task(st.session_state.hitl_thread, "approve")
            st.session_state.hitl_pending = False
            st.session_state.history_display.append(
                (snap.get("task", ""), final.get("final_answer", "")[:300])
            )
            st.success("Approved — task complete.")
            st.write_stream(_token_stream(final.get("final_answer", "")))
            _show_recent_plots(hitl_run_start)
            st.stop()
    with col_r:
        feedback_text = st.text_area("Revision instructions", height=80)
        if st.button("Request revision", use_container_width=True):
            if not feedback_text.strip():
                st.warning("Enter revision instructions first.")
            else:
                with st.spinner("Sending back to executor…"):
                    final = resume_task(
                        st.session_state.hitl_thread, "revise", feedback=feedback_text.strip()
                    )
                st.session_state.hitl_pending = False
                st.session_state.history_display.append(
                    (snap.get("task", ""), final.get("final_answer", "")[:300])
                )
                st.success("Revision complete.")
                st.write_stream(_token_stream(final.get("final_answer", "")))
                _show_recent_plots(hitl_run_start)
                st.stop()
    st.stop()


# ── Main input ────────────────────────────────────────────────────────────
task = st.text_area(
    "Coding task",
    placeholder="e.g. Implement binary search and test it on 20 numbers",
    height=100,
)
run_btn = st.button("Solve", type="primary", use_container_width=True)

if run_btn:
    if not task.strip():
        st.warning("Enter a coding task first.")
        st.stop()

    thread_id = st.session_state.thread_id
    run_start_time = time.time()
    final: dict = {}
    plan_shown = False
    exec_shown = False
    critic_shown = False

    if hitl_mode:
        # ── HITL streaming (stops at interrupt) ───────────────────────────
        try:
            with st.status("Agents working (HITL mode)…", expanded=True) as status_widget:
                last_snap: dict = {}
                for snapshot in stream_task_hitl(task.strip(), thread_id=thread_id):
                    last_snap = snapshot

                    if snapshot.get("plan") and not plan_shown:
                        plan_shown = True
                        st.write("✅ **Planner** — plan ready")
                        for i, step in enumerate(snapshot["plan"], 1):
                            st.write(f"&nbsp;&nbsp;&nbsp;{i}. {step}")

                    if snapshot.get("final_answer") and not exec_shown:
                        exec_shown = True
                        n = len(snapshot.get("code_runs", []))
                        st.write(f"⚙️ **Executor** — {n} code run(s) complete")

                status_widget.update(label="Waiting for human review…", state="running")
        except Exception as exc:
            _handle_api_error(exc)
            st.stop()

        # Graph interrupted before critic — hand off to HITL UI
        last_snap["task"] = task.strip()
        st.session_state.hitl_pending = True
        st.session_state.hitl_snapshot = last_snap
        st.session_state.hitl_thread = thread_id
        st.session_state.hitl_run_start = run_start_time
        st.rerun()

    else:
        # ── Normal streaming ──────────────────────────────────────────────
        try:
            with st.status("Agents working…", expanded=True) as status_widget:
                for snapshot in stream_task(task.strip(), thread_id=thread_id):
                    final = snapshot

                    if snapshot.get("plan") and not plan_shown:
                        plan_shown = True
                        st.write("✅ **Planner** — plan ready")
                        for i, step in enumerate(snapshot["plan"], 1):
                            st.write(f"&nbsp;&nbsp;&nbsp;{i}. {step}")

                    if snapshot.get("final_answer") and not exec_shown:
                        exec_shown = True
                        n = len(snapshot.get("code_runs", []))
                        st.write(f"⚙️ **Executor** — {n} code run(s) complete")

                    if snapshot.get("verdict") and not critic_shown:
                        critic_shown = True
                        v = snapshot["verdict"]
                        icon = "✅" if v == "APPROVE" else "🔄"
                        st.write(f"{icon} **Critic** — {v}: {snapshot.get('critique', '')[:120]}")

                status_widget.update(label="Complete!", state="complete")

        except Exception as exc:
            _handle_api_error(exc)
            st.stop()

        if not final:
            st.error("No result returned.")
            st.stop()

        # ── Results layout ────────────────────────────────────────────────
        st.divider()
        col_left, col_right = st.columns([1, 2])

        plan = final.get("plan", [])
        code_runs = final.get("code_runs", [])
        verdict = final.get("verdict", "")
        critique_text = final.get("critique", "")
        fix_attempts = final.get("fix_attempts", 0)

        with col_left:
            st.subheader("Agent Summary")
            st.markdown("**Planner**")
            for i, step in enumerate(plan, 1):
                st.markdown(f"&nbsp;&nbsp;{i}. {step}")
            st.markdown(f"**Executor** — {len(code_runs)} run(s)")
            st.markdown("**Critic**")
            if verdict == "APPROVE":
                st.success(f"APPROVED — {critique_text}")
            else:
                st.warning(f"REVISE after {fix_attempts} attempt(s) — {critique_text}")

        with col_right:
            st.subheader("Answer")
            answer = final.get("final_answer", "_No answer generated._")
            st.write_stream(_token_stream(answer))

        # ── Code runs ─────────────────────────────────────────────────────
        if final.get("code_runs"):
            st.divider()
            st.subheader(f"Executed code ({len(final['code_runs'])} run(s))")
            for i, run in enumerate(final["code_runs"], 1):
                with st.expander(f"Run {i}", expanded=(i == len(final["code_runs"]))):
                    st.code(run["code"], language="python")
                    st.text("Output:")
                    st.code(run["output"])

        # ── Metrics ───────────────────────────────────────────────────────
        st.divider()
        c1, c2, c3, c4 = st.columns(4)
        c1.metric("Plan steps", len(plan))
        c2.metric("Code runs", len(code_runs))
        c3.metric("Fix attempts", fix_attempts)
        c4.metric("Session tasks", len(st.session_state.history_display) + 1)

        # ── Generated plots ───────────────────────────────────────────────
        _show_recent_plots(run_start_time)

        st.session_state.history_display.append(
            (task.strip(), final.get("final_answer", "")[:300])
        )