dev/registry/extract_parameters.py

#!/usr/bin/env python3
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
"""
Airflow Registry Parameter & Module Extractor

Discovers provider modules (operators, hooks, sensors, triggers, etc.) at runtime
and extracts constructor/function parameters via MRO or signature inspection.
Produces both modules.json (the full module catalog) and per-provider
parameters.json files.

Must be run inside breeze where all providers are installed.

Usage:
    breeze run python dev/registry/extract_parameters.py

Output:
    - dev/registry/modules.json (+ registry/src/_data/modules.json on host)
    - dev/registry/output/versions/{provider_id}/{version}/parameters.json
    - registry/src/_data/versions/{provider_id}/{version}/parameters.json
    - dev/registry/runtime_modules.json (debug stats)
"""

from __future__ import annotations

import argparse
import concurrent.futures
import importlib
import inspect
import json
import logging
import re
import sys
import typing
from collections import defaultdict
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path

import yaml
from extract_metadata import fetch_provider_inventory, read_inventory
from registry_contract_models import validate_modules_catalog, validate_provider_parameters
from registry_tools.types import (
    BASE_CLASS_IMPORTS,
    CLASS_LEVEL_CATEGORY_OVERRIDES,
    CLASS_LEVEL_SECTIONS,
    DICT_SHAPED_CLASS_LEVEL_SECTIONS,
    MODULE_LEVEL_SECTIONS,
)

AIRFLOW_ROOT = Path(__file__).parent.parent.parent
SCRIPT_DIR = Path(__file__).parent
PROVIDERS_DIR = AIRFLOW_ROOT / "providers"

PROVIDERS_JSON_CANDIDATES = [
    SCRIPT_DIR / "providers.json",
    AIRFLOW_ROOT / "registry" / "src" / "_data" / "providers.json",
]

# Inside breeze, write to dev/registry/output/ (mounted).
# On host, also write to the registry data directory.
OUTPUT_DIRS = [
    SCRIPT_DIR / "output",
    AIRFLOW_ROOT / "registry" / "src" / "_data",
]


@dataclass
class Module:
    """A discovered provider module (operator, hook, sensor, etc.)."""

    id: str
    name: str  # Class name (e.g., SnowflakeOperator)
    type: str  # operator, hook, sensor, trigger, transfer, etc.
    import_path: str  # Full import path to the class
    module_path: str  # Module file path
    short_description: str
    docs_url: str
    source_url: str
    category: str
    provider_id: str
    provider_name: str
    supports_durable_execution: bool


def get_category(integration_name: str) -> str:
    """Slugify an integration name into a category ID."""
    cat_id = integration_name.lower().replace(" ", "-").replace("(", "").replace(")", "")
    return re.sub(r"[^a-z0-9-]", "", cat_id)


def format_annotation(annotation: type, _depth: int = 0) -> str | None:
    """Convert a type annotation to a human-readable string."""
    if _depth > 5:
        return str(annotation)

    if annotation is inspect.Parameter.empty:
        return None

    if annotation is type(None):
        return "None"

    origin = getattr(annotation, "__origin__", None)

    # typing.Union (includes X | Y on 3.10+)
    if origin is typing.Union:
        args = typing.get_args(annotation)
        parts = [format_annotation(a, _depth + 1) for a in args]
        return " | ".join(p for p in parts if p)

    if origin is not None:
        args = typing.get_args(annotation)
        origin_name = getattr(origin, "__name__", str(origin))
        if args:
            arg_strs = [format_annotation(a, _depth + 1) or "Any" for a in args]
            return f"{origin_name}[{', '.join(arg_strs)}]"
        return origin_name

    if hasattr(annotation, "__name__"):
        return annotation.__name__

    s = str(annotation)
    s = re.sub(r"\btyping\.", "", s)
    s = re.sub(r"\bcollections\.abc\.", "", s)
    return s


def format_default(default: object) -> object:
    """Convert a default value to a JSON-serializable representation."""
    if default is inspect.Parameter.empty:
        return None

    if default is None:
        return None

    if isinstance(default, (str, int, float, bool)):
        return default

    if isinstance(default, (list, tuple, dict)):
        try:
            json.dumps(default)
            return default
        except (TypeError, ValueError):
            pass

    try:
        return str(default)
    except Exception:
        return repr(default)


