"""YAML Pipeline Builder — slim coordinator with plugin architecture.

The builder orchestrates the pipeline from YAML to DirectedGraph:
1. Parse YAML → 2. Select environment → 3. Validate → 4. Preprocess → 5. Build graph

Plugins are extracted into focused modules:
- ``preprocessing/``: IncludePreprocessPlugin, EnvironmentVariablePlugin, TemplatePlugin
- ``plugins/``: MacroDefinitionPlugin, MacroEntityPlugin, NodeEntityPlugin
"""

from __future__ import annotations

from contextlib import contextmanager
from functools import lru_cache
from pathlib import Path
from typing import Any, Protocol

import yaml

from hexdag.kernel.domain.dag import DirectedGraph
from hexdag.kernel.exceptions import HexDAGError
from hexdag.kernel.logging import get_logger
from hexdag.kernel.pipeline_builder.pipeline_config import PipelineConfig

# Re-export extracted classes for backward compatibility.
# Existing code that imports from ``yaml_builder`` continues to work.
from hexdag.kernel.pipeline_builder.plugins.macro_definition import (  # noqa: F401
    MacroDefinitionPlugin,
)
from hexdag.kernel.pipeline_builder.plugins.macro_entity import MacroEntityPlugin  # noqa: F401
from hexdag.kernel.pipeline_builder.plugins.node_entity import NodeEntityPlugin  # noqa: F401
from hexdag.kernel.pipeline_builder.preprocessing.env_vars import (  # noqa: F401
    EnvironmentVariablePlugin,
)
from hexdag.kernel.pipeline_builder.preprocessing.include import (  # noqa: F401
    IncludePreprocessPlugin,
)
from hexdag.kernel.pipeline_builder.preprocessing.template import TemplatePlugin  # noqa: F401
from hexdag.kernel.pipeline_builder.yaml_validator import YamlValidator

logger = get_logger(__name__)


class YamlPipelineBuilderError(HexDAGError):
    """YAML pipeline building errors."""

    pass


# ============================================================================
# Plugin Protocols
# ============================================================================


class PreprocessPlugin(Protocol):
    """Plugin for preprocessing YAML before building (env vars, templating, etc.)."""

    def process(self, config: dict[str, Any]) -> dict[str, Any]:
        """Process entire config, returning modified version."""
        ...


class EntityPlugin(Protocol):
    """Plugin for building specific entity types (macros, nodes, etc.)."""

    def can_handle(self, node_config: dict[str, Any]) -> bool:
        """Return True if this plugin can handle the node config."""
        ...

    def build(
        self, node_config: dict[str, Any], builder: YamlPipelineBuilder, graph: DirectedGraph
    ) -> Any:
        """Build entity from config.

        Args
        -----
            node_config: The node configuration dictionary.
            builder: The YamlPipelineBuilder instance.
            graph: The DirectedGraph instance.

        Return NodeSpec or None if handled (e.g., macro merged into graph)."""
        ...


# ============================================================================
# Core Builder
# ============================================================================


