Blame for src/generator.py

"""Top-level orchestration: DatasetGenerator drives the full per-sample pipeline.

For every sample this wires together, in order:

1. :class:`~src.templates.template_loader.TemplateLoader` -- pick a template + signature box.
2. :class:`~src.signatures.signature_dataset.SignatureDataset` -- pick a signature image.
3. :class:`~src.rendering.signature_preprocessor.SignaturePreprocessor` -- clean it into an RGBA stamp.
4. :class:`~src.placement.placement_engine.PlacementEngine` -- fit/rotate/position it.
5. :class:`~src.augmentation.augmentation_engine.AugmentationEngine` -- signature-layer effects.
6. :class:`~src.rendering.renderer.Renderer` -- blend onto the document.
7. :class:`~src.augmentation.augmentation_engine.AugmentationEngine` -- document-level effects.
8. :class:`~src.analysis.geometry_analyzer.GeometryAnalyzer` -- detect + measure.
9. :class:`~src.analysis.difficulty_analyzer.DifficultyAnalyzer` -- score.
10. :class:`~src.metadata.metadata_exporter.MetadataExporter` -- assemble + write JSON.
11. :class:`~src.analysis.validation_engine.ValidationEngine` -- verify.

A single :class:`GeneratorContext` holds all stateless-between-calls engine
instances for one process; :class:`DatasetGenerator` either drives it
directly (``workers=1``) or fans it out across a
:class:`concurrent.futures.ProcessPoolExecutor` pool, where each worker
process builds exactly one context via ``_init_worker`` and reuses it for
every sample it is assigned. Because every random draw is derived purely
from ``(master_seed, sample_index, stage)``, results are identical
regardless of worker count or scheduling order.
"""

from __future__ import annotations

from concurrent.futures import ProcessPoolExecutor, as_completed
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import numpy as np
from tqdm import tqdm

from src.analysis.difficulty_analyzer import DifficultyAnalyzer
from src.analysis.geometry_analyzer import GeometryAnalyzer
from src.analysis.validation_engine import ValidationEngine, ValidationReport, summarize
from src.augmentation.augmentation_engine import AugmentationEngine
from src.config.schema import PipelineConfig
from src.metadata.metadata_exporter import MetadataExporter
from src.placement.placement_engine import PlacementEngine, PlacementResult
from src.rendering.printed_text import render_printed_name_stamp
from src.rendering.renderer import Renderer
from src.rendering.signature_preprocessor import SignaturePreprocessor
from src.signatures.signature_dataset import SignatureDataset, SignatureRecord
from src.templates.annotation_schema import SignatureBoxAnnotation
from src.templates.template_loader import Template, TemplateLoader
from src.utils.image_io import load_image_any, load_image_rgb, save_image_rgb
from src.utils.logging_utils import get_logger
from src.utils.names import generate_name, generate_role
from src.utils.random_utils import SeedManager, resolve_scalar_or_range
from src.verification.case_types import CASE_TABLE, CaseSpec
from src.verification.registry import SignerIdentityRegistry
from src.visualization.visualizer import save_visualization

logger = get_logger(__name__)

INK_COLOR_POOL: tuple[tuple[int, int, int], ...] = (
    (15, 15, 20),    # black
    (18, 22, 58),    # blue-black
    (24, 34, 96),    # ballpoint blue
    (12, 12, 14),    # deep black
)


@dataclass
class SampleOutcome:
    """Lightweight, picklable result of generating one sample (used across process boundaries)."""

    sample_id: str
    success: bool
    error: "str | None" = None
    difficulty_tier: "str | None" = None
    difficulty_score: "float | None" = None
    validation_passed: "bool | None" = None
    validation_errors: int = 0
    validation_warnings: int = 0