def get_params_from_signature(sig: inspect.Signature, qualified_origin: str) -> dict[str, dict]:
    """Convert an inspect.Signature into parameter metadata."""
    params: dict[str, dict] = {}

    for name, param in sig.parameters.items():
        if name in ("self", "cls"):
            continue
        if name.startswith("_"):
            continue
        if param.kind in (param.VAR_KEYWORD, param.VAR_POSITIONAL):
            continue

        params[name] = {
            "name": name,
            "type": format_annotation(param.annotation),
            "default": format_default(param.default),
            "required": param.default is inspect.Parameter.empty,
            "origin": qualified_origin,
        }

    return params


def get_params_from_class(cls: type) -> dict[str, dict]:
    """
    Extract all __init__ parameters by walking the MRO in reverse
    so child class overrides parent for the same parameter name.
    Records the full qualified origin (module.ClassName) for each param.
    """
    params: dict[str, dict] = {}

    for klass in reversed(cls.__mro__):
        if klass is object:
            continue

        init = klass.__dict__.get("__init__")
        if init is None:
            continue

        try:
            sig = inspect.signature(init)
        except (TypeError, ValueError):
            continue

        qualified_origin = f"{klass.__module__}.{klass.__qualname__}"
        params.update(get_params_from_signature(sig, qualified_origin))

    return params


def get_mro_chain(cls: type) -> list[str]:
    """Return the full MRO as a list of qualified class names."""
    return [f"{k.__module__}.{k.__qualname__}" for k in cls.__mro__ if k is not object]


def parse_param_descriptions(doc: str) -> dict[str, str]:
    """Parse ``:param name: description`` entries from a docstring."""
    descriptions: dict[str, str] = {}
    if not doc:
        return descriptions

    for match in re.finditer(
        r":param\s+(\w+):\s*(.+?)(?=\n\s*:\w|\n\s*\.\.|$)",
        doc,
        re.DOTALL,
    ):
        name = match.group(1)
        desc = match.group(2).strip()
        desc = re.sub(r"\s+", " ", desc)
        if name not in descriptions:
            descriptions[name] = desc

    return descriptions


def parse_docstring_params(cls: type) -> dict[str, str]:
    """
    Parse :param name: description from class and ancestor docstrings.
    Child class descriptions take priority.
    """
    descriptions: dict[str, str] = {}

    for klass in cls.__mro__:
        if klass is object:
            continue
        for name, desc in parse_param_descriptions(getattr(klass, "__doc__", None) or "").items():
            if name not in descriptions:
                descriptions[name] = desc

    return descriptions


def extract_class_params(cls: type) -> tuple[list[str], list[dict]]:
    """
    Extract parameter list for a class, merging signature + docstrings.
    Only includes params originating from provider classes (airflow.providers.*).
    Returns (mro_chain, filtered_params).
    """
    params = get_params_from_class(cls)
    descriptions = parse_docstring_params(cls)

    for name, param in params.items():
        if name in descriptions:
            param["description"] = descriptions[name]
        else:
            param["description"] = None

    provider_params = [p for p in params.values() if p["origin"].startswith("airflow.providers.")]
    mro = get_mro_chain(cls)

    return mro, provider_params


def extract_callable_params(func: typing.Callable[..., typing.Any]) -> tuple[list[str], list[dict]]:
    """Extract parameter metadata from a callable signature and docstring."""
    try:
        sig = inspect.signature(func)
    except (TypeError, ValueError) as e:
        raise TypeError(f"Could not inspect callable signature: {e}") from e

    qualified_origin = f"{func.__module__}.{func.__qualname__}"
    params = get_params_from_signature(sig, qualified_origin)
    descriptions = parse_param_descriptions(getattr(func, "__doc__", None) or "")

    for name, param in params.items():
        param["description"] = descriptions.get(name)

    provider_params = [p for p in params.values() if p["origin"].startswith("airflow.providers.")]
    return [], provider_params


def extract_params(obj: object) -> tuple[list[str], list[dict]]:
    """Extract parameter metadata from either a class or a callable."""
    if inspect.isclass(obj):
        return extract_class_params(obj)
    if callable(obj):
        return extract_callable_params(typing.cast("typing.Callable[..., typing.Any]", obj))
    raise TypeError(f"Unsupported import type: {type(obj)!r}")


