"""Masking Policy compiler.

Reuses:
    - ``src.datashield.snowflake_generator``
    - ``src.datashield.service.ShieldService``
"""

from __future__ import annotations

import logging
from typing import Any, Dict, List, Optional

from databridge_core.sdk.compiler.base import BaseCompiler
from databridge_core.sdk.core.exceptions import CompilationError
from databridge_core.sdk.core.types import CompiledArtifact

logger = logging.getLogger(__name__)


class MaskingPolicyCompiler(BaseCompiler):
    """Compile DataShield classifications into Snowflake shielded views.

    Generates UDF bundle DDL and per-table shielded view DDL based on
    PII classifications from a DataShield project.

    Reuses:
        - ``src.datashield.snowflake_generator`` for DDL generation
        - ``src.datashield.service.ShieldService`` for project data
    """

    object_type = "MASKING POLICY"

    def compile(
        self,
        workflow_name: str,
        params: Dict[str, Any],
        target_database: Optional[str] = None,
        target_schema: Optional[str] = None,
    ) -> List[CompiledArtifact]:
        """Compile into Masking Policy / Shielded View DDL.

        Required params:
            project_id: DataShield project ID.

        Optional params:
            key_ref: Key reference for UDF calls (defaults to project key_alias).
            policy_name_prefix: Prefix for generated object names.
            data_dir: Data directory (default 'data').
        """
        # Lazy imports to avoid hard dependencies
        try:
            from src.datashield.snowflake_generator import (
                generate_udfs,
                generate_view_ddl,
            )
            from src.datashield.service import ShieldService
        except ImportError as exc:
            raise CompilationError(
                f"MaskingPolicyCompiler requires DataShield modules: {exc}"
            ) from exc

        project_id = params.get("project_id")
        if not project_id:
            raise CompilationError(
                "MaskingPolicyCompiler requires 'project_id' param"
            )

        key_ref = params.get("key_ref")
        policy_name_prefix = params.get("policy_name_prefix", "")
        data_dir = params.get("data_dir", "data")

        # Load DataShield project
        service = ShieldService(data_dir=data_dir)
        project = service.get_project(project_id)

        if project is None:
            raise CompilationError(
                f"MaskingPolicyCompiler: DataShield project '{project_id}' not found"
            )

        # Use project key_alias as default key_ref
        if not key_ref:
            key_ref = getattr(project, "key_alias", project_id)

        artifacts: List[CompiledArtifact] = []

        # Artifact 1: UDF bundle (scramble functions)
        try:
            udf_ddl = generate_udfs()
        except Exception as exc:
            raise CompilationError(
                f"MaskingPolicyCompiler: UDF generation failed: {exc}"
            ) from exc

        udf_name = f"{policy_name_prefix}DATASHIELD_UDFS" if policy_name_prefix else "DATASHIELD_UDFS"
        udf_fqn = self._fqn(udf_name, target_database, target_schema)

        udf_artifact = CompiledArtifact(
            artifact_type="masking_policy",
            object_name=udf_fqn,
            ddl=self._header(udf_fqn) + udf_ddl,
            dependencies=[],
            source_workflow=workflow_name,
            metadata={
                "sub_type": "udf_bundle",
                "project_id": project_id,
            },
        )
        artifacts.append(udf_artifact)

        # Artifacts 2..N: one shielded view per table
        tables = getattr(project, "tables", []) or []
        for table_config in tables:
            try:
                view_ddl = generate_view_ddl(table_config, key_ref)
            except Exception as exc:
                raise CompilationError(
                    f"MaskingPolicyCompiler: view DDL generation failed "
                    f"for table '{getattr(table_config, 'table_name', '?')}': {exc}"
                ) from exc

            table_name = getattr(table_config, "table_name", "UNKNOWN")
            db = getattr(table_config, "database", target_database)
            schema = getattr(table_config, "schema_name", target_schema)

            view_name = f"{policy_name_prefix}VW_SHIELDED_{table_name}" if policy_name_prefix else f"VW_SHIELDED_{table_name}"
            view_fqn = self._fqn(view_name, db, schema)

            # Source table FQN for dependency tracking
            source_fqn = self._fqn(table_name, db, schema)

            view_artifact = CompiledArtifact(
                artifact_type="masking_policy",
                object_name=view_fqn,
                ddl=self._header(view_fqn) + view_ddl,
                dependencies=[udf_fqn, source_fqn],
                source_workflow=workflow_name,
                metadata={
                    "sub_type": "shielded_view",
                    "source_table": source_fqn,
                    "project_id": project_id,
                    "key_ref": key_ref,
                },
            )
            artifacts.append(view_artifact)

        logger.info(
            "MaskingPolicyCompiler produced %d artifacts "
            "(1 UDF bundle + %d shielded views) for project '%s'",
            len(artifacts),
            len(artifacts) - 1,
            project_id,
        )
        return artifacts
