cowork-local/presentation/folder/ai_output_writer.py

"""Ghi kết quả AI ra đĩa — R08-T12.

Tách khỏi ``ai_edit_runner.py`` vì đây là phần DUY NHẤT thật sự chạm vào
file của người dùng. Mọi thứ trước nó chỉ dựng bản xem trước.

Gồm cả nhánh sinh ảnh: lượt nào có ảnh thì phải chờ ảnh xong mới ghi, vì
nội dung có thể tham chiếu tới đường dẫn ảnh vừa tạo.
"""
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 AiOutputWriterMixin:
    """Trộn vào FolderTab."""

    def _ai_generate_then_finalize(self, p: dict) -> None:
        imgs = p.get("image_gens") or []
        root = os.path.normpath(self._root)
        img_model, img_base, img_key = self._ai_image_model()   # may target another provider
        self._ai_set_busy(True)
        self.status_message.emit(tr("folder.ai_generating"))

        def job(worker):
            from ...core import image_gen
            results = []
            for prompt, rel in imgs:
                dest = rel if os.path.isabs(rel) else os.path.join(root, rel)
                dest = os.path.normpath(dest)
                if os.path.commonpath([dest, root]) != root:
                    results.append((rel, False, "path escapes the folder"))
                    continue
                try:
                    os.makedirs(os.path.dirname(dest) or root, exist_ok=True)
                except OSError as exc:
                    results.append((rel, False, str(exc)))
                    continue
                ok, msg = image_gen.generate_image(self.ctx.config, prompt, dest,
                                                   model=img_model, base_url=img_base, api_key=img_key)
                results.append((dest, ok, msg))
            return {"results": results}

        worker = AgentWorker(job)
        worker.finished_ok.connect(lambda res, pp=p: self._ai_images_done(res, pp))
        worker.failed.connect(lambda err, pp=p: self._ai_images_done({"results": [], "err": err}, pp))
        self._ai_worker = worker
        worker.start()

    def _ai_images_done(self, res: dict, p: dict) -> None:
        self._ai_worker = None
        self._ai_set_busy(False)
        created = []
        for dest, ok, msg in res.get("results", []):
            if ok:
                created.append(dest)
                self.ai_chat.add_success("✓ " + tr("folder.ai_image_created", name=Path(dest).name))
            else:
                self.ai_chat.add_error(tr("folder.ai_image_failed", err=msg))
        # Now apply any text/file edit (pptx image: fields now point at real files).
        self._ai_finalize_apply(p, images_done=True)
        # If it was only image generation, open the first new image.
        if p.get("content") is None and not p.get("target") and created:
            self.open_file(created[0], reset=False)

    def _ai_finalize_apply(self, p: dict, images_done: bool = False) -> None:
        content = p.get("content")
        target = p.get("target")
        if content is None:
            self.ai_chat.scroll_to_bottom()
            self._ai_flag_done()
            self.status_message.emit(tr("folder.ai_done", name=getattr(self, "_ai_running_file", "")))
            return
        if target:
            dest = self._create_new_file(target, content)
            if dest is None:
                return
            self.ai_chat.add_success("✓ " + tr("folder.ai_created", name=Path(dest).name))
            self.status_message.emit(tr("folder.ai_created", name=Path(dest).name))
        else:
            self.editor.setPlainText(content)     # live update in the editor/preview
            self._ai_write_out(content, skip_image_confirm=images_done)
            self.ai_chat.add_success("✓ " + tr("folder.ai_applied"))
            self.status_message.emit(tr("folder.ai_done", name=getattr(self, "_ai_running_file", "")))
        self.ai_chat.scroll_to_bottom()
        self._ai_flag_done()

    def _create_new_file(self, target: str, content: str) -> Optional[str]:
        """Create ``target`` (relative to the folder root) with ``content`` and
        open it — like Cowork's save_file. Refuses paths escaping the root."""
        root = os.path.normpath(self._root)
        dest = target if os.path.isabs(target) else os.path.join(root, target)
        dest = os.path.normpath(dest)
        if os.path.commonpath([dest, root]) != root:
            self.status_message.emit(tr("folder.ai_error", err="path escapes the folder"))
            return None
        try:
            os.makedirs(os.path.dirname(dest) or root, exist_ok=True)
            if Path(dest).suffix.lower() in _PPTX_SUFFIXES and _pptx_available():
                # A .pptx is a binary package — build a real deck from the marker
                # text (writing text straight to .pptx would corrupt it).
                from ...core import pptx_edit
                pptx_edit.create_pptx_from_text(dest, content)
            else:
                Path(dest).write_text(content, encoding="utf-8")
        except Exception as exc:  # noqa: BLE001 - OS error or pptx build failure
            self.status_message.emit(tr("folder.save_error", err=str(exc)))
            return None
        self.open_file(dest, reset=False)   # show the new file; keep this AI chat
        return dest

    def _ai_write_out(self, content: str, skip_image_confirm: bool = False) -> None:
        """Persist the confirmed content to disk AND refresh the preview.
        pptx text is written back into the deck (no PowerPoint window)."""
        if not self._current_file:
            return
        try:
            if self._edit_kind == "pptx":
                if not self._write_pptx(content, skip_confirm=skip_image_confirm):
                    return
            else:
                Path(self._current_file).write_text(content, encoding="utf-8")
        except Exception as exc:  # noqa: BLE001
            self.status_message.emit(tr("folder.save_error", err=str(exc)))
            return
        # Refresh preview: HTML re-renders; pptx re-renders the slides; code stays
        # in the (now-saved) editor.
        suffix = Path(self._current_file).suffix.lower()
        if suffix in _HTML_SUFFIXES:
            self._show_html(self._current_file, mode_preview=True)
        elif suffix in _PPTX_SUFFIXES:
            self._show_pptx(self._current_file, mode_preview=True)