def import_symbol(import_path: str) -> object | None:
    """Import a symbol from its full dotted path."""
    parts = import_path.rsplit(".", 1)
    if len(parts) != 2:
        return None

    module_path, symbol_name = parts
    try:
        module = importlib.import_module(module_path)
        return getattr(module, symbol_name, None)
    except Exception as e:
        print(f"  WARN failed to import {import_path}: {e}")
        return None


def find_json(candidates: list[Path], name: str) -> Path:
    """Find first existing JSON file from candidates list."""
    for candidate in candidates:
        if candidate.exists():
            return candidate
    print(f"ERROR: {name} not found. Searched:")
    for c in candidates:
        print(f"  - {c}")
    print(f"\nCopy {name} to dev/registry/ or run extract_metadata.py first.")
    sys.exit(1)


log = logging.getLogger(__name__)


def load_base_classes() -> dict[str, type]:
    """Import base classes for issubclass checks.

    Returns a mapping of type name -> base class (e.g. "sensor" -> BaseSensorOperator).
    """
    base_classes: dict[str, type] = {}
    for type_name, import_path in BASE_CLASS_IMPORTS:
        module_path, class_name = import_path.rsplit(".", 1)
        try:
            mod = importlib.import_module(module_path)
            base_classes[type_name] = getattr(mod, class_name)
        except Exception:
            log.warning("Could not import base class %s", import_path)
    return base_classes


def _should_skip_class(name: str) -> bool:
    """Return True if a class name should be excluded from discovery."""
    if name.startswith("_"):
        return True
    if name.startswith("Base"):
        return True
    if "Abstract" in name or "Mixin" in name:
        return True
    return False


def _get_first_docstring_line(obj: object) -> str | None:
    """Return the first non-empty line of a class docstring, or None."""
    doc = getattr(obj, "__doc__", None)
    if not doc:
        return None
    for line in doc.strip().splitlines():
        stripped = line.strip()
        if stripped:
            return stripped
    return None


def _get_source_line(cls: type) -> int | None:
    """Return the source line number for a class, or None if unavailable."""
    try:
        return inspect.getsourcelines(cls)[1]
    except (OSError, TypeError):
        return None


def load_resumable_job_mixin() -> type | None:
    """Import ResumableJobMixin for durable-execution capability checks, or None if unavailable."""
    try:
        from airflow.sdk import ResumableJobMixin

        return ResumableJobMixin
    except ImportError:
        log.warning("Could not import ResumableJobMixin")
        return None


def is_durable_capable(cls: type, resumable_mixin: type | None) -> bool:
    """Return True if a class implements durable/crash-safe execution.

    Two ways to qualify:
    1. A class-level `__supports_durable_execution = True`
    declaration (for operators that implement this directly against
    task_state_store, without ResumableJobMixin -- e.g. KubernetesPodOperator,
    AgentOperator).
    2. Genuinely implementing ResumableJobMixin's contract.

    The first path deliberately looks up the class prefixed attribute
    (`_{ClassName}__supports_durable_execution`) rather than a fixed string.
    A subclass that overrides execute() itself (e.g. SparkKubernetesOperator)
    may not preserve the parent's task_state_store reconnect behavior, so the
    declaration must not be inherited -- only the exact class that wrote
    `__supports_durable_execution` in its own body qualifies this way.

    Inheriting the mixin alone is not sufficient for the second path: a
    complete override is inert unless execute() actually calls
    execute_resumable().
    """
    if getattr(cls, f"_{cls.__name__}__supports_durable_execution", None) is True:
        return True

    if resumable_mixin is None or resumable_mixin not in cls.__mro__:
        return False

    if inspect.isabstract(cls):
        return False

    execute = getattr(cls, "execute", None)
    if execute is None:
        return False
    try:
        source = inspect.getsource(execute)
    except (OSError, TypeError):
        return False

    return "execute_resumable" in source


def _resolve_dotted_path(class_path: str) -> tuple[str, str, object] | None:
    """Split a dotted ``module.name`` path and import ``name`` from that module.

    Returns ``None`` (without logging) if ``class_path`` has no dot, or ``None``
    (after logging a warning) if the module import fails. Otherwise returns
    ``(module_path, name, obj)``, where ``obj`` may be ``None`` if the module
    has no such attribute — callers decide how to treat a missing attribute.
    """
    parts = class_path.rsplit(".", 1)
    if len(parts) != 2:
        return None
    module_path, name = parts
    try:
        mod = importlib.import_module(module_path)
        obj = getattr(mod, name, None)
    except Exception:
        log.warning("Could not import %s", class_path)
        return None
    return module_path, name, obj


