"""
The ``cortex_tool_resources`` module provides low-level management of Snowflake
resources (stages, stored procedures) for Cortex agent RAI tools.
"""
import json
import logging
from dataclasses import dataclass
from typing import Callable, List, Tuple, Optional, NamedTuple, Union, cast

import pandas
from snowflake import snowpark
from snowflake.snowpark.types import StringType, IntegerType, MapType, VariantType, DataType

from relationalai.util.error import exc, warn
from relationalai.util.tracing import span
from .api.agent import CortexAgentApi
from .deployment_config import DeploymentConfig
from .tool import ToolRegistry, Capability


class SprocSignature(NamedTuple):
    """Stored procedure signature for registration and cleanup."""
    name: str
    input_types: List[DataType]
    return_type: DataType

    def signature_string(self) -> str:
        """Generate SQL signature string for DROP PROCEDURE."""
        if not self.input_types:
            return f"{self.name}()"
        type_names = []
        for t in self.input_types:
            if isinstance(t, StringType):
                type_names.append("VARCHAR")
            elif isinstance(t, IntegerType):
                type_names.append("INTEGER")
            else:
                type_names.append(str(t))
        return f"{self.name}({', '.join(type_names)})"


@dataclass
class ToolResourceStatus:
    """Status information of Snowflake resources (stage, sprocs) for agent tools.

    Attributes
    ----------
    stage_exists : bool
        ``True`` if the Snowflake stage exists.
    sprocs : dict
        Dictionary mapping sproc names to existence booleans.
    """
    stage_exists: bool
    sprocs: dict  # sproc_name -> exists

    @property
    def all_exist(self) -> bool:
        """``True`` if stage and all sprocs exist."""
        return self.stage_exists and all(self.sprocs.values())

    @property
    def any_exist(self) -> bool:
        """``True`` if any resources exist."""
        return self.stage_exists or any(self.sprocs.values())


