app.py

#!/usr/bin/env python3
"""
Recode — Spaced repetition for ML code.
Drop .py scripts into the writable problems directory shown by `recode --paths`.
Run: uv run python -m recode
"""
from __future__ import annotations

import difflib
import os
import subprocess
import tempfile
from pathlib import Path

from rich.markup import escape
from rich.syntax import Syntax
from textual.app import App, ComposeResult
from textual.binding import Binding
from textual.containers import Vertical
from textual.screen import Screen
from textual.widgets import DataTable, Footer, Header, Input, RichLog, Static

from ai import get_explain, get_hint, get_suggest_fix, opencode_chat, opencode_chat_health
from db import get_db, get_row, get_streak, log_mistake, recent_mistakes, reset_progress, sm2_update
from modals import AIModal, ChatModal, ConfirmModal, RatingModal, CollectionSelectModal, PaperGenerateModal, ImportProblemModal
from problems_utils import (
    build_side_by_side,
    get_problem_id,
    has_test_cases,
    is_marimo_problem,
    load_problem_meta,
    max_rating_for,
    problem_badges,
    scan_problems,
    scan_collections,
    status_label,
)
from test_runner import run_tests, format_test_results
from paper_generator import generate_problems, parse_arxiv_url, fetch_paper, extract_sections
from themes import TERMINAL_SEXY_THEMES
from recode.runtime import RuntimePaths, get_runtime, prepare_runtime

# ── Config ────────────────────────────────────────────────────────────────────
RUNTIME      = get_runtime()
PROBLEMS_DIR = RUNTIME.problems_dir
DB_PATH      = RUNTIME.db_path
EDITOR       = RUNTIME.editor
_TMP         = Path(tempfile.gettempdir())

RATING_LABELS = {1: "Again", 2: "Hard", 3: "Good", 4: "Easy"}


def configure_runtime(runtime: RuntimePaths | None = None) -> RuntimePaths:
    global RUNTIME, PROBLEMS_DIR, DB_PATH, EDITOR
    RUNTIME = runtime or prepare_runtime()
    PROBLEMS_DIR = RUNTIME.problems_dir
    DB_PATH = RUNTIME.db_path
    EDITOR = RUNTIME.editor
    return RUNTIME


