cowork-local/presentation/folder/ai_edit_runner.py

"""Một lượt AI sửa file, từ lúc gửi tới lúc ghi ra đĩa — R08-T12.

``_ai_run_edit`` dài (82 dòng) vì nó là cả một lượt: dựng ngữ cảnh từ
file đang mở, gọi provider, nhận nội dung phát dần, tách phần mã khỏi
phần giải thích, rồi dựng bản xem trước.

Không bao giờ ghi đè thẳng: kết quả hiện ra để người dùng xem, và chỉ
``_ai_apply`` mới chạm vào file.

Cùng kiểu mixin, xem ghi chú ở đầu ``presentation/shell/nav_rail.py``.
"""
from __future__ import annotations

from .file_helpers import (
    _HTML_SUFFIXES, _PPTX_SUFFIXES, _parse_ai_output, _pptx_available,
)

import os
from pathlib import Path
from typing import Optional
from PySide6.QtCore import Qt
from ...core.worker import AgentWorker
from ...i18n import tr
from ...theme import current_palette


class AiEditRunnerMixin:
    """Trộn vào FolderTab."""

    def _ai_send(self) -> None:
        if not self._root or not os.path.isdir(self._root):
            self.ai_chat.add_error(tr("folder.ai_no_file"))
            return
        instruction = self.ai_input.text().strip()
        if not instruction:
            return
        self.ai_input.clear()
        self.ai_chat.add_user(instruction)
        # QUEUE: while a run is active OR a proposal is awaiting Apply/Discard,
        # hold the new instruction and run it when the pipeline goes idle. Lets
        # the user line up several edits without waiting for each to finish.
        if self._ai_worker is not None or self._ai_pending is not None:
            self._ai_queue.append(instruction)
            self.ai_chat.add_status(tr("folder.ai_queued", n=len(self._ai_queue)))
            self._update_queue_status()
            return
        self._ai_start(instruction)

    def _ai_start(self, instruction: str) -> None:
        """Begin processing one instruction (plan → edit). Assumes the pipeline
        is idle (the queue calls this when the previous run finishes)."""
        # If a text/code/HTML file is open (even in Preview), switch it into the
        # editor so AI can edit it. If nothing editable is open, that's fine —
        # the request may be to CREATE a new file (the model names it via FILE:).
        editable = self.stack.currentWidget() is self.editor
        if not editable:
            editable = self._ensure_editor_for_ai()
        self._maybe_suggest_image_model(instruction)
        # Auto Model Routing (may switch to the best coding model for this run).
        self._ai_apply_routing(instruction)
        has_file = editable and bool(self._current_file)
        self._ai_running_file = Path(self._current_file).name if has_file else tr("folder.ai_new_file")
        self._ai_set_busy(True)
        # Announce start on the status bar so it's visible even from another tab —
        # the edit keeps running in the background until it finishes.
        self.status_message.emit(tr("folder.ai_running", name=self._ai_running_file))
        # Two phases so the PLAN is shown INLINE *before* the edit runs.
        self._ai_ctx = {
            "filename": Path(self._current_file).name if has_file else "",
            "content": self.editor.toPlainText() if has_file else "",
            "convo": self._cowork_context(),
            "instruction": instruction,
            "provider": self._ai_provider(),
            "plan": "",
        }
        # Reset the token/cost tally for THIS prompt (plan + edit calls sum into it).
        self._ai_prompt_usage = {"in": 0, "out": 0, "cache": 0, "cost": 0.0}
        self._ai_run_plan()

    def _ai_maybe_dequeue(self) -> None:
        """When the pipeline is fully idle, start the next queued instruction."""
        if self._ai_worker is not None or self._ai_pending is not None:
            return
        if not self._ai_queue:
            return
        nxt = self._ai_queue.pop(0)
        self._update_queue_status()
        self._ai_start(nxt)

    def _ai_add_usage(self, usage) -> None:
        """Add one model call's usage (plan or edit) to THIS prompt's tally."""
        if not isinstance(usage, dict):
            return
        tot = getattr(self, "_ai_prompt_usage", None)
        if tot is None:
            tot = self._ai_prompt_usage = {"in": 0, "out": 0, "cache": 0, "cost": 0.0}
        tot["in"] += int(usage.get("in", 0) or 0)
        tot["out"] += int(usage.get("out", 0) or 0)
        tot["cache"] += int(usage.get("cache", 0) or 0)
        tot["cost"] += float(usage.get("cost_usd", 0.0) or 0.0)

    def _ai_show_usage(self, bubble) -> None:
        """Footer under the AI-edit reply: ↓in ↑out ▤ctx $cost for the whole
        prompt (plan + edit), priced in the display currency — same as Cowork."""
        tot = getattr(self, "_ai_prompt_usage", None)
        if bubble is None or not tot or not (tot["in"] or tot["out"]):
            return
        from ...core import model_pricing as mp, usage_tracker as ut
        pricing = {**ut.DEFAULT_PRICING, **(self.ctx.config.data.get("usage") or {})}
        line = (f"↓{mp.format_tokens(int(tot['in']))} ↑{mp.format_tokens(int(tot['out']))} "
                f"▤{mp.format_tokens(int(tot['in'] + tot['out'] + tot['cache']))} "
                f"{ut.format_cost(tot['cost'], pricing)}")
        try:
            bubble.add_usage(line)
        except Exception:  # noqa: BLE001 - a usage footer must never break the edit
            pass

    def _ai_run_plan(self) -> None:
        c = self._ai_ctx
        plan_bubble = self.ai_chat.add_plan(tr("folder.ai_planning"))
        self.ai_chat.scroll_to_bottom()

        def job(worker):
            from ...core import usage_tracker as ut
            from ...core.co4e_runner import _usage_delta
            provider = c["provider"]
            messages = [{"role": "system", "content":
                         "You are an AI file editor. Give a SHORT numbered plan (2-4 steps) for "
                         "the requested change. Plan ONLY — do NOT output any code."}]
            if c["convo"]:
                messages.append({"role": "system",
                                 "content": "Context from the user's Cowork conversation:\n" + c["convo"]})
            messages.append({"role": "user", "content":
                             f"File: {c['filename']}\n\nCurrent content:\n```\n{c['content']}\n```\n\n"
                             f"Request: {c['instruction']}"})
            ut.set_context("folder", c.get("filename") or "AI edit")   # Dashboard + usage
            ut.begin_accumulation(); base = ut.accumulated()
            try:
                r = provider.chat(messages, tools=None, cancel=worker.is_cancelled)
                txt = r.get("content", "") if isinstance(r, dict) else str(r)
                usage = _usage_delta(base, self.ctx.config)
            finally:
                ut.end_accumulation()
            return {"plan": provider.strip_think(txt) or "", "usage": usage}

        worker = AgentWorker(job)
        worker.finished_ok.connect(lambda res, b=plan_bubble: self._ai_plan_done(res, b))
        worker.failed.connect(lambda err, b=plan_bubble: self._ai_failed(err, b))
        self._ai_worker = worker
        worker.start()

    def _ai_plan_done(self, result, plan_bubble) -> None:
        self._ai_add_usage((result or {}).get("usage"))    # plan-step tokens
        plan = ((result or {}).get("plan") or "").strip()
        self._ai_ctx["plan"] = plan
        plan_bubble.set_plain(plan or tr("folder.ai_empty"))
        self.ai_chat.scroll_to_bottom()
        self._ai_run_edit()   # now execute the plan

    def _ai_run_edit(self) -> None:
        c = self._ai_ctx
        bubble = self.ai_chat.add_assistant(tr("folder.ai_edit"))
        self.ai_chat.scroll_to_bottom()

        pptx_note = ("\nThis is a PPTX shown as marker blocks '### Slide N / Box M', where N is the "
                     "1-based SLIDE NUMBER and M the box on that slide. When the user refers to a "
                     "slide by number (e.g. 'edit slide 3'), change ONLY the blocks under '### Slide "
                     "3' and leave every other slide's block exactly as-is. Each block has fields "
                     "type/pos/size/font/text (and image for pictures). To change TEXT COLOUR or "
                     "FONT edit the `font:` line — e.g. `font: name=Arial size=28 bold=1 "
                     "color=FF0000`. Keep all block markers and structure.") if self._edit_kind == "pptx" else ""

        # When creating a NEW deck (request mentions slides/pptx and we're not
        # already editing one), tell the model the marker format to emit so we can
        # build a real .pptx from it.
        _pptx_words = ("pptx", "powerpoint", "slide", "presentation", "deck",
                       "スライド", "プレゼン", "trình chiếu", "trinh chieu", "bài thuyết trình")
        wants_new_pptx = (self._edit_kind != "pptx"
                          and any(w in c["instruction"].lower() for w in _pptx_words))
        new_pptx_note = ("\nTo CREATE a PowerPoint, name it `FILE: <name>.pptx` and output the slides "
                         "as marker blocks — one block per shape:\n"
                         "### Slide 1 / Box 1\ntype: text\npos: 0.5, 0.4\nsize: 9.0, 1.2\n"
                         "font: name=Calibri size=32 bold=1 color=1F3864\ntext:\nTitle here\n\n"
                         "### Slide 1 / Box 2\ntype: text\npos: 0.5, 1.8\nsize: 9.0, 4.5\n"
                         "text:\nBullet one\nBullet two\n\n"
                         "Increment the Slide number for each new slide; pos/size are in inches; "
                         "font color is RRGGBB hex.") if wants_new_pptx else ""

        imggen_note = ""
        try:
            from ...core import image_gen
            if image_gen.is_configured(self.ctx.config):
                imggen_note = ("\nYou can also GENERATE an illustration image: add a line "
                               "`IMAGE_GEN: <describe the image> => <relative/path.png>`. Use a "
                               "generated image e.g. as a new picture, or (for pptx) set a picture "
                               "box's `image:` field to that same path to insert it.")
        except Exception:  # noqa: BLE001
            pass

        def job(worker):
            provider = c["provider"]
            open_note = (f"the currently-open file '{c['filename']}'" if c["filename"]
                         else "no file is open")
            messages = [{"role": "system", "content":
                         "You are an AI file editor inside an app. Following the plan, output the "
                         "COMPLETE file content in ONE fenced code block (```), and nothing after "
                         "it. Preserve everything you were not asked to change.\n"
                         "If the request is to CREATE A NEW file (or a different file than the one "
                         "open), put a line `FILE: <relative/path/name.ext>` (relative to the "
                         "current folder) immediately before the code block. Omit FILE to edit the "
                         f"open file. Right now, {open_note}." + pptx_note + new_pptx_note + imggen_note}]
            if c["convo"]:
                messages.append({"role": "system",
                                 "content": "Context from the user's Cowork conversation:\n" + c["convo"]})
            if c["plan"]:
                messages.append({"role": "system", "content": "Plan to follow:\n" + c["plan"]})
            cur = (f"File: {c['filename']}\n\nCurrent content:\n```\n{c['content']}\n```\n\n"
                   if c["filename"] else "No file is currently open.\n\n")
            messages.append({"role": "user", "content": cur + f"Request: {c['instruction']}"})

            def on_text(piece: str) -> None:
                worker.emit_event({"type": "text", "delta": piece})

            from ...core import usage_tracker as ut
            from ...core.co4e_runner import _usage_delta
            ut.set_context("folder", c.get("filename") or "AI edit")   # Dashboard + usage
            ut.begin_accumulation(); base = ut.accumulated()
            try:
                r = provider.chat(messages, tools=None, on_text=on_text, cancel=worker.is_cancelled)
                txt = r.get("content", "") if isinstance(r, dict) else str(r)
                usage = _usage_delta(base, self.ctx.config)
            finally:
                ut.end_accumulation()
            return {"text": provider.strip_think(txt) or "", "usage": usage}

        worker = AgentWorker(job)
        worker.event.connect(lambda ev, b=bubble: self._ai_stream(ev, b))
        worker.finished_ok.connect(lambda res, b=bubble: self._ai_done(res, b))
        worker.failed.connect(lambda err, b=bubble: self._ai_failed(err, b))
        self._ai_worker = worker
        worker.start()

    def _ai_stream(self, ev, bubble) -> None:
        if isinstance(ev, dict) and ev.get("type") == "text":
            bubble.append_delta(ev.get("delta", ""))
            self.ai_chat.scroll_to_bottom()

    def _ai_done(self, result, bubble) -> None:
        self._ai_worker = None
        self._ai_set_busy(False)
        self._ai_add_usage((result or {}).get("usage"))   # edit-step tokens
        self._ai_show_usage(bubble)                        # footer: prompt total (plan+edit)
        text = ((result or {}).get("text") or "").strip()
        target, new_content, summary, image_gens = _parse_ai_output(text)
        if new_content is None and not image_gens:
            bubble.set_markdown(text or tr("folder.ai_empty"))
            self.ai_chat.scroll_to_bottom()
            self._ai_flag_done()
            return
        # Decide edit-current vs create-new. A FILE: naming a path different from
        # the open file (or when nothing is open) → CREATE a new file.
        create = bool(target) and (not self._current_file
                                   or Path(target).name != Path(self._current_file).name)
        # PROPOSE the change — nothing is written until the user clicks Apply.
        self._ai_pending = {"content": new_content,
                            "target": target if create else None,
                            "image_gens": image_gens}
        hint = tr("folder.ai_review_hint")
        bubble.set_markdown((summary + "\n\n" if summary else "") + "_" + hint + "_")
        if new_content is not None:
            import difflib
            old = "" if create else self.editor.toPlainText()
            diff = "".join(difflib.unified_diff(
                old.splitlines(keepends=True), new_content.splitlines(keepends=True),
                fromfile=("(new file)" if create else "current"),
                tofile=(target if create else "proposed"))) or "(no textual difference)"
            title = tr("folder.ai_proposed_new", name=target) if create else tr("folder.ai_proposed")
            self.ai_chat.add_diff(title, diff)
        if image_gens:
            listing = "\n".join(f"• {p}  ({prompt[:60]})" for prompt, p in image_gens)
            self.ai_chat.add_status(tr("folder.ai_image_plan") + "\n" + listing)
        self._ai_confirm_row.setVisible(True)
        self.ai_chat.scroll_to_bottom()
        name = target if create else getattr(self, "_ai_running_file", "")
        self.status_message.emit(tr("folder.ai_proposed_status", name=name))
        self._ai_status.setText("● " + hint)
        self._ai_status.setStyleSheet(f"color:{current_palette().warning};")

    def _ai_apply(self) -> None:
        """Confirmed by the user. If the edit GENERATES images, ask the image
        gate then generate them (off-thread) before finalising the file edit."""
        if not self._ai_pending:
            return
        p = self._ai_pending
        self._ai_pending = None
        self._ai_confirm_row.setVisible(False)
        if p.get("image_gens"):
            from PySide6.QtWidgets import QMessageBox
            if QMessageBox.question(self, tr("folder.ai_image_confirm_title"),
                                    tr("folder.ai_image_confirm_gen")) != QMessageBox.Yes:
                self.status_message.emit(tr("folder.ai_image_declined"))
                return
            self._ai_generate_then_finalize(p)
            return
        self._ai_finalize_apply(p)





    def _ai_discard(self) -> None:
        self._ai_pending = None
        self._ai_confirm_row.setVisible(False)
        self.ai_chat.add_status(tr("folder.ai_discarded"))
        self.ai_chat.scroll_to_bottom()
        self._ai_status.setText("")
        self._ai_maybe_dequeue()   # discarding resolves the gate → run the next queued edit


    def _ai_failed(self, err, bubble) -> None:
        self._ai_worker = None
        bubble.set_markdown(tr("folder.ai_error", err=err))
        self._ai_set_busy(False)
        self.status_message.emit(tr("folder.ai_error", err=err))
        self._ai_flag_done()