"""Provide pipeline stub fragments and project files for Spark runtimes."""

from __future__ import annotations

from textwrap import dedent
from typing import Callable, Mapping

from .. import PipelineFile, PipelineProject, PipelineStub, register_pipeline_stub


def _coerce_hint(hints: Mapping[str, object], key: str) -> str:
    value = hints.get(key)
    if value is None:
        return ""
    return str(value).strip()


def _spark_readme(hints: Mapping[str, object]) -> str:
    runtime = _coerce_hint(hints, "spark_runtime")
    workspace = _coerce_hint(hints, "spark_workspace_url")
    profile = _coerce_hint(hints, "spark_workspace_profile")
    cluster = _coerce_hint(hints, "spark_cluster")

    wizard_hints = ["Wizard hints captured during export:"]
    if runtime:
        wizard_hints.append(f"- Runtime: {runtime}")
    if workspace:
        wizard_hints.append(f"- Workspace URL: {workspace}")
    if profile:
        wizard_hints.append(f"- CLI profile: {profile}")
    if cluster:
        wizard_hints.append(f"- Cluster reference: {cluster}")
    if len(wizard_hints) == 1:
        wizard_hints.append("- (no runtime-specific hints were recorded)")

    return (
        dedent(
            """
            # Spark pipeline example

            This directory contains a full PySpark pipeline that uses the dc43
            integration helpers exported by the setup wizard. The `main.py`
            entrypoint shows how to:

            - load service backends via ``bootstrap_pipeline.load_backends``
            - build a Spark session with ``build_spark_context``
            - read governed product inputs with automatic casting and
              validation
            - apply a simple transformation
            - publish results through the data-product governance helpers

            Update the placeholder constants in ``main.py`` before running the
            pipeline and remove or replace the demo transformation with your
            business logic.
            """
        ).strip()
        + "\n\n"
        + "\n".join(wizard_hints)
        + "\n"
    )


def _spark_main_script() -> str:
    return dedent(
        '''
        #!/usr/bin/env python3
        """Run the Spark pipeline example generated by the dc43 setup wizard."""

        from __future__ import annotations

        from bootstrap_pipeline import build_spark_context, load_backends
        from spark_pipeline import io as spark_io
        from spark_pipeline import transformations as spark_transformations


        DATA_PRODUCT_ID = "replace-with-data-product-id"
        INPUT_PORT = "replace-with-input-port"
        OUTPUT_PORT = "replace-with-output-port"


        def main() -> None:
            suite = load_backends()
            context = build_spark_context(app_name="dc43-spark-example")
            spark = context.get("spark")
            if spark is None:
                raise RuntimeError(
                    "build_spark_context did not return a Spark session. "
                    "Install pyspark and ensure the runtime is available."
                )

            print("[spark] Spark session initialised:", spark)

            governance_service = context.get(
                "governance_service", suite.governance
            )

            source_df, status = spark_io.read_data_product_input(
                spark=spark,
                governance_service=governance_service,
                data_product=DATA_PRODUCT_ID,
                input_port=INPUT_PORT,
            )
            if status is not None:
                print(
                    "[data_quality] Input status:",
                    status.status,
                    status.reason or "",
                )

            enriched_df = spark_transformations.build_output_frame(source_df)

            write_validation = spark_io.write_data_product_output(
                dataframe=enriched_df,
                governance_service=governance_service,
                output_port=OUTPUT_PORT,
                data_product=DATA_PRODUCT_ID,
            )

            print(
                "[done] Example pipeline completed. Replace placeholder identifiers",
                "before running in production.",
            )


        if __name__ == "__main__":
            main()
        '''
    ).strip() + "\n"


def _spark_io_module() -> str:
    return dedent(
        '''
        """Read/write helpers used by the Spark pipeline example."""

        from __future__ import annotations

        from typing import Optional, Tuple

        from dc43_integrations.spark.io import (
            GovernanceSparkReadRequest,
            GovernanceSparkWriteRequest,
            read_with_governance,
            write_with_governance,
        )
        from dc43_integrations.spark.violation_strategy import SplitWriteViolationStrategy
        from dc43_service_clients.data_products import (
            DataProductInputBinding,
            DataProductOutputBinding,
        )
        from dc43_service_clients.data_quality import ValidationResult
        from dc43_service_clients.governance.client.interface import GovernanceServiceClient
        from dc43_service_clients.governance.models import GovernanceReadContext, GovernanceWriteContext
        from pyspark.sql import DataFrame, SparkSession


        def read_data_product_input(
            *,
            spark: SparkSession,
            governance_service: GovernanceServiceClient,
            data_product: str,
            input_port: str,
        ) -> Tuple[DataFrame, Optional[ValidationResult]]:
            """Load a product input port enforced through governance."""

            request = GovernanceSparkReadRequest(
                context=GovernanceReadContext(
                    input_binding=DataProductInputBinding(
                        data_product=data_product,
                        port_name=input_port,
                    )
                ),
            )
            df, status = read_with_governance(
                spark,
                request,
                governance_service=governance_service,
                enforce=True,
                auto_cast=True,
                return_status=True,
            )
            return df, status


        def write_data_product_output(
            *,
            dataframe: DataFrame,
            governance_service: GovernanceServiceClient,
            output_port: str,
            data_product: str,
        ) -> ValidationResult:
            """Persist pipeline results relying solely on governance."""

            request = GovernanceSparkWriteRequest(
                context=GovernanceWriteContext(
                    output_binding=DataProductOutputBinding(
                        data_product=data_product,
                        port_name=output_port,
                    ),
                ),
                mode="overwrite",
            )

            validation, reject_status = write_with_governance(
                df=dataframe,
                request=request,
                governance_service=governance_service,
                enforce=True,
                auto_cast=True,
                return_status=True,
                violation_strategy=SplitWriteViolationStrategy(
                    valid_suffix="valid",
                    reject_suffix="reject",
                ),
            )

            if reject_status is not None:
                print(
                    "[write] Reject status:",
                    reject_status.status,
                    reject_status.reason or "",
                )

            return validation
        '''
    ).strip() + "\n"