# ── Study Screen ──────────────────────────────────────────────────────────────
class StudyScreen(Screen):
    BINDINGS = [
        Binding("e", "edit",        "Edit"),
        Binding("s", "submit",      "Submit"),
        Binding("h", "hint",        "Hint"),
        Binding("f", "suggest_fix", "Fix"),
        Binding("c", "chat",        "Chat"),
        Binding("x", "explain",     "Explain"),
        Binding("q", "back",        "Menu"),
    ]

    def __init__(self, problem: Path) -> None:
        super().__init__()
        self.problem   = problem
        self.pid       = get_problem_id(problem, PROBLEMS_DIR)
        self.meta      = load_problem_meta(problem)
        # Use pid in temp file name to avoid collisions
        safe_pid       = self.pid.replace("/", "_").replace("\\", "_")
        ext            = problem.suffix  # preserve .py, .jl, .R etc.
        self.work_file = _TMP / f"recode_{safe_pid}{ext}"
        self.conn      = get_db(DB_PATH)
        self.attempts  = 0
        self.has_diff  = False
        self.chat_session_id: str | None = None

    def compose(self) -> ComposeResult:
        yield Header()
        yield Static(id="problem-bar")
        yield RichLog(id="diff-pane", classes="full-pane", highlight=False, markup=True, wrap=False, auto_scroll=False)
        yield Footer()

    def on_mount(self) -> None:
        desc = self.meta["description"]
        desc_part = f"  [dim italic]{escape(desc)}[/]" if desc else ""
        badges = problem_badges(self.problem)
        badge_str = " " + " ".join(f"[{b[1]}]" for b in badges) if badges else ""
        self.query_one("#problem-bar", Static).update(
            f"[bold white]{escape(self.problem.name)}[/]{badge_str}{desc_part}"
        )
        log = self.query_one("#diff-pane", RichLog)
        row = get_row(self.conn, self.pid)
        if row and row["last_output"] and row["last_output"] != "✓  perfect match":
            log.write("[dim]── last session ──[/]")
            log.write(Syntax(row["last_output"], "diff", theme="ansi_dark", word_wrap=False))
            log.write("[dim]press  e  to start new attempt[/]")
        else:
            log.write("[dim]press  e  to open editor[/]")

    def action_hint(self) -> None:
        ref  = self.meta["solution"]
        user = self.work_file.read_text() if self.work_file.exists() else None
        self.app.push_screen(AIModal(
            f"hint  ·  {self.problem.name}",
            lambda: get_hint(self.problem.name, ref, user),
        ))

    def action_suggest_fix(self) -> None:
        if not self.work_file.exists():
            self.query_one("#diff-pane", RichLog).write("\n[dim]edit first — press  e[/]")
            return
        ref  = self.meta["solution"]
        user = self.work_file.read_text()
        self.app.push_screen(AIModal(
            f"suggest fix  ·  {self.problem.name}",
            lambda: get_suggest_fix(self.problem.name, ref, user),
        ))

    def action_explain(self) -> None:
        ref = self.meta["solution"]
        self.app.push_screen(AIModal(
            f"explain  ·  {self.problem.name}",
            lambda: get_explain(self.problem.name, ref),
        ))

    def action_chat(self) -> None:
        ref = self.meta["solution"]

        def _send(msg: str) -> tuple[str, str]:
            if msg.strip() == "/health":
                return opencode_chat_health()
            user = self.work_file.read_text() if self.work_file.exists() else ""
            reply, sid, status = opencode_chat(
                self.problem.name,
                ref,
                user,
                msg,
                self.chat_session_id,
            )
            self.chat_session_id = sid or self.chat_session_id
            return (reply, status)

        self.app.push_screen(
            ChatModal(
                f"chat  ·  {self.problem.name}",
                _send,
                context_fn=self._chat_context_line,
                mistakes_fn=self._chat_recent_mistakes,
                diff_fn=self._chat_diff_preview,
                open_hint_fn=self.action_hint,
                open_fix_fn=self.action_suggest_fix,
            )
        )

    def _chat_context_line(self) -> str:
        row = get_row(self.conn, self.pid)
        last_rating = RATING_LABELS.get(row["last_rating"], "—") if row and row["last_rating"] else "—"
        draft_state = "present" if self.work_file.exists() else "empty"
        diff_state = "ready" if self.has_diff else "none"
        return (
            f"[dim]context: attempts {self.attempts} | last rating {last_rating} | "
            f"draft {draft_state} | diff {diff_state}[/]"
        )

    def _chat_recent_mistakes(self) -> list[str]:
        return recent_mistakes(self.conn, self.pid, limit=2)

    def _chat_diff_preview(self) -> str:
        if not self.has_diff and not self.work_file.exists():
            return ""
        user_code = self.work_file.read_text() if self.work_file.exists() else ""
        ref_code = self.meta["solution"]
        raw = "\n".join(
            difflib.unified_diff(
                ref_code.splitlines(),
                user_code.splitlines(),
                fromfile="reference",
                tofile="yours",
                lineterm="",
            )
        )
        lines = raw.splitlines()
        if not lines:
            return "No diff: your draft currently matches the reference."
        return "\n".join(lines[:48])

    def action_edit(self) -> None:
        self.attempts += 1
        if not self.work_file.exists():
            self.work_file.write_text(f"# {self.problem.name}\n\n")
        with self.app.suspend():
            subprocess.run([EDITOR, str(self.work_file)])
        self.call_after_refresh(self._show_diff)

    def _show_diff(self) -> None:
        log = self.query_one("#diff-pane", RichLog)
        log.clear()

        user_code = self.work_file.read_text() if self.work_file.exists() else ""
        ref_code  = self.meta["solution"]

        if user_code.splitlines() == ref_code.splitlines():
            log.write("[white]✓  perfect match[/]")
            summary = "✓  perfect match"
        else:
            log.write(build_side_by_side(ref_code, user_code))
            summary = "\n".join(difflib.unified_diff(
                ref_code.splitlines(), user_code.splitlines(),
                fromfile="reference", tofile="yours", lineterm="",
            ))
            small_summary = "\n".join(summary.splitlines()[:16]).strip()
            if small_summary:
                log_mistake(self.conn, self.pid, small_summary)

        # Run tests if available
        if has_test_cases(self.problem) and user_code.strip():
            log.write("\n[dim]── running tests ──[/]")
            try:
                test_results = run_tests(self.problem, user_code)
                if test_results:
                    log.write(format_test_results(test_results))
                    # Log test failures as mistakes
                    for tr in test_results:
                        if not tr.passed:
                            log_mistake(self.conn, self.pid, f"test failed: {tr.name} — {tr.detail}")
                else:
                    log.write("[dim]no tests ran[/]")
            except Exception as e:
                log.write(f"[red]test runner error: {e}[/]")

        max_r = max_rating_for(self.attempts)
        if self.attempts >= 4:
            log.write(f"\n[bold red]attempt {self.attempts} — press  s  to record (forced: Again)[/]")
        else:
            log.write(
                f"\n[dim]attempt {self.attempts}  ·  max: {RATING_LABELS[max_r]}"
                f"  ·  s = submit    e = retry    h = hint    f = fix    c = chat    x = explain[/]"
            )

        self.conn.execute(
            "INSERT INTO reviews (problem_id, last_output) VALUES (?,?) "
            "ON CONFLICT(problem_id) DO UPDATE SET last_output=excluded.last_output",
            (self.pid, summary),
        )
        self.conn.commit()
        self.has_diff = True

    def action_submit(self) -> None:
        if not self.has_diff:
            self.query_one("#diff-pane", RichLog).write("\n[dim]edit first — press  e[/]")
            return
        max_r = max_rating_for(self.attempts)
        if max_r == 1:
            log = self.query_one("#diff-pane", RichLog)
            log.write("\n[bold red]forced: Again  (4+ attempts)[/]")
            sm2_update(self.conn, self.pid, 1)
            self.work_file.unlink(missing_ok=True)
            self.app.pop_screen()
        else:
            self.app.push_screen(RatingModal(max_r, self.attempts), self._rated)

    def _rated(self, rating: int | None) -> None:
        if rating:
            sm2_update(self.conn, self.pid, rating)
        self.work_file.unlink(missing_ok=True)
        self.app.pop_screen()

    def action_back(self) -> None:
        if self.work_file.exists():
            self.app.push_screen(
                ConfirmModal("Discard in-progress work and go back?"),
                self._confirm_back,
            )
        else:
            self.app.pop_screen()

    def _confirm_back(self, confirmed: bool | None) -> None:
        if confirmed:
            self.work_file.unlink(missing_ok=True)
            self.app.pop_screen()