def discover_classes_from_provider(
    provider_yaml_path: Path,
    base_classes: dict[str, type],
    resumable_mixin: type | None = None,
    inventory: dict[str, str] | None = None,
    version: str = "",
) -> list[dict]:
    """Discover classes from a single provider by importing its modules at runtime.

    Reads the provider.yaml to find which modules/classes to inspect, imports them,
    and returns metadata for each discovered class with all 12 Module fields.
    """
    with open(provider_yaml_path) as f:
        provider_yaml = yaml.safe_load(f)

    provider_id = provider_yaml.get("package-name", "").replace("apache-airflow-providers-", "")
    if not provider_id:
        return []

    provider_name = provider_yaml.get("name", provider_id.replace("-", " ").title())
    provider_rel_path = provider_yaml_path.parent.relative_to(PROVIDERS_DIR)
    tag = f"providers-{provider_id}/{version}" if version else "main"
    base_docs_url = f"https://airflow.apache.org/docs/apache-airflow-providers-{provider_id}/stable"
    base_source_url = f"https://github.com/apache/airflow/blob/{tag}/providers/{provider_rel_path}/src"

    # Build integration-name lookup for module-level sections
    # Maps (section_name, module_path) -> integration_name
    integration_by_module: dict[tuple[str, str], str] = {}
    for section_name in list(MODULE_LEVEL_SECTIONS) + ["bundles"]:
        for group in provider_yaml.get(section_name, []):
            integration = group.get("integration-name", "")
            for mp in group.get("python-modules", []):
                integration_by_module[(section_name, mp)] = integration

    def resolve_docs_url(full_class_path: str, module_path: str) -> str:
        """Look up docs URL from inventory, falling back to manual construction."""
        if inventory and full_class_path in inventory:
            return f"{base_docs_url}/{inventory[full_class_path]}"
        api_ref_path = module_path.replace(".", "/")
        return f"{base_docs_url}/_api/{api_ref_path}/index.html#{full_class_path}"

    def make_source_url(cls: type, module_path: str) -> str:
        """Construct a GitHub source URL with line number when available."""
        url = f"{base_source_url}/{module_path.replace('.', '/')}.py"
        line = _get_source_line(cls)
        if line:
            url += f"#L{line}"
        return url

    def make_entry(
        cls_or_obj: type,
        name: str,
        module_type: str,
        import_path: str,
        module_path: str,
        integration: str = "",
        category: str = "",
        transfer_desc: str | None = None,
    ) -> dict:
        """Build a full module entry dict with all 12 fields."""
        module_name = module_path.split(".")[-1]
        docstring = _get_first_docstring_line(cls_or_obj)
        short_desc = docstring or transfer_desc or f"{integration} {module_type}".strip()

        return {
            "id": f"{provider_id}-{module_name}-{name}",
            "name": name,
            "type": module_type,
            "import_path": import_path,
            "module_path": module_path,
            "short_description": short_desc,
            "docs_url": resolve_docs_url(import_path, module_path),
            "source_url": make_source_url(cls_or_obj, module_path),
            "category": category or get_category(integration),
            "provider_id": provider_id,
            "provider_name": provider_name,
            "supports_durable_execution": is_durable_capable(cls_or_obj, resumable_mixin),
        }

    discovered: list[dict] = []

    # --- Module-level sections (operators, hooks, sensors, triggers, bundles) ---
    for section_name, module_type in MODULE_LEVEL_SECTIONS.items():
        expected_base = base_classes.get(module_type)
        for group in provider_yaml.get(section_name, []):
            integration = group.get("integration-name", "")
            category = get_category(integration)
            for module_path in group.get("python-modules", []):
                try:
                    mod = importlib.import_module(module_path)
                except Exception:
                    log.warning("Could not import module %s", module_path)
                    continue

                for name, cls in inspect.getmembers(mod, inspect.isclass):
                    if cls.__module__ != mod.__name__:
                        continue
                    if _should_skip_class(name):
                        continue
                    if expected_base and not issubclass(cls, expected_base):
                        continue

                    discovered.append(
                        make_entry(
                            cls,
                            name,
                            module_type,
                            f"{module_path}.{name}",
                            module_path,
                            integration,
                            category,
                        )
                    )

    # --- Transfers (module-level, singular python-module key) ---
    transfer_base = base_classes.get("operator")
    for transfer in provider_yaml.get("transfers", []):
        module_path = transfer.get("python-module", "")
        if not module_path:
            continue
        source = transfer.get("source-integration-name", "")
        target = transfer.get("target-integration-name", "")
        transfer_desc = f"Transfer from {source} to {target}" if source and target else None
        category = get_category(source) if source else ""

        try:
            mod = importlib.import_module(module_path)
        except Exception:
            log.warning("Could not import module %s", module_path)
            continue

        for name, cls in inspect.getmembers(mod, inspect.isclass):
            if cls.__module__ != mod.__name__:
                continue
            if _should_skip_class(name):
                continue
            if transfer_base and not issubclass(cls, transfer_base):
                continue

            discovered.append(
                make_entry(
                    cls,
                    name,
                    "transfer",
                    f"{module_path}.{name}",
                    module_path,
                    source,
                    category,
                    transfer_desc,
                )
            )

    # --- Class-level sections (notifications, secrets-backends, logging, executors) ---
    for section_name, module_type in CLASS_LEVEL_SECTIONS.items():
        for class_path in provider_yaml.get(section_name, []):
            if not class_path or not isinstance(class_path, str):
                continue
            if (resolved := _resolve_dotted_path(class_path)) is None:
                continue
            module_path, class_name, candidate = resolved
            if candidate is None or not inspect.isclass(candidate):
                log.warning("%s is not a class", class_path)
                continue
            cls = candidate

            discovered.append(
                make_entry(
                    cls,
                    class_name,
                    module_type,
                    class_path,
                    module_path,
                    category=CLASS_LEVEL_CATEGORY_OVERRIDES.get(section_name, section_name),
                )
            )

    # --- Dict-shaped class-level sections (plugins, dialects; each entry is a
    # dict carrying an integration-name field plus a class-path field, see types.py) ---
    for section_name, (
        module_type,
        class_field,
        integration_field,
    ) in DICT_SHAPED_CLASS_LEVEL_SECTIONS.items():
        for entry in provider_yaml.get(section_name, []):
            if not isinstance(entry, dict):
                continue
            if not (class_path := entry.get(class_field, "")):
                continue
            if (resolved := _resolve_dotted_path(class_path)) is None:
                continue
            module_path, class_name, candidate = resolved
            if candidate is None or not inspect.isclass(candidate):
                log.warning("%s is not a class", class_path)
                continue

            discovered.append(
                make_entry(
                    candidate,
                    class_name,
                    module_type,
                    class_path,
                    module_path,
                    integration=entry.get(integration_field, ""),
                    category=CLASS_LEVEL_CATEGORY_OVERRIDES.get(section_name, section_name),
                )
            )

    # --- Task decorators (class-name key in each entry) ---
    for decorator in provider_yaml.get("task-decorators", []):
        if not (class_path := decorator.get("class-name", "")):
            continue
        if (resolved := _resolve_dotted_path(class_path)) is None:
            continue
        module_path, func_name, obj = resolved
        if obj is None:
            continue

        decorator_name = decorator.get("name", "")
        display_name = f"@task.{decorator_name}" if decorator_name else func_name
        docstring = _get_first_docstring_line(obj) if hasattr(obj, "__doc__") else None
        short_desc = docstring or f"Task decorator for {decorator_name or func_name}"

        discovered.append(
            {
                "id": f"{provider_id}-decorator-{decorator_name or func_name}",
                "name": display_name,
                "type": "decorator",
                "import_path": class_path,
                "module_path": module_path,
                "short_description": short_desc,
                "docs_url": resolve_docs_url(class_path, module_path),
                "source_url": f"{base_source_url}/{module_path.replace('.', '/')}.py",
                "category": "decorators",
                "provider_id": provider_id,
                "provider_name": provider_name,
            }
        )

    return discovered


