from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path

from src.agents.annotation_agent import OfflineAnnotationAgent
from src.agents.editing_agent import OfflinePDFEditingAgent
from src.agents.export_agent import OfflineExportAgent
from src.agents.ocr_agent import OfflineOCRAgent
from src.agents.structure_inspector import OfflinePDFStructureInspector
from src.agents.validation_agent import OfflineValidationAgent
from src.core.models import InspectionReport, JobConfig, PDFOperation


@dataclass
class OperationBuckets:
    edit: list[PDFOperation]
    annotate: list[PDFOperation]


class PDFToolbox:
    """Shared PDF tools used by orchestrators and pipeline definitions."""

    EDIT_OPERATION_TYPES = {
        "insert_text",
        "whiteout_region",
        "rotate_page",
        "delete_page",
        "reorder_pages",
    }
    ANNOTATION_OPERATION_TYPES = {"highlight", "comment", "rectangle", "redact"}

    def __init__(
        self,
        inspector: OfflinePDFStructureInspector | None = None,
        ocr: OfflineOCRAgent | None = None,
        editor: OfflinePDFEditingAgent | None = None,
        annotator: OfflineAnnotationAgent | None = None,
        validator: OfflineValidationAgent | None = None,
        exporter: OfflineExportAgent | None = None,
    ):
        self.inspector = inspector or OfflinePDFStructureInspector()
        self.ocr_agent = ocr or OfflineOCRAgent()
        self.editing_agent = editor or OfflinePDFEditingAgent()
        self.annotation_agent = annotator or OfflineAnnotationAgent()
        self.validation_agent = validator or OfflineValidationAgent()
        self.export_agent = exporter or OfflineExportAgent()

    def inspect(self, pdf_path: Path) -> InspectionReport:
        return self.inspector.inspect(pdf_path)

    def run_ocr(self, input_pdf: Path, output_pdf: Path, language: str = "eng") -> Path:
        return self.ocr_agent.run(input_pdf, output_pdf, language)

    def edit_pdf(
        self, input_pdf: Path, output_pdf: Path, operations: list[PDFOperation]
    ) -> Path:
        return self.editing_agent.run(input_pdf, output_pdf, operations)

    def annotate_pdf(
        self, input_pdf: Path, output_pdf: Path, operations: list[PDFOperation]
    ) -> Path:
        return self.annotation_agent.run(input_pdf, output_pdf, operations)

    def validate_pdf(self, input_pdf: Path, output_pdf: Path) -> tuple[Path, dict]:
        return self.validation_agent.run(input_pdf, output_pdf)

    def export_job(
        self, final_pdf: Path, job_config: JobConfig, stages: list, log_file: Path
    ) -> tuple[Path, Path]:
        return self.export_agent.run(final_pdf, job_config, stages, log_file)

    def bucket_operations(self, operations: list[PDFOperation]) -> OperationBuckets:
        edit_ops = [op for op in operations if op.type in self.EDIT_OPERATION_TYPES]
        annotate_ops = [
            op for op in operations if op.type in self.ANNOTATION_OPERATION_TYPES
        ]
        return OperationBuckets(edit=edit_ops, annotate=annotate_ops)


def build_default_pdf_toolbox() -> PDFToolbox:
    return PDFToolbox()