# src/pdftl/cli/pipeline.py

"""Manage a pipeline of operations"""

import logging
import sys
import types
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Union

if TYPE_CHECKING:
    import pikepdf

logger = logging.getLogger(__name__)
import pdftl.core.constants as c
from pdftl.cli.constants import SUB_END, SUB_START
from pdftl.cli.whoami import WHOAMI
from pdftl.core.executor import run_operation
from pdftl.core.registry import register_help_topic, registry
from pdftl.core.types import HelpExample, OpResult
from pdftl.exceptions import MissingArgumentError, UserCommandLineError
from pdftl.output.save import save_content
from pdftl.utils.io_helpers import smart_pikepdf_open
from pdftl.utils.user_input import pdf_filename_completer


def _first_or_none(x: list):
    try:
        return x[0]
    except (IndexError, ValueError):
        return None


@dataclass
class InlineSubPipeline:
    """Wrapper for a nested list of stages representing an inline sub-pipeline."""

    stages: list["CliStage"]
    original_text: str = "<inline_sub_pipeline>"
    handle_name: str | None = None

    def __repr__(self):
        return self.original_text


@dataclass
class CliStage:
    """
    A structured representation of a single stage in a processing pipeline.
    """

    operation: str | None = None
    inputs: list[str] = field(default_factory=list)
    input_passwords: list[str | None] = field(default_factory=list)
    handles: dict[str, int] = field(default_factory=dict)
    operation_args: list[str] = field(default_factory=list)
    options: dict[str, Any] = field(default_factory=dict)

    def resolve_stage_io_prompts(self, get_input, stage_num):
        """
        Looks for "PROMPT" in a parsed stage's inputs
        and prompts the user to resolve them.
        """
        logger.debug("resolve_stage_io_prompts")
        # Create an inverse handle map for nice prompts
        handles_inverse = {index: handle for handle, index in self.handles.items()}
        for i, filename in enumerate(self.inputs):
            logger.debug("i=%s, filename=%s", i, filename)
            if filename == "PROMPT":
                logger.debug("Found a PROMPT, asking user")
                desc = f"input #{i + 1}"
                if (handle := handles_inverse.get(i, None)) is not None:
                    desc += f" with handle {handle}"

                if stage_num > 1:
                    desc = f"pipeline stage {stage_num}, {desc}"

                new_filename = get_input(
                    f"Enter a filename for an input PDF ({desc}): ",
                    completer=pdf_filename_completer,
                )

                self.inputs[i] = new_filename


class PipelineResult:
    """The final payload returned by a pipeline execution."""

    pdf: Union["pikepdf.Pdf", None] = None
    results: list[OpResult] = field(default_factory=list)