def compare_with_ast(
    runtime_classes: list[dict],
    modules_json_path: Path,
) -> dict:
    """Compare runtime-discovered classes against AST-produced modules.json.

    Returns a stats dict with counts of phantoms, misses, and type mismatches.
    """
    with open(modules_json_path) as f:
        ast_data = json.load(f)

    ast_modules = ast_data.get("modules", [])

    ast_by_path: dict[str, dict] = {}
    for m in ast_modules:
        path = m.get("import_path", "")
        if path:
            ast_by_path[path] = m

    runtime_by_path: dict[str, dict] = {}
    for r in runtime_classes:
        path = r.get("import_path", "")
        if path:
            runtime_by_path[path] = r

    ast_paths = set(ast_by_path)
    runtime_paths = set(runtime_by_path)

    phantoms = sorted(ast_paths - runtime_paths)
    misses = sorted(runtime_paths - ast_paths)
    common = ast_paths & runtime_paths

    type_mismatches = []
    for path in sorted(common):
        ast_type = ast_by_path[path].get("type", "")
        runtime_type = runtime_by_path[path].get("type", "")
        if ast_type != runtime_type:
            type_mismatches.append(
                {
                    "import_path": path,
                    "ast_type": ast_type,
                    "runtime_type": runtime_type,
                }
            )

    # Print comparison table
    print("\n" + "=" * 60)
    print("Runtime vs AST Comparison")
    print("=" * 60)
    print(f"  AST classes:      {len(ast_paths)}")
    print(f"  Runtime classes:  {len(runtime_paths)}")
    print(f"  In common:        {len(common)}")
    print(f"  AST phantoms:     {len(phantoms)} (in AST, not runtime)")
    print(f"  AST misses:       {len(misses)} (in runtime, not AST)")
    print(f"  Type mismatches:  {len(type_mismatches)}")

    if phantoms:
        print(f"\nAST Phantoms ({len(phantoms)}):")
        for p in phantoms[:20]:
            ast_type = ast_by_path[p].get("type", "?")
            print(f"  [{ast_type}] {p}")
        if len(phantoms) > 20:
            print(f"  ... and {len(phantoms) - 20} more")

    if misses:
        print(f"\nAST Misses ({len(misses)}):")
        for p in misses[:20]:
            rt_type = runtime_by_path[p].get("type", "?")
            print(f"  [{rt_type}] {p}")
        if len(misses) > 20:
            print(f"  ... and {len(misses) - 20} more")

    if type_mismatches:
        print(f"\nType Mismatches ({len(type_mismatches)}):")
        for m in type_mismatches[:20]:
            print(f"  {m['import_path']}: AST={m['ast_type']} Runtime={m['runtime_type']}")
        if len(type_mismatches) > 20:
            print(f"  ... and {len(type_mismatches) - 20} more")

    print("=" * 60)

    return {
        "ast_phantoms": len(phantoms),
        "ast_misses": len(misses),
        "type_mismatches": len(type_mismatches),
        "phantom_paths": phantoms,
        "miss_paths": misses,
        "mismatch_details": type_mismatches,
    }


