src/core/pipelines.py

from __future__ import annotations

from dataclasses import dataclass, field
from pathlib import Path
from typing import Callable

from src.core.models import JobConfig, JobResult, dataclass_to_dict
from src.core.toolbox import PDFToolbox
from src.services.local_logger import LocalLogger
from src.utils.file_utils import ensure_dir, stage_file


@dataclass
class PipelineContext:
    job_config: JobConfig
    job_dir: Path
    log_file: Path
    logger: LocalLogger
    toolbox: PDFToolbox
    current_pdf: Path
    stages: list = field(default_factory=list)
    errors: list = field(default_factory=list)
    state: dict = field(default_factory=dict)

    def make_stage_path(self, prefix: str, suffix: str = ".pdf") -> Path:
        return stage_file(self.job_dir, prefix, suffix=suffix)

    def finish_success(self, final_pdf: Path, manifest: Path) -> JobResult:
        self.logger.log("orchestrator", "info", "Job finished", success=True)
        return JobResult(
            self.job_config.job_id,
            True,
            final_pdf,
            self.log_file,
            manifest,
            self.stages,
            self.errors,
        )

    def finish_failure(self) -> JobResult:
        return JobResult(
            self.job_config.job_id,
            False,
            None,
            self.log_file,
            self.job_dir / "manifest.json",
            self.stages,
            self.errors,
        )


@dataclass
class PipelineStep:
    name: str
    action: Callable[[PipelineContext], None]
    enabled: Callable[[PipelineContext], bool] | None = None

    def should_run(self, context: PipelineContext) -> bool:
        if self.enabled is None:
            return True
        return self.enabled(context)


@dataclass
class PDFPipeline:
    name: str
    description: str
    steps: list[PipelineStep]

    def run(self, context: PipelineContext) -> PipelineContext:
        context.logger.log(
            "orchestrator",
            "info",
            "Pipeline started",
            pipeline=self.name,
            job_id=context.job_config.job_id,
        )
        for step in self.steps:
            if step.should_run(context):
                step.action(context)
        return context


def create_pipeline_context(
    job_config: JobConfig, toolbox: PDFToolbox, work_dir: Path
) -> PipelineContext:
    job_config.output_dir = ensure_dir(job_config.output_dir)
    job_dir = ensure_dir(job_config.output_dir / f"job-{job_config.job_id}")
    log_file = job_dir / "pipeline.jsonl"
    logger = LocalLogger(log_file)
    logger.log("orchestrator", "info", "Job started", job_id=job_config.job_id)
    return PipelineContext(
        job_config=job_config,
        job_dir=job_dir,
        log_file=log_file,
        logger=logger,
        toolbox=toolbox,
        current_pdf=job_config.input_pdf,
        state={"work_dir": str(work_dir)},
    )


def _inspect_step(context: PipelineContext) -> None:
    report = context.toolbox.inspect(context.current_pdf)
    buckets = context.toolbox.bucket_operations(context.job_config.operations)
    context.state["inspection_report"] = report
    context.state["operation_buckets"] = buckets
    context.stages.append({"stage": "inspect", "report": dataclass_to_dict(report)})
    context.logger.log(
        "inspect",
        "info",
        "Done",
        ocr_pages=report.ocr_required_pages,
        edit_operations=len(buckets.edit),
        annotation_operations=len(buckets.annotate),
    )


def _needs_ocr(context: PipelineContext) -> bool:
    report = context.state.get("inspection_report")
    if report is None:
        return context.job_config.force_ocr or context.job_config.run_ocr
    return (
        context.job_config.force_ocr
        or context.job_config.run_ocr
        or bool(report.ocr_required_pages)
    )


def _ocr_step(context: PipelineContext) -> None:
    output = context.make_stage_path("02-ocr")
    context.current_pdf = context.toolbox.run_ocr(
        context.current_pdf, output, context.job_config.ocr_language
    )
    context.stages.append({"stage": "ocr", "output": str(context.current_pdf)})
    context.logger.log("ocr", "info", "Done", output=str(context.current_pdf))