# pylint: disable=too-few-public-methods
class PipelineManager:
    """Orchestrates the execution of a multi-stage PDF processing pipeline."""

    def __init__(self, stages, input_context, is_inline=False, handles=None) -> None:
        self.stages: list[CliStage] = stages
        self.pipeline_pdf = None
        self.kept_id = None
        self.input_context = input_context
        self.results: list[OpResult] = []
        self.result_discardable = False
        self.is_inline = is_inline
        self.handles = handles or {}

    def save_pdf_file(self, pdf, filename, stage):
        # Pass stage.options directly.
        # This ensures intermediate files don't inherit global flags like 'uncompress'.
        save_options = stage.options

        # Call save. _save_kw_options will handle the isolation logic.
        save_content(
            pdf,
            filename,
            self.input_context,
            **self._save_kw_options(override_options=save_options),
        )

    def run(self):
        """Executes all stages in the pipeline."""
        logger.debug("Running pipeline with %s stages", len(self.stages))
        try:
            for i, stage in enumerate(self.stages):
                stage.resolve_stage_io_prompts(self.input_context.get_input, i + 1)
                self._validate_and_execute_numbered_stage(i, stage)
                stage_output = stage.options.get(c.OUTPUT)
                logger.debug("stage_output=%s", stage_output)

                op_entry = registry.operations.get(stage.operation, {})

                # Check if we have an output file AND a PDF to save.
                # Some operations (like dump_text) handle output themselves via hooks
                # and leave pipeline_pdf as None. We must skip save_pdf in that case.
                #
                # OR they can set skip_pipeline_save in the @register_operation decorator
                #
                skip_pipeline_save = op_entry.get("skip_pipeline_save", False)
                if stage_output and self.pipeline_pdf and not skip_pipeline_save:
                    self.save_pdf_file(self.pipeline_pdf, stage_output, stage)
                else:
                    logger.debug(
                        (
                            "Skipping save. skip_pipeline_save=%s, "
                            "stage_output='%s', self.pipeline_pdf is truthy-%s"
                        ),
                        skip_pipeline_save,
                        stage_output,
                        bool(self.pipeline_pdf),
                    )

        finally:
            if self.pipeline_pdf is not None and not self.is_inline:
                import pikepdf

                if isinstance(self.pipeline_pdf, pikepdf.Pdf):
                    self.pipeline_pdf.close()

    def _save_kw_options(self, override_options=None):
        """
        Construct the keyword arguments for the save_pdf function.
        """
        # If override_options is provided, we use it.
        # Otherwise, we default to empty (standard defaults apply).
        final_options = override_options.copy() if override_options else {}

        # Return the kwargs expected by save_pdf
        return {c.OPTIONS: final_options, "set_pdf_id": self.kept_id}

    def _validate_and_execute_numbered_stage(self, i, stage):
        if not stage.operation and i == len(self.stages) - 1:
            logger.debug("Final stage is empty, proceeding to save.")
            return

        is_first = i == 0
        is_last = i == len(self.stages) - 1

        logger.debug("--- PIPELINE: STAGE %d ---", i + 1)
        logger.debug("Parsed stage: %s", stage)
        # fixme: resolve - here for stdin! or fix validation
        self._validate_stage_args(stage, is_first, is_last)
        self._execute_stage(stage, is_first)

    def _execute_stage(self, stage, is_first):
        """Opens PDFs and runs the operation for a single stage."""
        logger.debug("_execute_stage")
        opened_pdfs = self._open_input_pdfs(stage, is_first)

        if self.pipeline_pdf and self.pipeline_pdf not in opened_pdfs:
            self.pipeline_pdf.close()

        result = self._run_operation(stage, opened_pdfs)
        self._process_result(result, stage, opened_pdfs)

    def _unpack_result_value_and_run_hooks(self, result, stage, opened_pdfs):
        # Unpack OpResult if present
        if isinstance(result, OpResult):
            self.results.append(result)
            self.result_discardable = result.is_discardable

            # CLI Hooks (like printing text to console)
            if not getattr(self.input_context, "is_api", False):
                op_entry = registry.operations.get(stage.operation, {})
                if hook := op_entry.get("cli_hook"):
                    hook(result, stage, self)

            # Update the pipeline data
            result_val = result.pdf
        else:
            self.result_discardable = False
            result_val = result

        if not result_val and opened_pdfs:
            result_val = opened_pdfs[0]

        return result_val

    def _process_result(self, result, stage, opened_pdfs):
        """
        Updates pipeline state and manages file cleanup.
        Crucially, it defers cleanup if the result is a generator.
        """

        import pikepdf

        result_val = self._unpack_result_value_and_run_hooks(result, stage, opened_pdfs)

        # Update the Pipeline State Variable
        self.pipeline_pdf = result_val

        # Smart Cleanup Logic
        # CASE A: Generator (Lazy Evaluation)
        # We CANNOT close opened_pdfs here because the generator hasn't run yet.
        # We rely on the generator having a `finally` block to close these when done.
        if isinstance(result_val, types.GeneratorType):
            logger.debug("Stage returned a generator. Deferring cleanup to the generator.")
            # Do nothing. The generator "owns" the opened_pdfs now.

        # CASE B: Standard Object (Immediate Evaluation)
        # We can safely close the inputs that aren't the result.
        else:
            for pdf in opened_pdfs:
                if pdf != result_val:
                    if isinstance(pdf, pikepdf.Pdf):
                        pdf.close()

    def _validate_stage_args(self, stage, is_first, is_last):
        """Validates arguments for a given stage."""
        if not stage.inputs and is_first:
            raise MissingArgumentError(
                "No initial input files provided. "
                "\n  Maybe you put an operation before the input file?"
                f"\n  Correct syntax: {WHOAMI} <input>... <operation> [<other arguments>]"
            )

        op_data = registry.operations.get(stage.operation, {})
        op_requires_output = " output " in op_data.get("usage", "")

        # Check if the stage has an output option
        has_output = bool(stage.options.get(c.OUTPUT))

        if (
            is_last
            and not self.is_inline
            and (stage.operation == "filter" or op_requires_output)
            and not has_output
        ):
            raise MissingArgumentError(
                f"The '{stage.operation}' operation requires 'output <file>' in the final stage."
            )

        num_explicit = len([i for i in stage.inputs if i != "_"])
        effective_inputs = num_explicit + (0 if is_first else 1)

        self._validate_number_of_effective_inputs(stage.operation, effective_inputs)

    def _validate_number_of_effective_inputs(self, operation, effective_inputs):
        if (op_data := registry.operations.get(operation)) is None:
            return
        op_type = op_data.get("type")
        logger.debug("operation=%s, op_type=%s", operation, op_type)
        if op_type == "single input operation" and effective_inputs != 1:
            raise UserCommandLineError(
                f"The '{operation}' operation requires one input, "
                f"but received {effective_inputs} effective input(s)."
            )
        if op_type == "multi input operation" and effective_inputs < 2:
            raise MissingArgumentError(
                f"The '{operation}' operation requires 2 or more inputs, "
                f"but received {effective_inputs} effective input(s)."
            )

    def _run_operation(self, stage, opened_pdfs):
        """Dispatches to the correct command function based on the operation."""
        operation = stage.operation
        op_data = registry.operations.get(operation)
        if op_data:
            op_function, arg_style = op_data.get("function"), op_data.get("args")
        else:
            op_function, arg_style = None, None
        if not op_function or not arg_style:
            raise ValueError(f"Operation '{operation}' is not fully configured.")

        # Determine output pattern from local stage options or default
        output_pattern = stage.options.get(c.OUTPUT, "pg_%04d.pdf")

        call_context = {
            c.OPERATION_NAME: operation,
            c.INPUTS: stage.inputs,
            c.OPENED_PDFS: opened_pdfs,
            c.INPUT_FILENAME: _first_or_none(stage.inputs),
            c.INPUT_PASSWORD: _first_or_none(stage.input_passwords),
            c.INPUT_PDF: _first_or_none(opened_pdfs),
            c.OPERATION_ARGS: stage.operation_args,
            c.ALIASES: stage.handles,
            c.OVERLAY_PDF: _first_or_none(stage.operation_args),
            c.OUTPUT: stage.options.get(c.OUTPUT, None),
            c.OUTPUT_PATTERN: output_pattern,
            c.GET_INPUT: self.input_context.get_input,
        }

        return run_operation(operation, call_context)

    def _make_op_args(self, arg_style, context):
        pos_arg_names, kw_arg_map = arg_style[:2]
        kw_const_arg_map = arg_style[2] if len(arg_style) > 2 else {}
        pos_args = [context[name] for name in pos_arg_names]
        kw_args = {key: context[val] for key, val in kw_arg_map.items()}
        kw_args.update(kw_const_arg_map)
        return pos_args, kw_args

    def _open_pdf_from_special_input(self, password: str | None, is_first: bool):
        """
        Handles opening a PDF from a special input source (stdin or a
        previous pipeline stage).
        """
        import pikepdf

        if is_first:
            logger.debug("Reading PDF from stdin for first stage")
            if sys.stdin.isatty():
                raise UserCommandLineError("Expected PDF data from stdin, but none was provided.")

            try:
                return smart_pikepdf_open(None, password)
            except pikepdf.PasswordError as exc:
                msg = (
                    str(exc)
                    if password
                    else f"PDF data on stdin is encrypted and requires a password. "
                    f"For help: {WHOAMI} help inputs"
                )
                raise UserCommandLineError(msg) from exc

        logger.debug("Using PDF from previous stage for input '_'")
        if not self.pipeline_pdf:
            raise UserCommandLineError(
                "Pipeline error: No PDF available from previous stage for input '_'."
            )
        return self.pipeline_pdf

    def _open_pdf_from_file(self, filename: str, password: str | None):
        """
        Opens a PDF from a file path, handling passwords and file-related errors.
        """
        import pikepdf

        try:
            logger.debug("Opening file '%s'", filename)
            return smart_pikepdf_open(filename, password)
        except pikepdf.PasswordError as exc:
            msg = (
                str(exc)
                if password
                else f"File '{filename}' is encrypted and requires a password. "
                f"For help: {WHOAMI} help inputs"
            )
            raise UserCommandLineError(msg) from exc
        except (FileNotFoundError, pikepdf.PdfError) as exception:
            raise UserCommandLineError(
                f"Could not open '{filename}': the file may be corrupted or is not a valid PDF.\n"
                f"  (Details: {exception})"
            ) from exception

    def _open_input_pdfs(self, stage, is_first):
        """Opens all PDF inputs required for a stage."""
        logger.debug("_open_input_pdfs")
        opened_pdfs = []

        # We assume global/final options (like keep_first_id) are attached to the LAST stage.
        final_stage_options = self.stages[-1].options if self.stages else {}

        for i, item in enumerate(stage.inputs):
            password = stage.input_passwords[i]

            if isinstance(item, InlineSubPipeline):
                logger.debug("Detected InlineSubPipeline input at index %s", i)
                pdf_obj = self._get_subpipeline_output_pdf(stage, i, item, opened_pdfs)
            # [CHANGE] Check if item is a sibling handle defined in THIS stage
            elif isinstance(item, str) and item in stage.handles and stage.handles[item] < i:
                pdf_obj = opened_pdfs[stage.handles[item]]
            elif isinstance(item, str) and item in self.handles:
                # Use the existing object from the outer scope
                pdf_obj = self.handles[item]
            elif item in ["-", "_"]:
                pdf_obj = self._open_pdf_from_special_input(password, is_first)
            else:
                pdf_obj = self._open_pdf_from_file(item, password)

            opened_pdfs.append(pdf_obj)
            if (
                final_stage_options.get("keep_first_id")
                and is_first
                and i == 0
                and len(opened_pdfs) > 0
            ):
                self.kept_id = list(opened_pdfs[0].trailer.ID)

        if final_stage_options.get("keep_final_id") and len(opened_pdfs) > 0:
            self.kept_id = list(opened_pdfs[-1].trailer.ID)

        return opened_pdfs

    def _get_subpipeline_output_pdf(self, stage, item_idx, item, opened_pdfs):
        # [CHANGE] Construct the handle scope for the child
        # 1. Start with inherited handles (from outer scopes)
        child_handles = self.handles.copy()
        # 2. Add sibling handles defined earlier in THIS stage
        # (e.g., in 'A=a.pdf B=JOB', B can see A)
        for h_name, h_idx in stage.handles.items():
            if h_idx < item_idx:
                child_handles[h_name] = opened_pdfs[h_idx]

        # 1. Create a FRESH PipelineManager for the inline stages.
        #    We pass the same input_context (for prompts/stdin),
        #    but the handle scope is implicitly reset because it's a new instance.
        sub_manager = PipelineManager(
            stages=item.stages,
            input_context=self.input_context,
            is_inline=True,
            handles=child_handles,
        )

        # 2. Execute the inner pipeline
        sub_manager.run()

        # 3. The result is our input.
        # If pipeline_pdf is None (inner pipe had no output), this logic handles it
        # by raising the appropriate error inside the sub-manager or returning None here.
        if sub_manager.pipeline_pdf is None:
            # This usually happens if the inner pipeline was purely side-effects
            # or failed silently. Treated as an error for process substitution.
            raise UserCommandLineError("Inline pipeline returned no output PDF.")

        pdf_obj = sub_manager.pipeline_pdf
        # Transfer ownership of the object so we don't double-close incorrectly later,
        # though sub_manager's finally block handles its own cleanup.
        # Ideally, we want the result object to stay open.
        sub_manager.pipeline_pdf = None

        return pdf_obj