def _extract_params_from_modules(
    modules: list[dict],
) -> tuple[dict[str, dict[str, dict]], dict[str, str], int, int, int]:
    """Extract parameters from a list of module dicts.

    Returns (provider_classes, provider_names, total_processed, total_failed, total_params).
    """
    provider_classes: dict[str, dict[str, dict]] = defaultdict(dict)
    provider_names: dict[str, str] = {}
    total_processed = 0
    total_failed = 0
    total_params = 0

    for i, module in enumerate(modules, 1):
        import_path = module.get("import_path", "")
        provider_id = module.get("provider_id", "")
        provider_name = module.get("provider_name", module.get("provider_id", ""))
        module_type = module.get("type", "")
        class_name = module.get("name", "")

        if not import_path or not provider_id:
            continue

        provider_names[provider_id] = provider_name

        obj = import_symbol(import_path)
        if obj is None:
            total_failed += 1
            continue

        try:
            mro, params = extract_params(obj)
        except Exception as e:
            print(f"  ERROR extracting params for {import_path}: {e}")
            total_failed += 1
            continue

        provider_classes[provider_id][import_path] = {
            "name": class_name,
            "type": module_type,
            "mro": mro,
            "parameters": params,
        }

        total_processed += 1
        total_params += len(params)

        if i % 100 == 0:
            print(f"  Processed {i}/{len(modules)} modules...")

    return provider_classes, provider_names, total_processed, total_failed, total_params