class GeneratorContext:
    """Holds every stateless-between-calls engine instance needed to generate a sample."""

    def __init__(self, config: PipelineConfig) -> None:
        self.config = config
        self.template_loader = TemplateLoader(config.paths.image_dir, config.paths.annotations_dir)
        self.signature_dataset = SignatureDataset(config.paths.dataset_dir)
        self.preprocessor = SignaturePreprocessor()
        self.placement_engine = PlacementEngine(config.placement)
        self.augmentation_engine = AugmentationEngine(config.augmentation)
        self.renderer = Renderer(config.rendering)
        self.geometry_analyzer = GeometryAnalyzer()
        self.difficulty_analyzer = DifficultyAnalyzer(config.difficulty)
        self.metadata_exporter = MetadataExporter()
        self.validation_engine = ValidationEngine(config.validation, config.difficulty)

        # Trigger discovery eagerly so the first sample doesn't pay (and doesn't
        # race another worker on) the filesystem-scan + annotation-cache cost.
        self.template_loader.discover()
        self.signature_dataset.discover()
        self.identity_registry = SignerIdentityRegistry(self.template_loader, self.signature_dataset, config.seed)

    def run_sample(self, sample_index: int) -> tuple[dict[str, Any], "np.ndarray", "ValidationReport | None"]:
        cfg = self.config
        sm = SeedManager(cfg.seed, sample_index)
        prefix = cfg.generation.filename_prefix
        digits = cfg.generation.filename_digits
        sample_id = f"{prefix}_{sample_index:0{digits}d}"

        case_spec = CASE_TABLE.get(cfg.quality.case_type)
        if case_spec is not None:
            return self._run_case_sample(sample_id, sm, case_spec)
        return self._run_generic_sample(sample_id, sm)

    def _run_generic_sample(self, sample_id: str, sm: SeedManager) -> tuple[dict[str, Any], "np.ndarray", "ValidationReport | None"]:
        cfg = self.config
        template = self._select_template(sm)
        page_annotation = self._select_page(template)
        box = self._select_box(page_annotation, sm)

        document_image = load_image_rgb(template.image_path)
        canvas_h, canvas_w = document_image.shape[:2]
        target_box = box.as_tuple()

        signer_name = cfg.signer.name or generate_name(sm.python_rng("signer_name"))
        signer_role = cfg.signer.role or box.role or generate_role(sm.python_rng("signer_role"))

        signature_record = self._select_signature(sm)
        raw_signature = load_image_any(signature_record.path)
        ink_color = INK_COLOR_POOL[int(sm.numpy_rng("ink_color").integers(0, len(INK_COLOR_POOL)))]
        signature_rgba = self.preprocessor.process(raw_signature, ink_color=ink_color)

        placement = self.placement_engine.place(
            signature_rgba, target_box=target_box, canvas_size=(canvas_w, canvas_h), seed_manager=sm,
        )
        signature_layer, signature_aug_records = self.augmentation_engine.apply_to_signature(placement.layer, sm)
        render_result = self.renderer.blend(document_image, signature_layer, sm)
        final_image, document_aug_records = self.augmentation_engine.apply_to_document(render_result.image, sm)

        geometry = self.geometry_analyzer.analyze(
            final_image, target_box=target_box, canvas_size=(canvas_w, canvas_h),
            placement_hint_bbox=placement.intended_bbox,
        )
        difficulty = self.difficulty_analyzer.analyze(geometry)

        quality_rng = sm.numpy_rng("quality")
        quality_values = {
            "scanner_quality": round(resolve_scalar_or_range(cfg.quality.scanner_quality, quality_rng), 4),
            "paper_quality": round(resolve_scalar_or_range(cfg.quality.paper_quality, quality_rng), 4),
            "render_quality": round(resolve_scalar_or_range(cfg.quality.render_quality, quality_rng), 4),
        }

        metadata = self.metadata_exporter.build(
            document_id=sample_id,
            template_id=template.template_id,
            document_type=template.annotation.document_type,
            page=page_annotation.page,
            page_size=(canvas_w, canvas_h),
            dpi=page_annotation.dpi,
            signer_name=signer_name,
            signer_role=signer_role,
            signature_source=signature_record,
            placement=placement,
            render=render_result,
            signature_aug=signature_aug_records,
            document_aug=document_aug_records,
            geometry=geometry,
            difficulty=difficulty,
            quality_cfg=cfg.quality,
            quality_values=quality_values,
            seed=cfg.seed,
            sample_index=sm.sample_index,
            verification=None,
        )

        validation_report = None
        if cfg.generation.run_validation:
            validation_report = self.validation_engine.validate(metadata)

        return metadata, final_image, validation_report

    # ------------------------------------------------------------------ #
    # Verification test-case generation (case_type-driven)
    # ------------------------------------------------------------------ #

    def _run_case_sample(
        self, sample_id: str, sm: SeedManager, case_spec: CaseSpec,
    ) -> tuple[dict[str, Any], "np.ndarray", "ValidationReport | None"]:
        cfg = self.config
        template = self._select_template_for_case(sm, case_spec)
        page_annotation = self._select_page(template)
        boxes = page_annotation.signature_boxes
        if not boxes:
            raise ValueError(f"Page {page_annotation.page} has no signature boxes defined.")
        target_box, companion_box = self._select_case_boxes(boxes, sm)

        document_image = load_image_rgb(template.image_path)
        canvas_h, canvas_w = document_image.shape[:2]
        target_box_tuple = target_box.as_tuple()
        canvas_size = (canvas_w, canvas_h)

        expected_group = self.identity_registry.authorized_signer_group(template.template_id, target_box.id)
        signer_name = cfg.signer.name or self.identity_registry.display_name(expected_group)
        signer_role = cfg.signer.role or target_box.role

        ink_color = INK_COLOR_POOL[int(sm.numpy_rng("ink_color").integers(0, len(INK_COLOR_POOL)))]

        document_after_companion = document_image
        companion_info: "dict[str, Any] | None" = None
        if case_spec.fill_companion:
            if companion_box is None:
                logger.warning(
                    "Case %s wants a companion signer but template %s has only one signature box; skipping.",
                    case_spec.code, template.template_id,
                )
            else:
                companion_group = self.identity_registry.authorized_signer_group(template.template_id, companion_box.id)
                sm_companion = SeedManager(cfg.seed, f"{sm.sample_index}_companion")
                companion_record = self.signature_dataset.sample_from_group(
                    sm_companion.numpy_rng("case:companion_signature"), companion_group, source_filter="genuine",
                )
                companion_raw = load_image_any(companion_record.path)
                companion_rgba = self.preprocessor.process(companion_raw, ink_color=ink_color)
                companion_placement = self.placement_engine.place(
                    companion_rgba, target_box=companion_box.as_tuple(), canvas_size=canvas_size, seed_manager=sm_companion,
                )
                companion_layer, _ = self.augmentation_engine.apply_to_signature(companion_placement.layer, sm_companion)
                companion_render = self.renderer.blend(document_after_companion, companion_layer, sm_companion)
                document_after_companion = companion_render.image
                companion_info = {
                    "box_id": companion_box.id, "filled": True,
                    "signer_group": companion_group, "script": self.identity_registry.script_for(companion_group),
                    "matches_expected": True,
                }

        actual_record, actual_group, actual_source_type, target_rgba = self._resolve_case_signature(
            case_spec, template, target_box, companion_box, expected_group, signer_name, ink_color, sm,
        )

        if target_rgba is None:
            placement = PlacementResult(
                layer=np.zeros((canvas_h, canvas_w, 4), dtype=np.uint8),
                target_box=target_box_tuple, intended_bbox=target_box_tuple,
                x_offset=0.0, y_offset=0.0, rotation_deg=0.0, scale_factor=0.0,
                anchor_point=cfg.placement.anchor_point, misplaced=False,
                misplacement_dx=0.0, misplacement_dy=0.0,
            )
        else:
            placement = self.placement_engine.place(
                target_rgba, target_box=target_box_tuple, canvas_size=canvas_size, seed_manager=sm,
            )
        signature_layer, signature_aug_records = self.augmentation_engine.apply_to_signature(placement.layer, sm)
        render_result = self.renderer.blend(document_after_companion, signature_layer, sm)
        final_image, document_aug_records = self.augmentation_engine.apply_to_document(render_result.image, sm)

        geometry = self.geometry_analyzer.analyze(
            final_image, target_box=target_box_tuple, canvas_size=canvas_size,
            placement_hint_bbox=placement.intended_bbox,
        )
        difficulty = self.difficulty_analyzer.analyze(geometry)

        quality_rng = sm.numpy_rng("quality")
        quality_values = {
            "scanner_quality": round(resolve_scalar_or_range(cfg.quality.scanner_quality, quality_rng), 4),
            "paper_quality": round(resolve_scalar_or_range(cfg.quality.paper_quality, quality_rng), 4),
            "render_quality": round(resolve_scalar_or_range(cfg.quality.render_quality, quality_rng), 4),
        }

        verification_record = {
            "case_type": case_spec.code,
            "verdict": case_spec.verdict,
            "reject_code": case_spec.reject_code,
            "target_box_id": target_box.id,
            "expected_signer_group": expected_group,
            "expected_signer_script": self.identity_registry.script_for(expected_group),
            "actual_signer_group": actual_group,
            "actual_source_type": actual_source_type,
            "actual_signer_script": self.identity_registry.script_for(actual_group) if actual_group else None,
            "companion": companion_info,
        }

        metadata = self.metadata_exporter.build(
            document_id=sample_id,
            template_id=template.template_id,
            document_type=template.annotation.document_type,
            page=page_annotation.page,
            page_size=(canvas_w, canvas_h),
            dpi=page_annotation.dpi,
            signer_name=signer_name,
            signer_role=signer_role,
            signature_source=actual_record,
            placement=placement,
            render=render_result,
            signature_aug=signature_aug_records,
            document_aug=document_aug_records,
            geometry=geometry,
            difficulty=difficulty,
            quality_cfg=cfg.quality,
            quality_values=quality_values,
            seed=cfg.seed,
            sample_index=sm.sample_index,
            verification=verification_record,
        )

        validation_report = None
        if cfg.generation.run_validation:
            validation_report = self.validation_engine.validate(metadata)

        return metadata, final_image, validation_report

    def _select_template_for_case(self, sm: SeedManager, case_spec: CaseSpec) -> Template:
        cfg = self.config
        if cfg.document.template:
            return self.template_loader.get(cfg.document.template)

        templates = self.template_loader.discover()
        if case_spec.fill_companion:
            candidates = sorted(
                tid for tid, t in templates.items() if len(t.annotation.pages[0].signature_boxes) >= 2
            )
            if not candidates:
                logger.warning(
                    "Case %s needs a multi-box template but none were found; using the full template pool.",
                    case_spec.code,
                )
                candidates = sorted(templates.keys())
        else:
            candidates = sorted(templates.keys())

        idx = int(sm.numpy_rng("template_selection").integers(0, len(candidates)))
        return self.template_loader.get(candidates[idx])

    def _select_case_boxes(
        self, boxes: "list[SignatureBoxAnnotation]", sm: SeedManager,
    ) -> "tuple[SignatureBoxAnnotation, SignatureBoxAnnotation | None]":
        idx = int(sm.numpy_rng("box_selection").integers(0, len(boxes)))
        target = boxes[idx]
        others = [b for i, b in enumerate(boxes) if i != idx]
        companion = others[0] if others else None
        return target, companion

    def _resolve_case_signature(
        self, case_spec: CaseSpec, template: Template, target_box: "SignatureBoxAnnotation",
        companion_box: "SignatureBoxAnnotation | None", expected_group: str, expected_name: str,
        ink_color: tuple[int, int, int], sm: SeedManager,
    ) -> "tuple[SignatureRecord | None, str | None, str, np.ndarray | None]":
        """Resolve what actually goes in the target box for one case type.

        Returns ``(signature_record_or_none, actual_signer_group_or_none,
        actual_source_type, rgba_layer_or_none)``.
        """
        mode = case_spec.signature_mode

        if mode == "blank":
            return None, None, "none", None

        if mode == "printed_text":
            layer_rgba = render_printed_name_stamp(expected_name, ink_color=ink_color)
            return None, None, "printed", layer_rgba

        if mode == "genuine_correct":
            group, source_filter, stage = expected_group, "genuine", "case:target_signature"
        elif mode == "genuine_variation":
            group, source_filter, stage = expected_group, "genuine", "case:target_signature_variant"
        elif mode == "forged_skilled":
            group, source_filter, stage = expected_group, "forged", "case:forged_signature"
        elif mode == "wrong_signer_genuine":
            wrong_group = self.identity_registry.random_wrong_signer_group(
                sm.python_rng("case:wrong_signer"), expected_group,
            )
            group, source_filter, stage = wrong_group, "genuine", "case:wrong_signature"
        elif mode == "role_swap":
            if companion_box is None:
                raise ValueError(f"Case {case_spec.code} requires a companion box but none is available.")
            group = self.identity_registry.authorized_signer_group(template.template_id, companion_box.id)
            source_filter, stage = "genuine", "case:role_swap_signature"
        elif mode == "foreign_template_signer":
            group, _foreign_template_id, _foreign_box_id = self.identity_registry.foreign_authorized(template.template_id)
            source_filter, stage = "genuine", "case:foreign_signature"
        else:
            raise ValueError(f"Unknown signature_mode: {mode!r}")

        record = self.signature_dataset.sample_from_group(sm.numpy_rng(stage), group, source_filter)
        raw = load_image_any(record.path)
        layer_rgba = self.preprocessor.process(raw, ink_color=ink_color)
        return record, group, record.source_type, layer_rgba

    def _select_template(self, sm: SeedManager):
        cfg = self.config
        if cfg.document.template:
            return self.template_loader.get(cfg.document.template)
        return self.template_loader.random_template(sm.numpy_rng("template_selection"))

    def _select_page(self, template):
        cfg = self.config
        if cfg.document.page is not None:
            return template.annotation.page_by_number(cfg.document.page)
        return template.annotation.pages[0]

    def _select_box(self, page_annotation, sm: SeedManager):
        boxes = page_annotation.signature_boxes
        if not boxes:
            raise ValueError(f"Page {page_annotation.page} has no signature boxes defined.")
        idx = int(sm.numpy_rng("box_selection").integers(0, len(boxes)))
        return boxes[idx]

    def _select_signature(self, sm: SeedManager):
        cfg = self.config
        if cfg.signature_selection.signature_id:
            return self.signature_dataset.get(cfg.signature_selection.signature_id)
        return self.signature_dataset.sample(
            sm.numpy_rng("signature_selection"), source_filter=cfg.signature_selection.source_filter,
        )