@register_help_topic(
    "pipeline",
    title="pipeline syntax",
    desc="Chaining operations, named handles, and pipeline substitution",
    examples=[
        HelpExample(
            desc="Shuffle two documents, then crop the resulting pages to A4",
            cmd="a.pdf b.pdf shuffle --- crop '(a4)' output out.pdf",
        ),
        HelpExample(
            desc=(
                "Crop all pages to A3 in landscape,\n"
                "and preview the effect of cropping odd pages to A4"
            ),
            cmd="in.pdf crop '(A3_l)' --- crop 'odd(A4)' preview output out.pdf",
        ),
        HelpExample(
            desc=(
                "Save a snapshot of a rotated file, then apply a stamp and save the final version"
            ),
            cmd=(
                "in.pdf rotate right output rotated_snapshot.pdf --- "
                "background watermark.pdf output final.pdf"
            ),
        ),
        HelpExample(
            desc=(
                "Use pipeline substitution (JOB...DONE) to rotate one file before "
                "merging it with another."
            ),
            cmd=f"{SUB_START} in.pdf cat right {SUB_END} main.pdf cat output final.pdf",
        ),
        HelpExample(
            desc=("Rotate and stamp a.pdf, crop b.pdf, then combine selected pages from both"),
            cmd=(
                f"A={SUB_START} a.pdf rotate right --- stamp logo.pdf {SUB_END} "
                f"B={SUB_START} b.pdf crop '(a4)' {SUB_END} "
                "cat A1-3 B2-end output combined.pdf"
            ),
        ),
        HelpExample(
            desc=("Join a contract with a stamped copy of itself"),
            cmd=(
                f"contract.pdf {SUB_START} contract.pdf stamp logo.pdf {SUB_END}"
                " output combined.pdf"
            ),
        ),
    ],
)
def _pipeline_help_topic():
    """
    Construct complex workflows using pipelines, handles, and substitution.

    **1. Chaining (`---`)**
    Multiple operations can be chained together using `---`. The output of one
    stage becomes the input for the next stage. If a stage has no explicit
    inputs, it automatically uses the result from the previous stage.

    **2. Named Handles (`X=...`)**
    You can assign single capital letter handles (A-Z) to inputs to refer
    to them later:
      `pdftl A=logo.pdf B=content.pdf ...`
    This allows you to re-use a specific file or result multiple times in
    different stages (e.g., `cat A B A`).

    **3. Pipeline Substitution (`X=JOB ... DONE`)**
    Similar to command substitution in a shell, you can process files in a
    temporary sub-pipeline before assigning them to a handle.
      `S=JOB in.pdf rotate right DONE`
    This runs the commands between JOB and DONE, and assigns the final result
    to the handle `S`. You can then use `S` like any other input file.
    """