class CortexToolResources:
    """Low-level manager for the Snowflake resources behind a Cortex agent's RAI tools.

    Creates and tears down the stage and four stored procedures
    (``RAI_DISCOVER_MODELS``, ``RAI_VERBALIZE_MODEL``, ``RAI_EXPLAIN_CONCEPT``,
    ``RAI_QUERY_MODEL``) that power the agent. Sprocs are registered with
    CALLER'S RIGHTS, so SI users execute them under their own role.

    Most deployers should use ``CortexAgentManager`` instead -- it wraps this
    class and also handles the Cortex agent lifecycle. Use this class directly
    only when you need fine-grained control over resource creation.

    Parameters
    ----------
    session : snowpark.Session
        Snowpark session with database and schema already set.
    stage_name : str
        Stage for sproc code storage. Default: ``"rai_sprocs"``.
    manage_stage : bool
        Create/drop the stage automatically. Default: ``True``.
    """

    # Sproc name constants
    RAI_DISCOVER_MODELS = "RAI_DISCOVER_MODELS"
    RAI_VERBALIZE_MODEL = "RAI_VERBALIZE_MODEL"
    RAI_EXPLAIN_CONCEPT = "RAI_EXPLAIN_CONCEPT"
    RAI_QUERY_MODEL = "RAI_QUERY_MODEL"

    # Sproc signatures (shared between create and cleanup)
    # Dictionary keyed by sproc name constants for easy lookup
    SPROC_SIGNATURES = {
        RAI_DISCOVER_MODELS: SprocSignature(RAI_DISCOVER_MODELS, [], MapType()),
        RAI_VERBALIZE_MODEL: SprocSignature(RAI_VERBALIZE_MODEL, [IntegerType()], StringType()),
        RAI_EXPLAIN_CONCEPT: SprocSignature(RAI_EXPLAIN_CONCEPT, [IntegerType(), StringType()], StringType()),
        RAI_QUERY_MODEL: SprocSignature(RAI_QUERY_MODEL, [IntegerType(), StringType()], VariantType()),
    }

    def __init__(
        self,
        session: snowpark.Session,
        stage_name: str = "rai_sprocs",
        manage_stage: bool = True,
    ):
        self._session = session
        self._stage_name = stage_name
        self._manage_stage = manage_stage

        # Extract current database and schema from session
        self._database = session.get_current_database()
        self._schema = session.get_current_schema()

    def create(
        self,
        init_tools: Callable,  # Callable[[Model], ToolRegistry]
        imports: List[Tuple[str, str]],
        config: DeploymentConfig,
        agent_spec: Optional[CortexAgentApi.AgentSpec] = None,
        extra_packages: Optional[List[str]] = None,
    ) -> CortexAgentApi.AgentSpec:
        """Register RAI tool sprocs and return an agent spec with tool configuration.

        Creates the stage (if ``manage_stage``), registers all four RAI sprocs,
        and merges tool definitions into the agent spec.

        Parameters
        ----------
        init_tools : Callable
            Callable accepting a ``Model`` and returning a ``ToolRegistry``.
            The framework creates the ``Model`` (with the sproc's session) and
            passes it in. The callable should initialize concepts/properties
            on the model and return a configured ``ToolRegistry``.
        imports : list of tuple of (str, str)
            ``(source_path, module_name)`` tuples for sproc dependencies.
        config : DeploymentConfig
            Configuration with model_name, warehouse, timeouts, integrations,
            and preview settings.
        agent_spec : CortexAgentApi.AgentSpec, optional
            Base spec to merge RAI tools into. Default: new empty spec.
        extra_packages : list of str, optional
            Additional PyPI packages for sprocs.

        Returns
        -------
        CortexAgentApi.AgentSpec
            Agent spec with RAI tools configured.

        Raises
        ------
        RAIException
            If Preview capabilities are detected without
            ``allow_preview=True``, or if stage creation fails.
        """
        with span("cortex.resources.create"):
            model_name = config.model_name if config.model_name else config.agent_name
            warehouse = config.warehouse
            query_timeout_s = config.query_timeout_s
            external_access_integration = config.external_access_integration
            artifact_repository = config.artifact_repository
            allow_preview = config.allow_preview

            if warehouse:
                if isinstance(warehouse, str):
                    warehouse = warehouse.upper()
                else:
                    exc(
                        "InvalidWarehouse",
                        f"warehouse must be a string if given, got {type(warehouse)}.",
                    )

            # Validate init_tools callable and result
            if not callable(init_tools):
                exc(
                    "InvalidInitTools",
                    f"init_tools must be callable, got {type(init_tools)}.",
                )

            try:
                from relationalai.semantics import Model
                validation_model = Model.all_models[0] if len(Model.all_models) == 1 else Model(model_name)
                registry = init_tools(validation_model)
            except Exception as e:
                exc("InitToolsException",
                    f"init_tools() raised an exception during verification: {e}\n"
                    f"Please ensure the function returns a valid ToolRegistry.")

            if not isinstance(registry, ToolRegistry):
                exc(
                    "InvalidInitTools",
                    f"init_tools() must return a ToolRegistry instance, got {type(registry)}.",
                )

            if not registry.has_tools():
                exc(
                    "NoToolsRegistered",
                    "init_tools() returned ToolRegistry with no tools registered.",
                    [
                        "Ensure your init_tools function registers a tool via ToolRegistry.add().\n\n"
                        "Example:\n"
                        "  def init_tools(model):\n"
                        "      return ToolRegistry().add(\n"
                        "          model=model,\n"
                        "          description='...'\n"
                        "      )"
                    ],
                )

            # Check for Preview capabilities before deployment
            preview_capabilities = self._check_preview_capabilities_cached(registry)
            if preview_capabilities and not allow_preview:
                cap_list = ", ".join(preview_capabilities)
                exc(
                    "PreviewCapabilities",
                    f"Preview capabilities detected in deployment: {cap_list}",
                    [
                        "Preview capabilities are not Generally Available and may have limitations.\n"
                        "To deploy Preview capabilities, set allow_preview=True in your DeploymentConfig.\n\n"
                        "Example:\n"
                        "  config = DeploymentConfig(\n"
                        "      ...,\n"
                        "      allow_preview=True  # Acknowledge Preview capabilities\n"
                        "  )"
                    ],
                )

            if preview_capabilities and allow_preview:
                cap_list = ", ".join(preview_capabilities)
                warn("PreviewCapabilities",
                     f"Deploying with Preview capabilities: {cap_list}. "
                     "These may change in future releases.")

            # Create or replace stage
            if self._manage_stage:
                with span("cortex.resources.create_stage"):
                    try:
                        self._session.sql(f"CREATE OR REPLACE STAGE {self._stage_name}").collect()
                    except Exception as e:
                        exc(
                            "StageCreationFailure",
                            f"Failed to create stage '{self._stage_name}': {e}",
                            [
                                f"Ensure you have CREATE STAGE privilege on schema {self._schema}. "
                                f"Or, you can rely on an existing stage by setting manage_stage=False."
                            ]
                        )

            # Get fully qualified schema for tool resources
            schema_fqn = self._session.get_fully_qualified_current_schema()

            # Helper to register sprocs with CALLER'S RIGHTS execution
            # This means sprocs execute with the warehouse and role of the calling session,
            # not the warehouse/role used during sproc creation
            def register_sproc(func: Callable, name: str, input_types, return_type):
                with span(f"cortex.resources.register_sproc.{name}"):
                    # Suppress Snowpark's "artifact_repository ... is experimental"
                    # log warning. Snowpark emits this via the snowflake.snowpark
                    # logger whenever artifact_repository is provided.
                    snowpark_logger = logging.getLogger("snowflake.snowpark")
                    previous_level = snowpark_logger.level
                    snowpark_logger.setLevel(logging.ERROR)
                    try:
                        import relationalai
                        self._session.sproc.register(
                            func,
                            name=name,
                            input_types=input_types,
                            return_type=return_type,
                            stage_location=self._stage_name,
                            is_permanent=True,
                            replace=True,
                            packages=[
                                "snowflake-snowpark-python",
                                "snowflake-telemetry-python",
                                f"relationalai=={relationalai.__version__}",
                                *(extra_packages if extra_packages else []),
                            ],
                            imports=cast(List[Union[str, Tuple[str, str]]], imports),
                            execute_as="caller",
                            artifact_repository=artifact_repository,
                            external_access_integrations=[external_access_integration],
                        )
                    except Exception as e:
                        exc(
                            "SprocRegistrationFailure",
                            f"Failed to register stored procedure '{name}': {e}",
                            [
                                f"Ensure your role has the following privileges:\n"
                                f"  GRANT CREATE PROCEDURE ON SCHEMA {self._schema} TO ROLE <your_role>;\n"
                                f"  GRANT DATABASE ROLE snowflake.pypi_repository_user TO ROLE <your_role>;\n\n"
                                f"Also check that the stage '{self._stage_name}' exists and is accessible."
                            ]
                        )
                    finally:
                        snowpark_logger.setLevel(previous_level)

            # Initialize or merge agent spec
            if agent_spec is None:
                agent_spec = CortexAgentApi.AgentSpec()

            agent_spec = agent_spec.merge(CortexAgentApi.AgentSpec(
                instructions=CortexAgentApi.Instructions(
                    response=self._build_response_instructions(registry),
                    orchestration=self._build_orchestration_instructions(registry),
                ),
                tools=[],
                tool_resources={},
            ))

            # Ensure tools and tool_resources are initialized after merge
            assert agent_spec.tools is not None, "tools should be initialized after merge"
            assert agent_spec.tool_resources is not None, "tool_resources should be initialized after merge"

            # Register RAI_DISCOVER_MODELS
            def rai_discover_models(session: snowpark.Session):
                try:
                    from relationalai.semantics import Model
                    model = Model(model_name)
                    registry = init_tools(model)
                    result = registry.discover()

                    if isinstance(result, dict):
                        has_queries = registry.has_capability(Capability.QUERY)
                        if not has_queries:
                            result['_note'] = (
                                'IMPORTANT: These RelationalAI (RAI) models support structure exploration only - data queries are not available. '
                                'RAI tools help users understand model structure and semantics. '
                                'Guide users to ask questions like: "What is the Customer concept?", "Explain the relationship between Orders and Customers", '
                                '"What are the key concepts in this model?", or "Show me the structure of the Order concept". '
                                'Make it clear these are RAI-powered capabilities focused on semantic understanding. '
                                'DO NOT suggest analytical questions like "show me top customers" or "calculate total revenue".'
                            )
                        else:
                            result['_note'] = (
                                'IMPORTANT: This model has a PREVIEW query capability with a limited catalog of specific pre-defined queries. '
                                'You can ONLY execute the exact queries listed in the model metadata. '
                                'For any analytical question, first check if it matches one of the available pre-defined queries. '
                                'If there is no matching pre-defined query, you CANNOT answer the question with data - '
                                'instead, guide users to understand the data model structure using RAI_VERBALIZE_MODEL and RAI_EXPLAIN_CONCEPT. '
                                'Frame this as: "This model provides specific curated analytical queries" - these are intentional limitations, not bugs.'
                            )

                    return result
                except Exception as e:
                    import traceback
                    error_msg = f"Error in rai_discover_models: {str(e)}\n{traceback.format_exc()}"
                    print(error_msg)
                    return {"error": error_msg, "models": []}

            sig = self.SPROC_SIGNATURES[self.RAI_DISCOVER_MODELS]
            register_sproc(rai_discover_models, sig.name, sig.input_types, sig.return_type)
            agent_spec.tools = [t for t in agent_spec.tools if t.tool_spec.name != sig.name]
            agent_spec.tools.append(CortexAgentApi.Tool(
                CortexAgentApi.ToolSpec(
                    type="generic",
                    name=sig.name,
                    description="Discover available data models in the RelationalAI registry. Returns models that can be explored to answer business questions.",
                    input_schema=None,
                )
            ))
            agent_spec.tool_resources[sig.name] = CortexAgentApi.ToolResourceSpec(
                execution_environment=CortexAgentApi.ExecutionEnvironment(
                    query_timeout=query_timeout_s,
                    type="warehouse",
                    warehouse=warehouse
                ),
                identifier=f"{schema_fqn}.{sig.name}",
                name=f"{sig.name}()",
                type="procedure",
            )

            # Register RAI_VERBALIZE_MODEL
            def rai_verbalize_model(session: snowpark.Session, param1: int):
                try:
                    from relationalai.semantics import Model
                    model = Model(model_name)
                    return init_tools(model).verbalize(param1)
                except Exception as e:
                    import traceback
                    error_msg = f"Error in rai_verbalize_model for model '{param1}': {str(e)}\n{traceback.format_exc()}"
                    print(error_msg)
                    return error_msg

            sig = self.SPROC_SIGNATURES[self.RAI_VERBALIZE_MODEL]
            register_sproc(rai_verbalize_model, sig.name, sig.input_types, sig.return_type)
            agent_spec.tools = [t for t in agent_spec.tools if t.tool_spec.name != sig.name]
            agent_spec.tools.append(CortexAgentApi.Tool(
                CortexAgentApi.ToolSpec(
                    type="generic",
                    name=sig.name,
                    description="Get detailed natural language descriptions of a data model from the RelationalAI registry. Use this RAI tool to understand model structure and present business-friendly explanations to users.",
                    input_schema={
                        "type": "object",
                        "properties": {
                            "arg1": {
                                "description": f"Numeric identifier obtained from {self.RAI_DISCOVER_MODELS}",
                                "type": "integer"
                            }
                        },
                        "required": ["arg1"]
                    },
                )
            ))
            agent_spec.tool_resources[sig.name] = CortexAgentApi.ToolResourceSpec(
                execution_environment=CortexAgentApi.ExecutionEnvironment(
                    query_timeout=query_timeout_s,
                    type="warehouse",
                    warehouse=warehouse
                ),
                identifier=f"{schema_fqn}.{sig.name}",
                name=f"{sig.name}(INTEGER)",
                type="procedure",
            )

            # Register RAI_EXPLAIN_CONCEPT
            def rai_explain_concept(session: snowpark.Session, param1: int, param2: str):
                try:
                    from relationalai.semantics import Model
                    model = Model(model_name)
                    return init_tools(model).explain(param1, param2)
                except Exception as e:
                    import traceback
                    error_msg = f"Error in rai_explain_concept for model '{param1}', concept '{param2}': {str(e)}\n{traceback.format_exc()}"
                    print(error_msg)
                    return error_msg

            sig = self.SPROC_SIGNATURES[self.RAI_EXPLAIN_CONCEPT]
            register_sproc(rai_explain_concept, sig.name, sig.input_types, sig.return_type)
            agent_spec.tools = [t for t in agent_spec.tools if t.tool_spec.name != sig.name]
            agent_spec.tools.append(CortexAgentApi.Tool(
                CortexAgentApi.ToolSpec(
                    type="generic",
                    name=sig.name,
                    description="Get explanations for specific concepts or pre-defined queries in a RelationalAI model. Pass a concept name or a query ID — both obtained from RAI_DISCOVER_MODELS — to receive the appropriate explanation.",
                    input_schema={
                        "type": "object",
                        "properties": {
                            "arg1": {
                                "description": f"Numeric model_id acquired from {self.RAI_DISCOVER_MODELS}",
                                "type": "integer"
                            },
                            "arg2": {
                                "description": f"A concept name or query ID, both obtained from {self.RAI_DISCOVER_MODELS}",
                                "type": "string"
                            }
                        },
                        "required": ["arg1", "arg2"]
                    },
                )
            ))
            agent_spec.tool_resources[sig.name] = CortexAgentApi.ToolResourceSpec(
                execution_environment=CortexAgentApi.ExecutionEnvironment(
                    query_timeout=query_timeout_s,
                    type="warehouse",
                    warehouse=warehouse
                ),
                identifier=f"{schema_fqn}.{sig.name}",
                name=f"{sig.name}(INTEGER, VARCHAR)",
                type="procedure",
            )

            # Register RAI_QUERY_MODEL
            def rai_query_model(session: snowpark.Session, param1: int, param2: str):
                import traceback
                try:
                    from relationalai.semantics import Model
                    model = Model(model_name)
                    registry = init_tools(model)
                except Exception as init_err:
                    error_trace = traceback.format_exc()
                    print(f"Error in rai_query_model: {str(init_err)}\n{error_trace}")
                    return [{"error": f"Query execution failed: {str(init_err)}", "model_id": param1,
                             "body_preview": param2[:100] if param2 else None, "traceback": error_trace[:500]}]

                model_id = param1
                body = param2

                if not registry.tool_has_capability(model_id, Capability.QUERY):
                    return [{
                        "error": "This model does not have query execution capability.",
                        "model_id": model_id,
                        "available_capabilities": ["discover", "verbalize", "explain"],
                        "guidance": "You can help users understand the data model structure and semantics, but cannot execute data queries."
                    }]

                try:
                    query_dict = json.loads(body)
                except json.JSONDecodeError as json_err:
                    return [{"error": f"Invalid JSON in body parameter: {str(json_err)}", "body_received": body[:100],
                             "expected_format": '{"id": "<query_id>", "args": {<params>}}'}]

                if not isinstance(query_dict, dict) or "id" not in query_dict:
                    # Retrieve available queries to guide the agent
                    available_queries = []
                    try:
                        tool = registry._tools[model_id]
                        query_meta = tool.queries()
                        if "queries" in query_meta:
                            for qid, qinfo in query_meta["queries"].items():
                                entry = {"id": qid}
                                if "args" in qinfo:
                                    entry["args"] = qinfo["args"]
                                if "description" in qinfo:
                                    entry["description"] = qinfo["description"]
                                available_queries.append(entry)
                    except Exception:
                        pass
                    return [{
                        "error": "Invalid query payload format.",
                        "expected_format": '{"id": "<query_id>", "args": {<params>}}',
                        "body_received": body[:200] if body else None,
                        "available_queries": available_queries,
                        "guidance": "The payload must be a JSON object with an 'id' key matching one of the available query IDs, and an 'args' key with the query parameters."
                    }]

                try:
                    result = registry.query(model_id, query_dict)
                except KeyError as key_err:
                    # Could be model not found or query ID not found — include available queries
                    available_queries = []
                    try:
                        tool = registry._tools[model_id]
                        query_meta = tool.queries()
                        if "queries" in query_meta:
                            for qid, qinfo in query_meta["queries"].items():
                                entry = {"id": qid}
                                if "args" in qinfo:
                                    entry["args"] = qinfo["args"]
                                if "description" in qinfo:
                                    entry["description"] = qinfo["description"]
                                available_queries.append(entry)
                    except (IndexError, KeyError):
                        try:
                            available_models = list(range(len(registry._tools)))
                        except Exception:
                            available_models = ["Unable to retrieve available models"]
                        return [{"error": f"Model '{model_id}' not found. Available models: {available_models}"}]
                    return [{
                        "error": f"Query not found: {str(key_err)}",
                        "model_id": model_id,
                        "available_queries": available_queries,
                        "expected_format": '{"id": "<query_id>", "args": {<params>}}'
                    }]
                except Exception as exec_err:
                    error_trace = traceback.format_exc()
                    print(f"Error in rai_query_model: {str(exec_err)}\n{error_trace}")
                    return [{"error": f"Query execution failed: {str(exec_err)}", "model_id": model_id,
                             "body_preview": body[:100] if body else None, "traceback": error_trace[:500]}]

                if isinstance(result, pandas.DataFrame):
                    return result.to_dict(orient="records")
                return result

            sig = self.SPROC_SIGNATURES[self.RAI_QUERY_MODEL]
            register_sproc(rai_query_model, sig.name, sig.input_types, sig.return_type)
            agent_spec.tools = [t for t in agent_spec.tools if t.tool_spec.name != sig.name]

            query_description = "Execute a pre-defined RAI query from the model's query catalog. ONLY call this for questions that match specific pre-defined queries discovered from RAI_DISCOVER_MODELS."
            if Capability.QUERY.is_preview:
                query_description += " [PREVIEW - Limited catalog of curated queries. Cannot create custom queries.]"

            if not registry.has_capability(Capability.QUERY):
                query_description += " [NO MODELS WITH QUERY CAPABILITY - This tool will return an error if called]"

            agent_spec.tools.append(CortexAgentApi.Tool(
                CortexAgentApi.ToolSpec(
                    type="generic",
                    name=sig.name,
                    description=query_description,
                    input_schema={
                        "type": "object",
                        "properties": {
                            "arg1": {
                                "description": f"Numeric model_id obtained from {self.RAI_DISCOVER_MODELS}",
                                "type": "integer"
                            },
                            "arg2": {
                                "description": (
                                    f'JSON string with the query payload. '
                                    f'Must contain "id" (the query ID from {self.RAI_DISCOVER_MODELS}) '
                                    f'and "args" (an object mapping parameter names to values). '
                                    f'Example: {{"id": "my_query", "args": {{"limit": 10}}}}. '
                                    f'For queries with no parameters, use empty args: {{"id": "my_query", "args": {{}}}}.'
                                ),
                                "type": "string"
                            }
                        },
                        "required": ["arg1", "arg2"]
                    },
                )
            ))
            agent_spec.tool_resources[sig.name] = CortexAgentApi.ToolResourceSpec(
                execution_environment=CortexAgentApi.ExecutionEnvironment(
                    query_timeout=query_timeout_s,
                    type="warehouse",
                    warehouse=warehouse
                ),
                identifier=f"{schema_fqn}.{sig.name}",
                name=f"{sig.name}(INTEGER, VARCHAR)",
                type="procedure",
            )

            return agent_spec

    def _check_preview_capabilities_cached(self, registry: ToolRegistry) -> List[str]:
        """Check if any Preview capabilities are configured in the tool registry."""
        preview_caps = []

        try:
            for capability in Capability:
                if capability.is_preview and registry.has_capability(capability):
                    preview_caps.append(capability.value)
        except Exception as e:
            warn("PreviewCheck", f"Could not check for Preview capabilities: {e}")

        return preview_caps

    def _build_response_instructions(self, registry: ToolRegistry) -> str:
        """Build response instructions based on available capabilities."""
        base = [
            "You are a professional assistant that uses tools to answer business questions from non-technical users.\n"
            "Your response must be based off the results of some tool. If none of the available tools can be used to answer a question, "
            "your response should indicate that and nothing else. Do not present general, non-tool based knowledge in the response. "
            "Focus more on clearly displaying the results, rather than providing exhaustive analysis. "
            "The user should be able to understand the results and ask follow-up questions.\n\n"
            "CRITICAL PRESENTATION RULES:\n"
            "- NEVER show Python code, SQL DDL, or any programming syntax to end users\n"
            "- NEVER display code snippets, function definitions, or implementation details\n"
            "- Translate all technical content to natural language descriptions that business users can understand\n"
            "- Express formulas using mathematical notation (e.g., \u00d7 for multiplication, \u00f7 for division) NOT Python operators (*, /)\n"
            "- Focus on WHAT the business logic does, not HOW it is implemented in code\n"
            "- Describe relationships, rules, and calculations in plain business language\n"
            "- If tool output contains code, summarize its meaning in natural language instead of showing the code\n\n"
            "BRANDING REQUIREMENTS:\n"
            "- When introducing capabilities, mention they are 'powered by RelationalAI' or reference 'RAI tools'\n"
            "- When explaining what models are available, note they are from 'the RelationalAI registry'\n"
            "- In refusal explanations, frame limitations as 'RAI tools focus on...' rather than generic 'available tools'\n"
            "- Natural mentions preferred: integrate 'RelationalAI' or 'RAI' into explanations organically\n\n"
        ]

        if registry.has_capability(Capability.QUERY):
            base.append(
                "QUERIES CAPABILITY - CRITICAL COMMUNICATION RULES:\n"
                "- When introducing query capabilities, ALWAYS clarify these are 'specific pre-defined queries' not arbitrary query execution\n"
                "- When listing available queries, frame them as 'These specific analytical queries are available' NOT as examples\n"
                "- For questions that don't match pre-defined queries, explain: 'This model has a limited set of pre-defined queries. Your question doesn't match any available queries.'\n"
                "- NEVER suggest you can answer arbitrary analytical questions - only the pre-defined queries listed in the catalog\n"
                "- Frame as intentional design: 'This model provides specific curated analytical queries' NOT as a limitation or gap\n\n"
            )

        if not registry.has_capability(Capability.QUERY):
            base.append(
                "IMPORTANT - When responding to onboarding prompts (like 'what can I ask about?' or 'help me get started'):\n"
                "- DO NOT suggest analytical queries like 'show me top customers' or 'calculate revenue' - you cannot answer these\n"
                "- ONLY guide users to ask questions about model STRUCTURE and SEMANTICS\n"
                "- Emphasize that you can explain relationships, definitions, and business rules, but NOT execute data queries\n"
                "- Make it crystal clear that data analysis and aggregation queries are not available\n\n"
            )

        return "".join(base)

    def _build_orchestration_instructions(self, registry: ToolRegistry) -> str:
        """Build orchestration instructions based on available capabilities."""
        base = ("Use RAI tools to understand data models, add context to answers, "
                "and ensure your answers are semantically correct. "
                "When presenting information to users, focus on semantic meaning for business users and avoid showing code snippets. "
                "When explaining capabilities to users, naturally mention that these are 'RAI tools' or 'RelationalAI-powered capabilities'. "
                "Frame model discovery as accessing 'the RelationalAI registry' or 'RAI models'. ")

        instructions = [base]

        if not registry.has_capability(Capability.QUERY):
            instructions.append(
                "CRITICAL: Query capability is NOT available for any models in the registry. "
                "Do NOT attempt to call RAI_QUERY_MODEL under any circumstances. "
                "You can ONLY discover models, verbalize their structure, and explain concepts. "
                "You CANNOT execute data queries, aggregations, or analytical operations. "
                "When explaining these limitations, position your capabilities as: 'RAI tools focus on helping you understand data model structure and semantics.' "
            )

        instructions.append(
            "Always discover RAI models first when beginning a conversation so you understand what you can answer questions about.\n\n"
        )

        if registry.has_capability(Capability.VERBALIZE):
            instructions.append(
                "Use RAI_VERBALIZE_MODEL to get detailed descriptions of data model structure, relationships, and semantics. "
            )

            if registry.has_capability(Capability.QUERY):
                instructions.append("This helps you understand the model's concepts and relationships.\n\n")
            else:
                instructions.append("Focus on explaining the model structure and relationships to users.\n\n")

        if registry.has_capability(Capability.EXPLAIN):
            instructions.append(
                "Use RAI_EXPLAIN_CONCEPT to understand business rules for a specific concept or to get a description of a pre-defined query. Both concept names and query IDs come from RAI_DISCOVER_MODELS.\n\n"
            )

        if registry.has_capability(Capability.QUERY):
            instructions.append(
                "Use RAI_QUERY_MODEL to execute ONLY the specific pre-defined analytical queries available in the model's query catalog. "
                "BEFORE attempting any query, check if the user's question matches one of the pre-defined queries discovered from RAI_DISCOVER_MODELS. "
                "If the question does NOT match a pre-defined query, do NOT call RAI_QUERY_MODEL. Instead, clearly explain that only specific pre-defined queries are available. "
                "Never attempt to construct custom queries - only use the exact pre-defined queries from the catalog.\n\n"
                "QUERY PAYLOAD FORMAT: The arg2 parameter must be a JSON string with this exact structure: "
                '{\"id\": \"<query_id>\", \"args\": {<params>}}. '
                "The 'id' must match a query ID from RAI_DISCOVER_MODELS. "
                "The 'args' must be an object mapping parameter names to values as described in the query metadata. "
                "For queries with no parameters, use empty args: "
                '{\"id\": \"<query_id>\", \"args\": {}}.\n\n'
            )

        return "".join(instructions)

    def cleanup(self, silent: bool = True) -> None:
        """Drop all RAI sprocs and the stage (if ``manage_stage``).

        Parameters
        ----------
        silent : bool
            Ignore errors for non-existent resources. Default: ``True``.

        Returns
        -------
        None
        """
        with span("cortex.resources.cleanup"):
            # Delete all sprocs with full signatures
            for sig in self.SPROC_SIGNATURES.values():
                try:
                    # Use full signature to handle overloaded procedures
                    self._session.sql(
                        f"DROP PROCEDURE IF EXISTS {self._database}.{self._schema}.{sig.signature_string()}"
                    ).collect()
                except Exception as e:
                    if not silent:
                        raise
                    warn("Cleanup", f"Failed to delete sproc {sig.name}: {e}")

            # Delete stage only if we manage it
            if self._manage_stage:
                try:
                    self._session.sql(
                        f"DROP STAGE IF EXISTS {self._database}.{self._schema}.{self._stage_name}"
                    ).collect()
                except Exception as e:
                    if not silent:
                        raise
                    warn("Cleanup", f"Failed to delete stage {self._stage_name}: {e}")

    def exists(self) -> ToolResourceStatus:
        """Check existence of the stage and all RAI sprocs.

        Returns
        -------
        ToolResourceStatus
            Status of each Snowflake resource.
        """
        # Check stage existence
        stage_exists = False
        try:
            result = self._session.sql(
                f"SHOW STAGES LIKE '{self._stage_name}' IN {self._database}.{self._schema}"
            ).collect()
            stage_exists = len(result) > 0
        except Exception:
            pass

        # Check sproc existence
        sproc_status = {}
        for sproc_name in self.SPROC_SIGNATURES.keys():
            exists = False
            try:
                result = self._session.sql(
                    f"SHOW PROCEDURES LIKE '{sproc_name}' IN {self._database}.{self._schema}"
                ).collect()
                exists = len(result) > 0
            except Exception:
                pass
            sproc_status[sproc_name] = exists

        return ToolResourceStatus(stage_exists=stage_exists, sprocs=sproc_status)