# ── Search bar ────────────────────────────────────────────────────────────────
class SearchBar(Static):
    DEFAULT_CSS = """
    SearchBar { height: 1; padding: 0 2; background: $surface; }
    SearchBar Input {
        border: none; height: 1;
        background: $surface; color: $foreground; padding: 0;
    }
    """

    def compose(self) -> ComposeResult:
        yield Input(placeholder="search…", id="search-input")


# ── Menu Screen ───────────────────────────────────────────────────────────────
class MenuScreen(Screen):
    BINDINGS = [
        Binding("r",      "refresh",      "Refresh"),
        Binding("/",      "focus_search", "Search"),
        Binding("c",      "change_collection", "Collection"),
        Binding("g",      "generate_from_paper", "Generate"),
        Binding("i",      "import_problems", "Import"),
        Binding("escape", "clear_search", "Clear",  show=False),
        Binding("d",      "reset_row",    "Reset",  show=False),
        Binding("q",      "quit_app",     "Quit"),
    ]

    def __init__(self) -> None:
        super().__init__()
        self.conn        = get_db(DB_PATH)
        self._all_rows: list[tuple] = []
        self._visible_paths: list[Path] = []
        self._filter     = ""
        self.current_collection = PROBLEMS_DIR

    def compose(self) -> ComposeResult:
        yield Header()
        yield Static(id="stats-bar")
        yield SearchBar(id="search-bar")
        yield DataTable(id="table", cursor_type="row")
        yield Footer()

    def on_mount(self) -> None:
        t = self.query_one(DataTable)
        t.add_columns("Status", " ", "Problem", "Reps", "Interval", "Next review")
        self._refresh()

    def _refresh(self) -> None:
        problems = scan_problems(self.current_collection)
        due = new = upcoming = 0
        rows: list[tuple] = []

        for p in problems:
            # Determine ID based on root PROBLEMS_DIR
            pid          = get_problem_id(p, PROBLEMS_DIR)
            row          = get_row(self.conn, pid)
            label, color = status_label(row)
            reps         = str(row["reps"])      if row else "0"
            interval     = f"{row['interval']}d" if row else "—"
            nxt          = row["next_review"][:10] if row else "—"
            sort_key     = 0 if label in ("Due", "Due soon") else (1 if label == "New" else 2)
            if label in ("Due", "Due soon"): due      += 1
            elif label == "New":             new      += 1
            else:                            upcoming += 1
            rows.append((sort_key, label, color, p, reps, interval, nxt))

        self._all_rows = sorted(rows, key=lambda x: x[0])

        streak = get_streak(self.conn)
        streak_str = f"  [bold yellow]🔥 {streak}d streak[/]" if streak >= 2 else ""

        # Determine collection display name
        if self.current_collection == PROBLEMS_DIR:
            col_name = "Main"
        elif self.current_collection.parent == PROBLEMS_DIR:
            col_name = self.current_collection.name
        else:
            try:
                col_name = str(self.current_collection.relative_to(PROBLEMS_DIR))
            except ValueError:
                col_name = self.current_collection.name

        self.query_one("#stats-bar", Static).update(
            f"  [bold cyan]📂 {col_name}[/]  "
            f"[bold red]{due} due[/]  [bold]{new} new[/]"
            f"  [dim]{upcoming} upcoming  ·  {len(problems)} total[/]"
            + streak_str
        )
        self._render_table()

    def _render_table(self) -> None:
        t = self.query_one(DataTable)
        t.clear()
        self._visible_paths = []
        q = self._filter.lower()
        for _, label, color, p, reps, interval, nxt in self._all_rows:
            if q and q not in p.name.lower():
                continue
            self._visible_paths.append(p)
            badges = problem_badges(p)
            badge_str = "".join(b[0] for b in badges) if badges else ""
            t.add_row(
                f"[{color}]{label}[/]", badge_str, p.name, reps, interval, nxt,
                key=str(p),
            )

    def action_focus_search(self) -> None:
        self.query_one("#search-input", Input).focus()

    def action_change_collection(self) -> None:
        cols = scan_collections(PROBLEMS_DIR)
        self.app.push_screen(
            CollectionSelectModal(cols, self.current_collection),
            self._on_collection_selected
        )

    def _on_collection_selected(self, collection: Path | None) -> None:
        if collection:
            self.current_collection = collection
            self._refresh()

    def action_generate_from_paper(self) -> None:
        self.app.push_screen(PaperGenerateModal(), self._on_paper_config)

    def _on_paper_config(self, config: dict | None) -> None:
        if not config:
            return
        
        # Show generating status
        stats = self.query_one("#stats-bar", Static)
        original_text = str(stats.renderable)
        stats.update("[bold yellow]  Generating problems from paper...[/]")
        
        # Run generation in a thread to not block UI
        import threading
        threading.Thread(
            target=self._run_generation,
            args=(config,),
            daemon=True,
        ).start()

    def _run_generation(self, config: dict) -> None:
        from paper_generator import generate_problems
        
        output_dir = PROBLEMS_DIR / "generated"
        paper, files = generate_problems(
            arxiv_url=config["url"],
            output_dir=output_dir,
            num_problems=config.get("num_problems", 3),
            language=config.get("language", "python"),
            use_marimo=True,
        )
        
        def _done():
            if paper and files:
                self.query_one("#stats-bar", Static).update(
                    f"[bold green]  Generated {len(files)} problems from: {paper.title}[/]"
                )
                # Switch to generated collection
                self.current_collection = output_dir
                self._refresh()
            else:
                self.query_one("#stats-bar", Static).update(
                    "[bold red]  Failed to generate problems. Check the paper URL.[/]"
                )
        
        self.app.call_from_thread(_done)

    def action_import_problems(self) -> None:
        self.app.push_screen(ImportProblemModal(), self._on_import_result)

    def _on_import_result(self, result: dict | None) -> None:
        if result:
            self._refresh()

    def action_clear_search(self) -> None:
        inp = self.query_one("#search-input", Input)
        inp.value = ""
        self._filter = ""
        self._render_table()
        self.query_one(DataTable).focus()

    def on_input_changed(self, event: Input.Changed) -> None:
        self._filter = event.value
        self._render_table()

    def on_input_submitted(self, _: Input.Submitted) -> None:
        self.query_one(DataTable).focus()

    def on_data_table_row_selected(self, event: DataTable.RowSelected) -> None:
        self.app.push_screen(StudyScreen(Path(str(event.row_key.value))))

    def on_screen_resume(self) -> None:
        self._refresh()

    def action_refresh(self) -> None:
        self._refresh()

    def action_quit_app(self) -> None:
        self.app.exit()

    def action_reset_row(self) -> None:
        t = self.query_one(DataTable)
        if t.cursor_row is None:
            return
        if t.cursor_row < 0 or t.cursor_row >= len(self._visible_paths):
            return
        p_path = self._visible_paths[t.cursor_row]
        pid = get_problem_id(p_path, PROBLEMS_DIR)
        self.app.push_screen(
            ConfirmModal(f"Reset progress for  {p_path.name}?"),
            lambda confirmed: self._do_reset(confirmed, pid),
        )

    def _do_reset(self, confirmed: bool | None, pid: str) -> None:
        if confirmed:
            reset_progress(self.conn, pid)
            self._refresh()