@register_help_topic(
    "input",
    title=c.INPUTS,
    desc="Specifying input files and passwords",
)
def _inputs_help_topic():
    """
    The general syntax for providing input to an operation is:

    ```
    <inputs> [ input_pw <password>... ]
    ```

    `<inputs>` is a space-separated list of one or more input PDF
    sources. Each source can be:

      - A file path: `my_doc.pdf`

      - A handle assignment (for referring to files in
        operations): `A=my_doc.pdf`

      - A single dash `-` to read from standard input (stdin).

      - The keyword `PROMPT` to be interactively asked for a
        file path.

    `[ input_pw <password>... ]` is an optional block to provide
    owner passwords for encrypted files. The passwords in the
    `<password>...` list can be assigned in two ways:

      - By position: Passwords are applied sequentially to the
        encrypted input files in the order they appear, as in:

        `enc1.pdf plain.pdf enc2.pdf input_pw pass1 pass2`

      - By handle: If an input file has a handle (e.g.,
        `A=file.pdf`), its password can be assigned using the same
        handle. This is the most reliable method when using
        multiple encrypted files. As in:

        `A=enc1.pdf B=enc2.pdf input_pw B=pass2 A=pass1`

    The keyword `PROMPT` can be used in the list to be securely
    prompted for a password. This is recommended.
    """
