"""Snowpark Python Stored Procedure compiler — NEW capability.

Generates CREATE PROCEDURE statements wrapping SDK workflow logic
as Snowpark Python stored procedures.
"""

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__)

_DEFAULT_PACKAGES = ["snowflake-snowpark-python", "pandas"]


class StoredProcCompiler(BaseCompiler):
    """Compile SDK workflows into Snowpark Python Stored Procedures.

    Generates CREATE PROCEDURE DDL with:
        - Python handler function
        - Package dependencies (snowflake-snowpark-python, pandas + custom)
        - Input/output parameter signatures
        - Stage references for code upload
    """

    object_type = "PROCEDURE"

    def compile(
        self,
        workflow_name: str,
        params: Dict[str, Any],
        target_database: Optional[str] = None,
        target_schema: Optional[str] = None,
    ) -> List[CompiledArtifact]:
        """Compile into Snowpark Python Stored Procedure DDL.

        Required params:
            procedure_name: Name of the stored procedure.
            handler: Python handler function name (e.g. 'main').

        Optional params:
            packages: Additional Python packages (list of str).
            stage: Stage name for code upload.
            imports: Files to import from stage (list of str).
            runtime_version: Python version (default '3.11').
            input_params: List of {name, type} dicts for procedure args.
            return_type: Return type (default 'VARIANT').
            execute_as: CALLER or OWNER (default 'CALLER').
            inline_code: Python source to embed inline.
            comment: DDL comment.
        """
        procedure_name = params.get("procedure_name")
        handler = params.get("handler")

        if not procedure_name:
            raise CompilationError(
                "StoredProcCompiler requires 'procedure_name' param"
            )
        if not handler:
            raise CompilationError("StoredProcCompiler requires 'handler' param")

        packages = list(params.get("packages") or [])
        stage = params.get("stage")
        imports = params.get("imports") or []
        runtime_version = params.get("runtime_version", "3.11")
        input_params = params.get("input_params") or []
        return_type = params.get("return_type", "VARIANT")
        execute_as = params.get("execute_as", "CALLER").upper()
        inline_code = params.get("inline_code")
        comment = params.get("comment")

        fqn = self._fqn(procedure_name, target_database, target_schema)

        # Build parameter signature
        param_parts = []
        for p in input_params:
            pname = p.get("name", "")
            ptype = p.get("type", "VARCHAR")
            if not pname:
                raise CompilationError(
                    f"StoredProcCompiler: input_params entry missing 'name': {p}"
                )
            param_parts.append(f"{pname} {ptype}")
        param_sig = ", ".join(param_parts)

        # Merge packages: defaults + user-specified (deduplicated)
        all_packages = list(_DEFAULT_PACKAGES)
        for pkg in packages:
            if pkg not in all_packages:
                all_packages.append(pkg)
        packages_str = ", ".join(f"'{p}'" for p in all_packages)

        # Build DDL
        ddl_parts = [self._header(fqn)]
        ddl_parts.append(f"CREATE OR REPLACE PROCEDURE {fqn}({param_sig})")
        ddl_parts.append(f"  RETURNS {return_type}")
        ddl_parts.append("  LANGUAGE PYTHON")
        ddl_parts.append(f"  RUNTIME_VERSION = '{runtime_version}'")
        ddl_parts.append(f"  PACKAGES = ({packages_str})")
        ddl_parts.append(f"  HANDLER = '{handler}'")
        ddl_parts.append(f"  EXECUTE AS {execute_as}")

        if imports:
            if stage:
                import_refs = ", ".join(
                    f"'@{stage}/{f}'" for f in imports
                )
            else:
                import_refs = ", ".join(f"'{f}'" for f in imports)
            ddl_parts.append(f"  IMPORTS = ({import_refs})")

        if comment:
            ddl_parts.append(f"  COMMENT = '{comment}'")

        if inline_code:
            ddl_parts.append(f"AS\n$$\n{inline_code}\n$$;")
        else:
            # External handler — no inline code block
            ddl_parts[-1] = ddl_parts[-1] + ";"

        ddl = "\n".join(ddl_parts)

        # Dependencies: stage if provided
        dependencies: List[str] = []
        if stage:
            stage_fqn = self._fqn(stage, target_database, target_schema)
            dependencies.append(stage_fqn)

        artifact = CompiledArtifact(
            artifact_type="stored_proc",
            object_name=fqn,
            ddl=ddl,
            dependencies=dependencies,
            source_workflow=workflow_name,
            metadata={
                "handler": handler,
                "runtime_version": runtime_version,
                "packages": all_packages,
                "execute_as": execute_as,
                "param_count": len(input_params),
                "has_inline_code": bool(inline_code),
            },
        )

        logger.info("StoredProcCompiler produced 1 artifact: %s", fqn)
        return [artifact]
