from __future__ import annotations
from pathlib import Path
import fitz
from pypdf import PdfReader, PdfWriter
from src.core.models import PDFOperation

class OfflinePDFEditingAgent:
    name = "editing_agent"

    def run(self, input_pdf: Path, output_pdf: Path, operations: list[PDFOperation]) -> Path:
        doc = fitz.open(input_pdf)
        page_ops = [op for op in operations if op.type in ("insert_text","whiteout_region","rotate_page")]
        delete_ops = [op for op in operations if op.type == "delete_page"]
        reorder_ops = [op for op in operations if op.type == "reorder_pages"]

        for op in page_ops:
            if op.page is None or op.page >= len(doc): continue
            page = doc[op.page]
            if op.type == "insert_text" and op.position and op.text:
                color = tuple(op.color) if op.color else (0, 0, 0)
                page.insert_text(op.position, op.text, fontsize=12, color=color, overlay=True)
            elif op.type == "whiteout_region" and op.rect:
                page.draw_rect(fitz.Rect(op.rect), color=(1,1,1), fill=(1,1,1), overlay=True)
                if op.text:
                    x0, y0 = op.rect[0], op.rect[1]
                    page.insert_text((x0+4, y0+14), op.text, fontsize=12, color=(0,0,0), overlay=True)
            elif op.type == "rotate_page" and op.rotation is not None:
                page.set_rotation(op.rotation)

        stage = output_pdf.with_name(output_pdf.stem + "-pre.pdf")
        stage.parent.mkdir(parents=True, exist_ok=True)
        doc.save(stage, garbage=4, deflate=True, clean=True)
        doc.close()

        reader = PdfReader(str(stage))
        writer = PdfWriter()
        delete_indices = {op.page for op in delete_ops if op.page is not None}
        order = reorder_ops[-1].pages if reorder_ops else None
        indices = [i for i in range(len(reader.pages)) if i not in delete_indices]
        if order:
            indices = [i for i in order if i not in delete_indices]
        for idx in indices:
            writer.add_page(reader.pages[idx])
        output_pdf.parent.mkdir(parents=True, exist_ok=True)
        with output_pdf.open("wb") as f:
            writer.write(f)
        stage.unlink(missing_ok=True)
        return output_pdf