def _has_edit_ops(context: PipelineContext) -> bool:
    buckets = context.state.get("operation_buckets")
    return bool(buckets and buckets.edit)


def _edit_step(context: PipelineContext) -> None:
    buckets = context.state["operation_buckets"]
    output = context.make_stage_path("03-edit")
    context.current_pdf = context.toolbox.edit_pdf(
        context.current_pdf, output, buckets.edit
    )
    context.stages.append(
        {
            "stage": "edit",
            "output": str(context.current_pdf),
            "count": len(buckets.edit),
        }
    )
    context.logger.log("edit", "info", "Done", count=len(buckets.edit))


def _has_annotation_ops(context: PipelineContext) -> bool:
    buckets = context.state.get("operation_buckets")
    return bool(buckets and buckets.annotate)


def _annotate_step(context: PipelineContext) -> None:
    buckets = context.state["operation_buckets"]
    output = context.make_stage_path("04-annotate")
    context.current_pdf = context.toolbox.annotate_pdf(
        context.current_pdf, output, buckets.annotate
    )
    context.stages.append(
        {
            "stage": "annotate",
            "output": str(context.current_pdf),
            "count": len(buckets.annotate),
        }
    )
    context.logger.log("annotate", "info", "Done", count=len(buckets.annotate))


def _validate_step(context: PipelineContext) -> None:
    output = context.make_stage_path("05-validated")
    context.current_pdf, report = context.toolbox.validate_pdf(
        context.current_pdf, output
    )
    context.state["validation_report"] = report
    context.stages.append(
        {"stage": "validate", "output": str(context.current_pdf), "report": report}
    )
    context.logger.log("validate", "info", "Done", **report)


def _export_step(context: PipelineContext) -> None:
    final_pdf, manifest = context.toolbox.export_job(
        context.current_pdf,
        context.job_config,
        context.stages,
        context.log_file,
    )
    context.state["final_pdf"] = final_pdf
    context.state["manifest"] = manifest
    context.stages.append(
        {"stage": "export", "output": str(final_pdf), "manifest": str(manifest)}
    )
    context.logger.log("export", "info", "Done", final=str(final_pdf))


def build_pipeline_catalog(toolbox: PDFToolbox) -> dict[str, PDFPipeline]:
    inspect = PipelineStep("inspect", _inspect_step)
    ocr = PipelineStep("ocr", _ocr_step, enabled=_needs_ocr)
    edit = PipelineStep("edit", _edit_step, enabled=_has_edit_ops)
    annotate = PipelineStep("annotate", _annotate_step, enabled=_has_annotation_ops)
    validate = PipelineStep("validate", _validate_step)
    export = PipelineStep("export", _export_step)

    return {
        "full_document": PDFPipeline(
            name="full_document",
            description="Inspects the document, runs OCR if needed, applies edits and annotations, validates, then exports.",
            steps=[inspect, ocr, edit, annotate, validate, export],
        ),
        "text_edit": PDFPipeline(
            name="text_edit",
            description="Runs the edit-oriented document pipeline without annotation steps.",
            steps=[inspect, ocr, edit, validate, export],
        ),
        "annotation_review": PDFPipeline(
            name="annotation_review",
            description="Applies annotation and review operations on top of the current PDF, then validates and exports.",
            steps=[inspect, annotate, validate, export],
        ),
        "page_restructure": PDFPipeline(
            name="page_restructure",
            description="Handles page-level editing operations like rotate, delete, and reorder before validation and export.",
            steps=[inspect, edit, validate, export],
        ),
        "ocr_preflight": PDFPipeline(
            name="ocr_preflight",
            description="Runs inspection and OCR preparation only, then validates and exports an OCR-ready PDF.",
            steps=[inspect, ocr, validate, export],
        ),
    }


def list_available_pipeline_names(toolbox: PDFToolbox | None = None) -> list[str]:
    resolved_toolbox = toolbox or PDFToolbox()
    return list(build_pipeline_catalog(resolved_toolbox).keys())