# ── App ───────────────────────────────────────────────────────────────────────
class MLStudyApp(App):
    TITLE = "RECODE"
    CSS = """
    Screen        { background: $background; color: $foreground; }
    Header        { background: $surface; color: $foreground; }
    Footer        { background: $surface; color: $accent; }

    #stats-bar    { height: 1; padding: 0 2; background: $surface; color: $accent; }
    #table        { height: 1fr; border: solid $primary; }

    #problem-bar  { height: 2; padding: 0 2; background: $surface; color: $accent; }
    .full-pane    { height: 1fr; border: solid $primary; padding: 1 2; overflow-y: auto; }

    #modal-box {
        background: $surface;
        border: double $primary;
        padding: 2 4;
        width: 56;
        height: 15;
        align: center middle;
    }
    #modal-title { text-style: bold; margin-bottom: 1; }
    #modal-skip  { color: $accent; }
    RatingModal  { align: center middle; background: $background 70%; }
    ConfirmModal { align: center middle; background: $background 70%; }

    #hint-box {
        background: $surface;
        border: double $primary;
        padding: 2 4;
        width: 80%;
        height: 60%;
        align: center middle;
    }
    #hint-title { text-style: bold; margin-bottom: 1; }
    #hint-md    { height: 1fr; overflow-y: auto; background: $surface; }
    AIModal     { align: center middle; background: $background 70%; }

    /* Collection Select Modal */
    #collection-box {
        background: $surface;
        border: double $primary;
        padding: 2 4;
        width: 60;
        height: 20;
        align: center middle;
    }
    #collection-list {
        height: 1fr;
        border: solid $accent;
    }
    CollectionSelectModal { align: center middle; background: $background 70%; }

    #chat-box {
        background: $surface;
        border: double $primary;
        padding: 1 2;
        width: 92%;
        height: 82%;
        align: center middle;
    }
    #chat-header {
        height: 1;
        margin-bottom: 1;
    }
    #chat-title {
        width: 1fr;
        text-style: bold;
    }
    #chat-status {
        width: auto;
        content-align: right middle;
    }
    #chat-context {
        height: 1;
        margin-bottom: 1;
    }
    #chat-main {
        height: 1fr;
        margin-bottom: 1;
    }
    #chat-log {
        width: 2fr;
        border: round $primary;
        background: $background;
        padding: 1;
    }
    #chat-diff {
        width: 1fr;
        border: round $accent;
        background: $surface;
        padding: 1;
        margin-left: 1;
        overflow-y: auto;
    }
    .hidden { display: none; }
    #chat-input {
        dock: bottom;
        margin-bottom: 1;
    }
    #chat-help {
        height: 1;
        color: $accent;
    }
    ChatModal   { align: center middle; background: $background 70%; }
    """

    def on_mount(self) -> None:
        for theme in TERMINAL_SEXY_THEMES:
            self.register_theme(theme)
        self.push_screen(MenuScreen())


def main(
    *,
    problems_dir: str | Path | None = None,
    db_path: str | Path | None = None,
    editor: str | None = None,
) -> None:
    configure_runtime(
        prepare_runtime(problems_dir=problems_dir, db_path=db_path, editor=editor)
    )
    MLStudyApp().run()


if __name__ == "__main__":
    main()