_WORKER_CONTEXT: "GeneratorContext | None" = None
_WORKER_CONFIG: "PipelineConfig | None" = None


def _init_worker(config: PipelineConfig) -> None:
    global _WORKER_CONTEXT, _WORKER_CONFIG
    _WORKER_CONFIG = config
    _WORKER_CONTEXT = GeneratorContext(config)


def _generate_and_export(sample_index: int) -> SampleOutcome:
    """Module-level worker entry point: generate one sample and write it to disk."""
    assert _WORKER_CONTEXT is not None and _WORKER_CONFIG is not None, "Worker not initialized"
    cfg = _WORKER_CONFIG
    ctx = _WORKER_CONTEXT
    prefix = cfg.generation.filename_prefix
    digits = cfg.generation.filename_digits
    sample_id = f"{prefix}_{sample_index:0{digits}d}"
    output_dir = Path(cfg.paths.output_dir)
    png_path = output_dir / f"{sample_id}.png"
    json_path = output_dir / f"{sample_id}.json"

    if not cfg.generation.overwrite and png_path.exists() and json_path.exists():
        return SampleOutcome(sample_id=sample_id, success=True, error="skipped_existing")

    try:
        metadata, final_image, validation_report = ctx.run_sample(sample_index)
        save_image_rgb(png_path, final_image)
        ctx.metadata_exporter.export(metadata, json_path)

        if cfg.generation.save_visualization:
            viz_path = output_dir / "visualizations" / f"{sample_id}_viz.png"
            save_visualization(viz_path, final_image, metadata)

        outcome = SampleOutcome(
            sample_id=sample_id, success=True,
            difficulty_tier=metadata["difficulty"]["tier"],
            difficulty_score=metadata["difficulty"]["score"],
        )
        if validation_report is not None:
            outcome.validation_passed = validation_report.passed
            outcome.validation_errors = len(validation_report.errors)
            outcome.validation_warnings = len(validation_report.warnings)
            if not validation_report.passed:
                logger.warning(
                    "Sample %s failed validation: %s", sample_id,
                    "; ".join(i.message for i in validation_report.errors),
                )
        return outcome
    except Exception as exc:  # noqa: BLE001 - report and continue the batch
        logger.exception("Failed to generate sample %s", sample_id)
        return SampleOutcome(sample_id=sample_id, success=False, error=str(exc))