def _spark_transformations_module() -> str:
    return dedent(
        '''
        """Demo transformations applied by the Spark example pipeline."""

        from __future__ import annotations

        from pyspark.sql import DataFrame, functions as F


        def build_output_frame(source: DataFrame) -> DataFrame:
            """Return a demo dataset derived from ``source``.

            Replace this logic with your business-specific transformations once the
            scaffold is in place.
            """

            enriched = source.withColumn(
                "processed_at",
                F.current_timestamp(),
            )
            if "status" in source.columns:
                enriched = enriched.withColumn(
                    "is_valid",
                    F.when(F.col("status") == F.lit("valid"), F.lit(True)).otherwise(False),
                )
            return enriched
        '''
    ).strip() + "\n"


def _spark_project(hints: Mapping[str, object]) -> PipelineProject:
    return PipelineProject(
        root="spark_pipeline",
        entrypoint="main.py",
        files=(
            PipelineFile(path="README.md", content=_spark_readme(hints)),
            PipelineFile(
                path="__init__.py",
                content='"""Spark pipeline example package."""\n',
            ),
            PipelineFile(path="main.py", content=_spark_main_script(), executable=True),
            PipelineFile(path="io.py", content=_spark_io_module()),
            PipelineFile(path="transformations.py", content=_spark_transformations_module()),
        ),
    )


def _spark_pipeline_stub(
    *,
    hints: Mapping[str, object],
    flags: Mapping[str, bool],
    json_literal: Callable[[object | None], str],
) -> PipelineStub:
    runtime_hint = json_literal(hints.get("spark_runtime"))
    workspace_hint = json_literal(hints.get("spark_workspace_url"))
    profile_hint = json_literal(hints.get("spark_workspace_profile"))
    cluster_hint = json_literal(hints.get("spark_cluster"))
    data_product_id = json_literal("replace-with-data-product-id")
    input_port = json_literal("replace-with-input-port")
    output_port = json_literal("replace-with-output-port")

    main_lines = (
        "    if integration == 'spark':",
        "        context = build_spark_context(app_name=\"dc43-pipeline-example\")",
        "        spark = context.get('spark')",
        "        if spark is None:",
        "            raise RuntimeError('Spark session was not initialised by build_spark_context')",
        "        print(\"[spark] Spark session initialised:\", spark)",
        f"        runtime_hint = {runtime_hint}",
        "        if runtime_hint:",
        "            print(\"[spark] Runtime configured in setup:\", runtime_hint)",
        f"        workspace_hint = {workspace_hint}",
        "        if workspace_hint:",
        "            print(\"[spark] Workspace URL:\", workspace_hint)",
        f"        profile_hint = {profile_hint}",
        "        if profile_hint:",
        "            print(\"[spark] CLI profile:\", profile_hint)",
        f"        cluster_hint = {cluster_hint}",
        "        if cluster_hint:",
        "            print(\"[spark] Cluster reference:\", cluster_hint)",
        "        governance_service = context.get('governance_service', governance_service)",
        f"        data_product_id = {data_product_id}",
        f"        input_port = {input_port}",
        f"        output_port = {output_port}",
        "        source_df, input_status = spark_io.read_data_product_input(",
        "            spark=spark,",
        "            governance_service=governance_service,",
        "            data_product=data_product_id,",
        "            input_port=input_port,",
        "        )",
        "        if input_status:",
        "            print(\"[spark] Input status:\", input_status.status, input_status.reason or \"\")",
        "        enriched_df = spark_transformations.build_output_frame(source_df)",
        "        write_validation = spark_io.write_data_product_output(",
        "            dataframe=enriched_df,",
        "            governance_service=governance_service,",
        "            output_port=output_port,",
        "            data_product=data_product_id,",
        "        )",
        "        print(\"[spark] Output validation status:\", write_validation.status)",
    )

    tail_lines = (
        "    if integration == 'spark':",
        "        print(\"[spark] Explore examples/spark_pipeline for modular helpers.\")",
    )

    return PipelineStub(
        bootstrap_imports=("build_spark_context",),
        additional_imports=(
            "from spark_pipeline import io as spark_io",
            "from spark_pipeline import transformations as spark_transformations",
        ),
        main_lines=main_lines,
        tail_lines=tail_lines,
        project=_spark_project(hints),
    )


register_pipeline_stub("spark", _spark_pipeline_stub)


__all__ = ["_spark_pipeline_stub"]