def _write_parameter_files(
    provider_classes: dict[str, dict[str, dict]],
    provider_names: dict[str, str],
    provider_versions: dict[str, str],
    generated_at: str,
) -> None:
    """Write per-provider parameter JSON files."""
    for output_dir in OUTPUT_DIRS:
        if not output_dir.parent.exists():
            continue

        written = 0
        for pid, classes in provider_classes.items():
            version = provider_versions.get(pid)
            if not version:
                print(f"  WARN: no version found for {pid}, skipping")
                continue

            version_dir = output_dir / "versions" / pid / version
            version_dir.mkdir(parents=True, exist_ok=True)

            provider_data = validate_provider_parameters(
                {
                    "provider_id": pid,
                    "provider_name": provider_names.get(pid, pid),
                    "version": version,
                    "generated_at": generated_at,
                    "classes": classes,
                }
            )
            with open(version_dir / "parameters.json", "w") as f:
                json.dump(provider_data, f, separators=(",", ":"))
            written += 1

        print(f"Wrote {written} provider parameter files to {output_dir}/versions/")


def _fetch_inventories(
    provider_ids: set[str],
    provider_yamls: dict[str, dict],
) -> dict[str, dict[str, str]]:
    """Fetch Sphinx inventory files in parallel for all providers."""
    package_names: dict[str, str] = {}
    for pid in provider_ids:
        py = provider_yamls.get(pid, {})
        package_names[pid] = py.get("package-name", f"apache-airflow-providers-{pid}")

    def _fetch_and_parse(pid: str) -> tuple[str, dict[str, str] | None]:
        inv_path = fetch_provider_inventory(package_names[pid])
        if inv_path:
            try:
                return pid, read_inventory(inv_path)
            except Exception as e:
                print(f"    Warning: Could not parse inventory for {pid}: {e}")
                return pid, None
        return pid, None

    inventories: dict[str, dict[str, str]] = {}
    with concurrent.futures.ThreadPoolExecutor(max_workers=8) as executor:
        futures = {executor.submit(_fetch_and_parse, pid): pid for pid in provider_ids}
        for future in concurrent.futures.as_completed(futures):
            pid, inv = future.result()
            if inv:
                inventories[pid] = inv

    return inventories


def main():
    parser = argparse.ArgumentParser(description="Extract provider parameters and modules")
    parser.add_argument(
        "--provider",
        default=None,
        help=(
            "Only process these provider ID(s) (space-separated, e.g. 'amazon google'). "
            "Writes a partial modules.json containing only the requested providers; "
            "merge_registry_data.py replaces those providers' entries in the global catalog "
            "while preserving everyone else."
        ),
    )
    parser.add_argument(
        "--providers-json",
        default=None,
        help="Path to providers.json (overrides default search paths).",
    )
    args = parser.parse_args()

    print("Airflow Registry Parameter & Module Extractor")
    print("=" * 50)

    if args.providers_json:
        providers_json_path = Path(args.providers_json)
    else:
        providers_json_path = find_json(PROVIDERS_JSON_CANDIDATES, "providers.json")
    with open(providers_json_path) as f:
        providers_data = json.load(f)

    provider_versions: dict[str, str] = {}
    for p in providers_data.get("providers", []):
        provider_versions[p["id"]] = p["version"]

    generated_at = datetime.now(timezone.utc).isoformat()
    _main_discover(
        provider_versions,
        generated_at,
        requested_providers=_parse_requested_providers(args.provider),
    )

    print("\nDone!")


def _parse_requested_providers(provider_arg: str | None) -> set[str] | None:
    """Parse --provider argument into a set of provider IDs.

    Accepts a space-separated string (matching extract_metadata.py and
    extract_connections.py). Returns None when the argument is empty so
    callers can distinguish "all providers" from "explicit empty set".
    """
    if not provider_arg:
        return None
    return {pid.strip() for pid in provider_arg.split() if pid.strip()}