@dataclass
class GenerationSummary:
    """Batch-level summary returned by :meth:`DatasetGenerator.generate`."""

    total: int
    succeeded: int
    failed: int
    skipped_existing: int
    validation_passed: int
    validation_failed: int
    tier_counts: dict[str, int]
    failed_samples: list[str]

    def as_dict(self) -> dict[str, Any]:
        return {
            "total": self.total,
            "succeeded": self.succeeded,
            "failed": self.failed,
            "skipped_existing": self.skipped_existing,
            "validation_passed": self.validation_passed,
            "validation_failed": self.validation_failed,
            "tier_counts": self.tier_counts,
            "failed_samples": self.failed_samples,
        }


class DatasetGenerator:
    """Drives batch generation of a synthetic signed-document dataset.

    Example:
        >>> generator = DatasetGenerator(config)
        >>> summary = generator.generate()
    """

    def __init__(self, config: PipelineConfig) -> None:
        self.config = config
        Path(config.paths.output_dir).mkdir(parents=True, exist_ok=True)

    def generate(self) -> GenerationSummary:
        """Generate ``config.generation.count`` samples, sequentially or via a worker pool."""
        cfg = self.config
        indices = list(range(cfg.generation.start_index, cfg.generation.start_index + cfg.generation.count))
        outcomes: list[SampleOutcome] = []

        if cfg.generation.workers <= 1:
            _init_worker(cfg)
            for idx in tqdm(indices, desc="Generating samples", unit="doc"):
                outcomes.append(_generate_and_export(idx))
        else:
            with ProcessPoolExecutor(
                max_workers=cfg.generation.workers, initializer=_init_worker, initargs=(cfg,),
            ) as pool:
                futures = {pool.submit(_generate_and_export, idx): idx for idx in indices}
                for future in tqdm(as_completed(futures), total=len(futures), desc="Generating samples", unit="doc"):
                    outcomes.append(future.result())

        return self._summarize(outcomes)

    @staticmethod
    def _summarize(outcomes: list[SampleOutcome]) -> GenerationSummary:
        succeeded = [o for o in outcomes if o.success]
        failed = [o for o in outcomes if not o.success]
        skipped = [o for o in succeeded if o.error == "skipped_existing"]
        validated = [o for o in succeeded if o.validation_passed is not None]
        validation_passed = sum(1 for o in validated if o.validation_passed)
        validation_failed = sum(1 for o in validated if not o.validation_passed)

        tier_counts: dict[str, int] = {}
        for o in succeeded:
            if o.difficulty_tier:
                tier_counts[o.difficulty_tier] = tier_counts.get(o.difficulty_tier, 0) + 1

        summary = GenerationSummary(
            total=len(outcomes),
            succeeded=len(succeeded),
            failed=len(failed),
            skipped_existing=len(skipped),
            validation_passed=validation_passed,
            validation_failed=validation_failed,
            tier_counts=tier_counts,
            failed_samples=[o.sample_id for o in failed],
        )
        logger.info("Generation summary: %s", summary.as_dict())
        return summary