"""DataShield workflow — PII classification and masking policy.

Reuses:
    - ``src.datashield.unified`` (shield_project, shield_table, shield_deploy)
    - ``src.datashield.snowflake_generator`` (masking policy DDL)
"""

from __future__ import annotations

from typing import Any, Dict, List, Optional

from databridge_core.sdk.core.context import SDKContext
from databridge_core.sdk.core.exceptions import NotImplementedStubError
from databridge_core.sdk.core.types import CompiledArtifact, Result
from databridge_core.sdk.workflows.base import BaseWorkflow
from databridge_core.sdk.workflows.result import WorkflowResult
from databridge_core.sdk.workflows.step import workflow_step


class DataShieldWorkflow(BaseWorkflow):
    """Classify columns for PII and generate masking policies.

    Full lifecycle: project creation, auto-classification of
    columns, masking policy generation, and Snowflake deployment.

    Reuses:
        - ``src.datashield.unified`` (shield_project, shield_table, shield_deploy)
        - ``src.datashield.snowflake_generator``
    """

    name = "quality_governance.datashield"
    category = "quality_governance"
    description = "PII classification and masking policy generation"
    required_params = ["project_name"]
    optional_params = {
        "tables": [],
        "auto_classify": True,
        "generate_ddl": True,
    }

    def validate(self, params: Dict[str, Any]) -> Dict[str, Any]:
        self._check_required(params)
        return self._apply_defaults(params)

    def execute(
        self, params: Dict[str, Any], ctx: Optional[SDKContext] = None
    ) -> Result:
        steps = [self._create_project(params, ctx)]
        if params.get("auto_classify") and params.get("tables"):
            steps.append(self._classify(params, ctx))
        if params.get("generate_ddl"):
            steps.append(self._generate_ddl(params, ctx))
        # Graph integration
        if ctx and ctx.graphs.available:
            steps.append(self._graph_record_lineage(params, ctx))
            steps.append(self._graph_index_vectors(params, ctx))
            steps.append(self._graph_ingest_kb(params, ctx))
        return WorkflowResult.from_steps(self.name, steps)

    def compile_to_snowflake(
        self, params: Dict[str, Any], ctx: Optional[SDKContext] = None
    ) -> List[CompiledArtifact]:
        # REUSE: src.datashield.snowflake_generator to produce
        #        CREATE MASKING POLICY DDL and ALTER TABLE ... SET MASKING POLICY
        raise NotImplementedStubError(
            self.name,
            "compile_to_snowflake",
            "src.datashield.snowflake_generator — "
            "CREATE MASKING POLICY + ALTER TABLE SET MASKING POLICY",
        )

    @workflow_step("create_project", "Create DataShield project")
    def _create_project(self, params: Dict[str, Any], ctx: Optional[SDKContext] = None):
        # REUSE: src.datashield.unified._project_create(
        #     settings, name=params["project_name"]
        # )
        raise NotImplementedStubError(
            self.name,
            "create_project",
            "src.datashield.unified._project_create(settings, name, passphrase)",
        )

    @workflow_step("classify", "Auto-classify columns for PII")
    def _classify(self, params: Dict[str, Any], ctx: Optional[SDKContext] = None):
        # REUSE: For each table in params["tables"]:
        #   src.datashield.unified._table_auto_classify(settings, project_id, table_name)
        raise NotImplementedStubError(
            self.name,
            "classify",
            "src.datashield.unified — shield_table action='auto_classify'",
        )

    @workflow_step("generate_ddl", "Generate masking policy DDL")
    def _generate_ddl(self, params: Dict[str, Any], ctx: Optional[SDKContext] = None):
        # REUSE: src.datashield.unified._deploy_generate_ddl(settings, project_id)
        raise NotImplementedStubError(
            self.name,
            "generate_ddl",
            "src.datashield.unified — shield_deploy action='generate_ddl'",
        )

    # --- Graph integration ---

    @workflow_step("graph_record_lineage", "Record RULE and EXPORT nodes in lineage")
    def _graph_record_lineage(self, params: Dict[str, Any], ctx: Optional[SDKContext] = None):
        # ctx.graphs.record_lineage_node("sdk", f"datashield:{project}", "EXTERNAL",
        #     description="DataShield masking policy")
        raise NotImplementedStubError(
            self.name,
            "graph_record_lineage",
            "ctx.graphs.record_lineage_node — RULE + EXPORT nodes for masking policies",
        )

    @workflow_step("graph_index_vectors", "Index PII classifications in vector store")
    def _graph_index_vectors(self, params: Dict[str, Any], ctx: Optional[SDKContext] = None):
        # ctx.graphs.index_content(
        #     doc_id=f"shield_{project}", content=classification_summary,
        #     source_type="pii_classification", metadata={...},
        # )
        raise NotImplementedStubError(
            self.name,
            "graph_index_vectors",
            "ctx.graphs.index_content(doc_id, content, 'pii_classification', metadata)",
        )

    @workflow_step("graph_ingest_kb", "Ingest masking rules into KB graph")
    def _graph_ingest_kb(self, params: Dict[str, Any], ctx: Optional[SDKContext] = None):
        # ctx.graphs.ingest_kb_nodes(
        #     nodes=[{"id": policy_id, "type": "RULE", "properties": {"pii_type": ...}}],
        # )
        raise NotImplementedStubError(
            self.name,
            "graph_ingest_kb",
            "ctx.graphs.ingest_kb_nodes(nodes) — masking rules as KB nodes",
        )
