"""
The ``deployment_config`` module provides configuration and status types for
Cortex agent deployments.
"""
from dataclasses import dataclass
from typing import Optional, Dict

from relationalai.util.docutils import include_in_docs

__include_in_docs__ = True


@include_in_docs
@dataclass
class DeploymentConfig:
    """Configuration for Cortex agent deployment to Snowflake.

    Parameters
    ----------
    database : str
        Snowflake database where agent will be deployed.
    schema : str
        Snowflake schema where agent will be deployed. To enable interaction
        through the Snowflake Intelligence UI, deploy to the schema configured
        for your account (typically AGENTS).
    agent_name : str
        Unique name for the Cortex agent within the schema.
    model_name : str, optional
        Name for the ``Model`` instance created inside each stored procedure.
        Default: same as ``agent_name``.
    warehouse : str, optional
        Snowflake warehouse the Cortex agent will use when executing RAI tools.
        If provided, SI users need USAGE on this warehouse. If not provided,
        tools use the caller's session warehouse. Stored procedures are created
        with CALLER'S RIGHTS.
    stage_name : str
        Name of the Snowflake stage for storing sproc dependencies.
        Default: ``"rai_sprocs"``.
    manage_stage : bool
        If ``True``, automatically create/drop the stage during deploy/cleanup.
        Set to ``False`` if using a pre-existing stage. Default: ``True``.
    llm : str
        Language model to use for agent orchestration. Must be a model available
        in Snowflake Cortex. Default: ``"claude-sonnet-4-5"``.
    query_timeout_s : int
        Timeout in seconds for stored procedure execution. Default: 300.
    external_access_integration : str
        Snowflake external access integration required at sproc runtime for
        RAI API access. Established by the default RAI installation.
        Default: ``"S3_RAI_INTERNAL_BUCKET_EGRESS_INTEGRATION"``.
    artifact_repository : str
        Snowflake artifact repository required at sproc runtime for Python
        package resolution. Established by the default RAI installation.
        Default: ``"snowflake.snowpark.pypi_shared_repository"``.
    budget_seconds : int, optional
        Time budget in seconds for agent execution. If specified, the agent
        will terminate when this limit is hit. Can be used with
        ``budget_tokens``.
    budget_tokens : int, optional
        Token budget for model consumption. If specified, the agent will
        terminate when this limit is hit. Can be used with ``budget_seconds``.
    allow_preview : bool
        Allow Preview capabilities (e.g., queries). Default: ``False``.

    Examples
    --------

    >>> config = DeploymentConfig(
    ...     database="APPS",
    ...     schema="AUSTIN",
    ...     agent_name="jaffle_assistant",
    ...     warehouse="COMPUTE_WH",
    ...     llm="claude-sonnet-4-5",
    ... )
    """
    database: str
    schema: str
    agent_name: str
    model_name: Optional[str] = None
    warehouse: Optional[str] = None
    stage_name: str = "rai_sprocs"
    manage_stage: bool = True
    llm: str = "claude-sonnet-4-5"
    query_timeout_s: int = 300
    external_access_integration: str = "S3_RAI_INTERNAL_BUCKET_EGRESS_INTEGRATION"
    artifact_repository: str = "snowflake.snowpark.pypi_shared_repository"
    budget_seconds: Optional[int] = None
    budget_tokens: Optional[int] = None
    allow_preview: bool = False

    def __post_init__(self):
        if self.model_name is None:
            self.model_name = self.agent_name


@include_in_docs
@dataclass
class DeploymentStatus:
    """Status information for a Cortex agent deployment.

    Returned by ``CortexAgentManager.status()`` to indicate which components exist.

    Attributes
    ----------
    agent_exists : bool
        ``True`` if the Cortex agent is registered.
    stage_exists : bool
        ``True`` if the Snowflake stage exists.
    sprocs : dict of str to bool
        Dictionary mapping stored procedure names to existence booleans.
        Keys are the four RAI tool names: ``RAI_DISCOVER_MODELS``,
        ``RAI_VERBALIZE_MODEL``, ``RAI_EXPLAIN_CONCEPT``, ``RAI_QUERY_MODEL``.

    Examples
    --------

    >>> status = manager.status()
    >>> status.fully_deployed()
    True
    """
    agent_exists: bool
    stage_exists: bool
    sprocs: Dict[str, bool]

    @include_in_docs
    def fully_deployed(self) -> bool:
        """Check if all components are deployed and ready.

        Returns
        -------
        bool
            ``True`` if agent, stage, and all stored procedures exist.
        """
        return self.agent_exists and self.stage_exists and all(self.sprocs.values())

    @include_in_docs
    def partially_deployed(self) -> bool:
        """Check if deployment is incomplete.

        Returns
        -------
        bool
            ``True`` if at least one component exists but not all.
            ``False`` if either all components exist or none exist.

        Notes
        -----
        Partial deployments usually indicate a failed deployment that needs
        cleanup before retrying.
        """
        any_exist = self.agent_exists or self.stage_exists or any(self.sprocs.values())
        return any_exist and not self.fully_deployed()

    @include_in_docs
    def clean(self) -> bool:
        """Check if all components have been removed.

        Returns
        -------
        bool
            ``True`` if no components exist (ready for fresh deployment).
        """
        return not self.fully_deployed() and not self.partially_deployed()
