# ============================================================
# src/accountable_record/pipeline/stages.py
# ============================================================
"""Stage registry for the Accountable Record pipeline (orchestration).

WHY:
    The pipeline is a fixed, one-way, nine-stage spine owned by AR. A stage is
    a single-argument transform: PipelineState -> PipelineState. Registration
    is co-located with the transform via the @stage decorator, so a registry
    entry cannot point at a moved, renamed, or missing function. The
    registration IS the function. The runner walks the registry in declared
    order.

OBS:
    Stage modules must be imported for their @stage decorators to run and
    populate the registry. register_stage_modules() imports them once.

REQ:
    Public stage codes are canonicalized as "s010", "s020", ..., "s090".
    The decorator accepts either "010" or "s010", but the registry stores the
    canonical public form.
"""

from collections.abc import Callable
from dataclasses import dataclass
from importlib import import_module

from accountable_record.pipeline.context import PipelineState
from accountable_record.pipeline.diagnostics import Concern

__all__ = [
    "PipelineStage",
    "get_stage",
    "normalize_stage_target",
    "stages_through",
    "valid_stage_targets",
]

StageTransform = Callable[[PipelineState], PipelineState]


@dataclass(frozen=True)
class PipelineStage:
    """One registered pipeline stage transform."""

    code: str
    concern: Concern
    description: str
    transform: StageTransform


# Registry, populated by the @stage decorator at import time, kept in
# declaration/import order.
_REGISTRY: list[PipelineStage] = []
_REGISTERED_CODES: set[str] = set()


def normalize_stage_code(code: str) -> str:
    """Normalize a stage code to the public form, e.g. s010."""
    normalized = code.strip().lower()

    if normalized == "all":
        return normalized

    if not normalized:
        msg = "Stage code must not be empty."
        raise ValueError(msg)

    if not normalized.startswith("s"):
        normalized = f"s{normalized}"

    if len(normalized) != 4 or not normalized[1:].isdigit():
        msg = f"Invalid stage code {code!r}. Expected a code like '010' or 's010'."
        raise ValueError(msg)

    return normalized


def stage(
    code: str,
    concern: Concern,
    description: str,
) -> Callable[[StageTransform], StageTransform]:
    """Register a stage transform under a stage code and concern.

    Usage:
        @stage("s010", Concern.IDENTITY, "Validate repository identity metadata.")
        def transform_identity(state: PipelineState) -> PipelineState: ...

    The shorter form also works:

        @stage("010", Concern.IDENTITY, "Validate repository identity metadata.")
    """
    normalized_code = normalize_stage_code(code)

    def decorator(fn: StageTransform) -> StageTransform:
        if normalized_code in _REGISTERED_CODES:
            msg = f"Duplicate pipeline stage registration: {normalized_code}"
            raise ValueError(msg)

        _REGISTERED_CODES.add(normalized_code)
        _REGISTRY.append(
            PipelineStage(
                code=normalized_code,
                concern=concern,
                description=description,
                transform=fn,
            )
        )
        return fn

    return decorator


# Stage modules whose import triggers @stage registration. Listed in pipeline
# order; this is the one place that enumerates the stages.
_STAGE_MODULES: tuple[str, ...] = (
    "accountable_record.stage_010_identity.transform",
    "accountable_record.stage_020_contract.transform",
    "accountable_record.stage_030_sources.transform",
    "accountable_record.stage_040_elements.transform",
    "accountable_record.stage_050_catalog.transform",
    "accountable_record.stage_060_exports.transform",
    "accountable_record.stage_070_resolution.transform",
    "accountable_record.stage_080_generated.transform",
    "accountable_record.stage_090_docs.transform",
)

_modules_loaded = False


def register_stage_modules() -> None:
    """Import all stage modules once so their @stage decorators run."""
    global _modules_loaded

    if _modules_loaded:
        return

    for module_path in _STAGE_MODULES:
        import_module(module_path)

    _modules_loaded = True


def pipeline_stages() -> tuple[PipelineStage, ...]:
    """Return all registered stages in canonical stage-code order."""
    register_stage_modules()
    return tuple(sorted(_REGISTRY, key=lambda stage_: stage_.code))


def normalize_stage_target(target: str) -> str:
    """Normalize a stage target supplied by the CLI."""
    return normalize_stage_code(target)


def get_stage(code: str) -> PipelineStage:
    """Return a registered stage by public stage code, e.g. s010."""
    normalized = normalize_stage_target(code)

    for stage_ in pipeline_stages():
        if stage_.code == normalized:
            return stage_

    valid = ", ".join(valid_stage_targets())
    msg = f"Unknown stage {code!r}. Expected one of: {valid}."
    raise ValueError(msg)


def stages_through(target: str) -> tuple[PipelineStage, ...]:
    """Return all stages from the first through the requested target, inclusive."""
    normalized = target.strip().lower()
    stages = pipeline_stages()

    if normalized == "all":
        return stages

    target_stage = get_stage(normalized)

    selected: list[PipelineStage] = []
    for stage_ in stages:
        selected.append(stage_)
        if stage_.code == target_stage.code:
            break

    return tuple(selected)


def valid_stage_targets() -> tuple[str, ...]:
    """Return valid public CLI stage targets."""
    return tuple(stage_.code for stage_ in pipeline_stages()) + ("all",)