def _main_discover(
    provider_versions: dict[str, str],
    generated_at: str,
    requested_providers: set[str] | None = None,
) -> None:
    """Runtime discovery: find classes from provider.yaml files, produce modules.json and parameters.

    When ``requested_providers`` is set, only those providers are scanned and the resulting
    modules.json is partial (covers only the requested providers). ``merge_registry_data.py``
    handles incremental merges by replacing entries for provider IDs present in the new
    modules.json while preserving all others, so partial output is safe.
    """
    provider_yaml_paths = sorted(PROVIDERS_DIR.rglob("provider.yaml"))
    print(f"Found {len(provider_yaml_paths)} provider.yaml files")

    base_classes = load_base_classes()
    print(f"Loaded {len(base_classes)} base classes: {', '.join(sorted(base_classes))}")

    resumable_mixin = load_resumable_job_mixin()

    # Load all provider.yaml data and map provider_id -> yaml dict / path
    provider_yamls_by_id: dict[str, dict] = {}
    provider_paths_by_id: dict[str, Path] = {}
    for yaml_path in provider_yaml_paths:
        with open(yaml_path) as f:
            py = yaml.safe_load(f)
        pid = py.get("package-name", "").replace("apache-airflow-providers-", "")
        if pid:
            provider_yamls_by_id[pid] = py
            provider_paths_by_id[pid] = yaml_path

    # Filter to requested provider(s) if specified
    if requested_providers:
        missing = requested_providers - set(provider_paths_by_id)
        if missing:
            print(f"ERROR: provider(s) not found in provider.yaml files: {sorted(missing)}")
            sys.exit(1)
        provider_paths_by_id = {pid: provider_paths_by_id[pid] for pid in requested_providers}
        provider_yamls_by_id = {pid: provider_yamls_by_id[pid] for pid in requested_providers}
        print(f"Filtering to provider(s): {', '.join(sorted(requested_providers))}")

    # Fetch Sphinx inventories in parallel
    print("Fetching Sphinx inventory files ...")
    inventories = _fetch_inventories(set(provider_yamls_by_id), provider_yamls_by_id)
    print(f"  {len(inventories)}/{len(provider_yamls_by_id)} inventories loaded")

    all_discovered: list[dict] = []
    providers_seen: set[str] = set()

    for pid, yaml_path in sorted(provider_paths_by_id.items()):
        version = provider_versions.get(pid, "")
        discovered = discover_classes_from_provider(
            yaml_path,
            base_classes,
            resumable_mixin,
            inventory=inventories.get(pid),
            version=version,
        )
        all_discovered.extend(discovered)
        for d in discovered:
            providers_seen.add(d["provider_id"])

    print(f"\nDiscovered {len(all_discovered)} classes from {len(providers_seen)} providers")

    # Deduplicate by ID
    seen_ids: set[str] = set()
    unique_modules: list[dict] = []
    for m in all_discovered:
        mid = m["id"]
        if mid not in seen_ids:
            seen_ids.add(mid)
            unique_modules.append(m)
    all_discovered = unique_modules
    print(f"Deduplicated to {len(all_discovered)} unique modules")

    # Write modules.json. In --provider mode this is partial (covers only the
    # requested providers); merge_registry_data.py drives module replacement
    # off the provider IDs present in this file, so non-requested providers'
    # entries are preserved untouched in the global catalog.
    modules_json = validate_modules_catalog({"modules": all_discovered})
    scope_label = (
        f"partial, providers: {', '.join(sorted(requested_providers))}" if requested_providers else "full"
    )
    output_dirs = [SCRIPT_DIR, AIRFLOW_ROOT / "registry" / "src" / "_data"]
    for out_dir in output_dirs:
        if not out_dir.parent.exists():
            continue
        out_dir.mkdir(parents=True, exist_ok=True)
        with open(out_dir / "modules.json", "w") as f:
            json.dump(modules_json, f, indent=2)
        print(f"Wrote {len(all_discovered)} modules ({scope_label}) to {out_dir / 'modules.json'}")

    # Write runtime_modules.json (debug/stats file). Only meaningful for full
    # builds; skip in --provider mode since it would only show partial stats.
    if not requested_providers:
        runtime_output = {
            "generated_at": generated_at,
            "discovery_method": "runtime",
            "stats": {
                "total_classes": len(all_discovered),
                "total_providers": len(providers_seen),
            },
            "classes": all_discovered,
        }
        runtime_json_path = SCRIPT_DIR / "runtime_modules.json"
        with open(runtime_json_path, "w") as f:
            json.dump(runtime_output, f, indent=2)
        print(f"Wrote {runtime_json_path}")

    # Extract parameters
    print("\nExtracting parameters from runtime-discovered classes...")
    provider_classes, provider_names, total_processed, total_failed, total_params = (
        _extract_params_from_modules(all_discovered)
    )

    print(f"\nProcessed {total_processed} classes, {total_failed} failed imports")
    print(f"Extracted {total_params} total parameters")
    print(f"Across {len(provider_classes)} providers")

    _write_parameter_files(provider_classes, provider_names, provider_versions, generated_at)


if __name__ == "__main__":
    sys.exit(main() or 0)