class YamlPipelineBuilder:
    """YAML -> DirectedGraph builder with plugin support.

    Template Rendering Phases
    -------------------------
    hexDAG uses a multi-phase rendering strategy:

    **Phase 1** (build-time): ``${VAR}`` env var substitution via
    ``EnvironmentVariablePlugin``. Secret patterns (``*_API_KEY``, etc.)
    are deferred to Phase 3b.

    **Phase 2** (build-time): ``{{ var }}`` Jinja2 rendering via
    ``TemplatePlugin``. Node ``spec`` fields and pipeline ``outputs``
    are **skipped** — preserved for Phase 3.

    **Phase 3** (runtime): ``{{var}}`` rendering via ``PromptTemplate``
    when nodes execute with actual dependency outputs.

    **Phase 3b** (runtime): Deferred ``${VAR}`` resolution at adapter
    instantiation via ``_resolve_deferred_env_vars``.

    Workflow:
    1. Parse YAML
    2. Select environment
    3. Validate structure
    4. Run preprocessing plugins (Phase 1 → Phase 2)
    5. Build graph using entity plugins
    6. Extract pipeline config
    """

    def __init__(self, base_path: Path | None = None) -> None:
        """Initialize builder.

        Args:
            base_path: Base directory for resolving includes (default: cwd)
        """
        self.base_path = base_path or Path.cwd()
        self.validator = YamlValidator()

        # Plugins
        self.preprocess_plugins: list[PreprocessPlugin] = []
        self.entity_plugins: list[EntityPlugin] = []

        self._register_default_plugins()

    def _register_default_plugins(self) -> None:
        """Register default plugins for common use cases."""
        # Preprocessing plugins (run before building)
        self.preprocess_plugins.append(IncludePreprocessPlugin(base_path=self.base_path))
        self.preprocess_plugins.append(EnvironmentVariablePlugin())
        self.preprocess_plugins.append(TemplatePlugin())

        # Entity plugins (build specific entity types)
        self.entity_plugins.append(MacroDefinitionPlugin())  # Process Macro definitions first
        self.entity_plugins.append(MacroEntityPlugin())  # Then macro invocations
        self.entity_plugins.append(NodeEntityPlugin(self))  # Finally regular nodes

    @contextmanager
    def _temporary_base_path(self, new_base: Path) -> Any:
        """Context manager for temporarily changing base_path.

        Args:
            new_base: Temporary base path to use

        Yields:
            None
        """
        original_base = self.base_path
        self.base_path = new_base

        # Update include plugin base paths
        for plugin in self.preprocess_plugins:
            if isinstance(plugin, IncludePreprocessPlugin):
                plugin.base_path = new_base

        try:
            yield
        finally:
            # Always restore original state
            self.base_path = original_base
            for plugin in self.preprocess_plugins:
                if isinstance(plugin, IncludePreprocessPlugin):
                    plugin.base_path = original_base

    # --- Public API ---

    def build_from_yaml_file(
        self, yaml_path: str, use_cache: bool = True
    ) -> tuple[DirectedGraph, PipelineConfig]:
        """Build from YAML file.

        Args:
            yaml_path: Path to YAML file
            use_cache: Whether to use cached YAML parsing

        Returns:
            Tuple of (DirectedGraph, PipelineConfig)
        """
        yaml_file = Path(yaml_path)
        yaml_content = yaml_file.read_text(encoding="utf-8")

        # Use context manager to temporarily change base_path for relative includes
        with self._temporary_base_path(yaml_file.parent):
            return self.build_from_yaml_string(yaml_content, use_cache=use_cache)

    def build_from_yaml_string(
        self, yaml_content: str, use_cache: bool = True, environment: str | None = None
    ) -> tuple[DirectedGraph, PipelineConfig]:
        """Build DirectedGraph + PipelineConfig from YAML string."""
        # Step 1: Parse YAML
        documents = self._parse_yaml(yaml_content, use_cache=use_cache)

        # Step 2: Process ALL documents for macro definitions first
        #         This allows macros to be defined in multi-document YAML
        #         before the pipeline that uses them
        for doc in documents:
            if isinstance(doc, dict) and doc.get("kind") == "Macro":
                # Process includes for macro definitions (but skip template rendering)
                # Templates in macros should be preserved for expansion-time rendering
                processed_doc = doc
                for plugin in self.preprocess_plugins:
                    # Skip TemplatePlugin for Macro definitions - templates should be preserved
                    if not isinstance(plugin, TemplatePlugin):
                        processed_doc = plugin.process(processed_doc)
                # Validate and process macro definition
                processed_doc = self._validate_config(processed_doc)
                # Register macro (MacroDefinitionPlugin handles this)
                self._process_macro_definitions([processed_doc])

        # Step 3: Select pipeline environment
        config = self._select_environment(documents, environment)

        # Skip if selected document is a Macro (already processed above)
        if config.get("kind") == "Macro":
            # Return empty graph for macro-only documents
            logger.info("Document contains only macro definitions, no pipeline to build")
            return DirectedGraph(), PipelineConfig(
                ports={}, type_ports={}, policies={}, metadata={}, nodes=[]
            )

        # Step 4: Preprocess FIRST (includes must resolve before validation)
        for plugin in self.preprocess_plugins:
            config = plugin.process(config)

        # Step 4.5: Register user-defined aliases BEFORE validation
        # so that custom node kinds can pass validation
        self._register_aliases(config)

        # Step 4.6: Register custom sanitized types BEFORE validation
        # so that output_schema type names can pass validation
        self._register_custom_types(config)

        # Step 5: Validate structure (after includes are resolved)
        config = self._validate_config(config)

        # Step 6: Build graph
        graph = self._build_graph(config)

        # Step 7: Extract pipeline config
        pipeline_config = self._extract_pipeline_config(config)

        logger.info(
            "Built pipeline '{name}' with {nodes} nodes, {ports} ports, {policies} policies",
            name=pipeline_config.metadata.get("name", "unknown"),
            nodes=len(graph.nodes),
            ports=len(pipeline_config.ports),
            policies=len(pipeline_config.policies),
        )

        return graph, pipeline_config

    # --- Core Logic ---

    def _parse_yaml(self, yaml_content: str, use_cache: bool) -> list[dict[str, Any]]:
        """Parse YAML into list of documents."""
        if "---" in yaml_content:
            return list(yaml.safe_load_all(yaml_content))
        parsed = _parse_yaml_cached(yaml_content) if use_cache else yaml.safe_load(yaml_content)
        return [parsed]

    def _select_environment(
        self, documents: list[dict[str, Any]], environment: str | None
    ) -> dict[str, Any]:
        """Select document by environment name.

        Skips Macro definitions when selecting the Pipeline document.
        """
        # Filter out Macro definitions - they're processed separately
        pipeline_docs = [doc for doc in documents if doc.get("kind") != "Macro"]

        if not pipeline_docs:
            # No pipeline documents, return first macro (for macro-only files)
            return documents[0]

        if environment:
            for doc in pipeline_docs:
                if doc.get("metadata", {}).get("namespace") == environment:
                    logger.info(f"Selected environment '{environment}' from multi-document YAML")
                    return doc

            available_envs = [
                doc.get("metadata", {}).get("namespace", "default") for doc in pipeline_docs
            ]
            raise YamlPipelineBuilderError(
                f"Environment '{environment}' not found in YAML. "
                f"Available environments: {', '.join(available_envs)}"
            )

        if len(pipeline_docs) > 1:
            logger.warning(
                f"Multi-document YAML detected ({len(pipeline_docs)} pipeline documents) "
                "but no environment specified. "
                "Using first pipeline. Specify environment parameter to select specific config."
            )

        return pipeline_docs[0]

    def _validate_config(self, config: Any) -> dict[str, Any]:
        """Validate YAML structure."""
        if not isinstance(config, dict):
            raise YamlPipelineBuilderError(
                f"YAML document must be a dictionary, got {type(config).__name__}"
            )

        self._validate_manifest_format(config)

        result = self.validator.validate(config)
        if not result.is_valid:
            errors = "\n".join(f"  ERROR: {error}" for error in result.errors)
            raise YamlPipelineBuilderError(f"YAML validation failed:\n{errors}")

        for warning in result.warnings:
            logger.warning(f"YAML validation warning: {warning}")

        return config

    @staticmethod
    def _validate_manifest_format(config: dict[str, Any]) -> None:
        """Validate declarative manifest format."""
        if "kind" not in config:
            raise YamlPipelineBuilderError(
                "YAML must use declarative manifest format with 'kind' field. "
                "Example:\n"
                "apiVersion: v1\n"
                "kind: Pipeline\n"
                "metadata:\n"
                "  name: my-pipeline\n"
                "spec:\n"
                "  nodes: [...]"
            )

        # Macro definitions have different structure (no spec field)
        kind = config.get("kind")
        if kind == "Macro":
            # Macro has: metadata, parameters, nodes (no spec)
            if "metadata" not in config:
                raise YamlPipelineBuilderError("Macro definition must have 'metadata' field")
            if "nodes" not in config:
                raise YamlPipelineBuilderError("Macro definition must have 'nodes' field")
        else:
            # Pipeline and other kinds require spec
            if "spec" not in config:
                raise YamlPipelineBuilderError("Manifest YAML must have 'spec' field")

            if "metadata" not in config:
                raise YamlPipelineBuilderError("Manifest YAML must have 'metadata' field")

    def _process_macro_definitions(self, macro_configs: list[dict[str, Any]]) -> None:
        """Process macro definitions and register them.

        Parameters
        ----------
        macro_configs : list[dict[str, Any]]
            List of validated macro configuration dictionaries
        """
        # Create temporary graph (not used for macros, but required by plugin interface)
        temp_graph = DirectedGraph()

        # Find MacroDefinitionPlugin
        macro_plugin = next(
            (p for p in self.entity_plugins if isinstance(p, MacroDefinitionPlugin)), None
        )

        if macro_plugin is None:
            raise YamlPipelineBuilderError(
                "MacroDefinitionPlugin not found. Cannot process macro definitions."
            )

        for macro_config in macro_configs:
            # Wrap in a fake node config format expected by entity plugin
            node_config = macro_config
            macro_plugin.build(node_config, self, temp_graph)

    def _register_aliases(self, config: dict[str, Any]) -> None:
        """Register user-defined aliases from spec.aliases before validation.

        This allows custom node kinds to be used in YAML without requiring
        full module paths.

        Parameters
        ----------
        config : dict[str, Any]
            The YAML configuration dict
        """
        from hexdag.kernel.resolver import register_alias

        spec = config.get("spec", {})
        aliases = spec.get("aliases", {})

        for alias, full_path in aliases.items():
            register_alias(alias, full_path)
            logger.debug(f"Registered alias: {alias} -> {full_path}")

    def _register_custom_types(self, config: dict[str, Any]) -> None:
        """Register custom sanitized types from ``spec.custom_types`` before validation.

        This allows output_schema fields to use custom type names defined
        declaratively in the pipeline YAML.

        Parameters
        ----------
        config : dict[str, Any]
            The YAML configuration dict
        """
        from hexdag.kernel.validation.sanitized_types import register_type_from_config

        spec = config.get("spec", {})
        custom_types = spec.get("custom_types", {})

        for type_name, type_config in custom_types.items():
            register_type_from_config(type_name, type_config)
            logger.debug(f"Registered custom type: {type_name}")

    def _build_graph(self, config: dict[str, Any]) -> DirectedGraph:
        """Build DirectedGraph using entity plugins."""
        graph = DirectedGraph()
        spec = config.get("spec", {})
        nodes_list = spec.get("nodes", [])

        for node_config in nodes_list:
            # Find plugin that can handle this entity
            for plugin in self.entity_plugins:
                if plugin.can_handle(node_config):
                    result = plugin.build(node_config, self, graph)
                    if result is not None:
                        graph += result
                    break
            else:
                # No plugin handled it - error
                kind = node_config.get("kind", "unknown")
                raise YamlPipelineBuilderError(f"No plugin can handle kind: {kind}")

        return graph

    @staticmethod
    def _extract_pipeline_config(config: dict[str, Any]) -> PipelineConfig:
        """Extract PipelineConfig from YAML."""
        spec = config.get("spec", {})
        metadata = config.get("metadata", {})

        return PipelineConfig(
            ports=spec.get("ports", {}),
            type_ports=spec.get("type_ports", {}),
            policies=spec.get("policies", {}),
            metadata=metadata,
            nodes=spec.get("nodes", []),
        )


# ============================================================================
# Utilities
# ============================================================================


@lru_cache(maxsize=32)
def _parse_yaml_cached(yaml_content: str) -> Any:
    """Cached YAML parsing.

    Returns dict[str, Any] in practice, but yaml.safe_load returns Any.
    """
    return yaml.safe_load(yaml_content)
