"""
The ``cortex_agent_manager`` module provides a unified interface for managing
Cortex agents with RAI-powered tools.
"""
from typing import Callable, List, Tuple, Optional

from snowflake import snowpark

from relationalai.util.docutils import include_in_docs
from relationalai.util.error import exc, RAIException
from relationalai.util.tracing import span

from .api.agent import CortexAgentApi
from .chat import CortexAgentChat
from .cortex_tool_resources import CortexToolResources
from .deployment_config import DeploymentConfig, DeploymentStatus
from .api.client import http_client

__include_in_docs__ = True


@include_in_docs
class CortexAgentManager:
    """Deploy and manage Cortex AI agents powered by RAI semantic models.

    Handles the full lifecycle of a Snowflake Cortex agent that powers a
    Snowflake Intelligence (SI) experience. Non-technical SI users interact
    with the agent through the SI UI; this class is for the deployer/admin
    who creates and maintains the agent behind it.

    Managed Snowflake resources
    ---------------------------
    - **Stage** (``rai_sprocs`` by default): Stores sproc code and
      dependencies. Created/dropped automatically when ``manage_stage=True``.
    - **Stored procedures**: Four CALLER'S RIGHTS sprocs that power the
      agent's RAI tools: ``RAI_DISCOVER_MODELS``, ``RAI_VERBALIZE_MODEL``,
      ``RAI_EXPLAIN_CONCEPT``, and ``RAI_QUERY_MODEL`` (PREVIEW).
    - **Cortex agent**: The agent registered in the target database/schema,
      visible in the Snowflake Intelligence UI.

    Required privileges
    -------------------
    Both the deployer and SI users require the ``rai_developer`` role
    (created during RAI Native App installation).

    The deployer's role (the session's active role) must have:

    - ``CREATE STAGE`` on the target schema (if ``manage_stage=True``).
    - ``CREATE PROCEDURE`` on the target schema.
    - ``CREATE AGENT`` on the target schema.
    - Database role ``snowflake.cortex_user``.
    - Database role ``snowflake.pypi_repository_user``.
    - Application role ``snowflake.ai_observability_events_lookup``.
    - ``USAGE`` on the target database and schema.

    SI end-users need:

    - ``USAGE`` on the warehouse specified in ``DeploymentConfig``.
    - Database role ``snowflake.cortex_user``.
    - ``USAGE`` on the database and schema containing the agent.
    - ``SELECT`` on tables accessed by the model.
    - ``EXECUTE`` on the stored procedures.

    Parameters
    ----------
    session : snowpark.Session
        Authenticated ``snowpark.Session`` with access to the target database
        and schema.
    config : DeploymentConfig
        Configuration specifying where and how to deploy the agent.

    Examples
    --------

    >>> session = create_config().get_session(SnowflakeConnection)
    >>> manager = CortexAgentManager(
    ...     session=session,
    ...     config=DeploymentConfig(
    ...         agent_name="EXAMPLE_CORTEX",
    ...         database="EXAMPLE",
    ...         schema="CORTEX",
    ...         warehouse="DEMOWAREHOUSE",
    ...     ),
    ... )
    >>> def init_tools(model):
    ...     init_model(model)
    ...     return ToolRegistry().add(
    ...         model=model, description="Customers and orders")
    >>> manager.deploy(
    ...     init_tools=init_tools, imports=discover_imports())
    """

    def __init__(
        self,
        session: snowpark.Session,
        config: DeploymentConfig,
    ):
        self._config = config

        # Convert connection to session and configure it
        self._session = session
        try:
            self._session.use_database(config.database)
        except Exception as e:
            exc(
                "InvalidDatabase",
                f"Database '{config.database}' does not exist or is not accessible: {e}",
                [
                    "Check the 'database' field in your DeploymentConfig. "
                    "Ensure the database exists and your role has USAGE on it.\n\n"
                    "Example:\n"
                    "  config = DeploymentConfig(\n"
                    "      database='MY_DB',  # must already exist\n"
                    "      schema='MY_SCHEMA',\n"
                    "      agent_name='MY_ASSISTANT',\n"
                    "  )"
                ],
            )
        try:
            self._session.use_schema(config.schema)
        except Exception as e:
            exc(
                "InvalidSchema",
                f"Schema '{config.database}.{config.schema}' does not exist or is not accessible: {e}",
                [
                    "Check the 'schema' field in your DeploymentConfig. "
                    "Ensure the schema exists and your role has USAGE on it.\n\n"
                    "Example:\n"
                    "  config = DeploymentConfig(\n"
                    "      database='MY_DB',\n"
                    "      schema='MY_SCHEMA',  # must already exist\n"
                    "      agent_name='MY_ASSISTANT',\n"
                    "  )"
                ],
            )

        # Create HTTP client for Cortex API
        self._conn = session.connection
        self._http_client = http_client(self._conn)

        # Initialize internal components
        self._resources = CortexToolResources(
            session=self._session,
            stage_name=config.stage_name,
            manage_stage=config.manage_stage,
        )
        self._api = CortexAgentApi(self._http_client)

        # Track agent specification
        self._agent_spec: Optional[CortexAgentApi.AgentSpec] = None

    @include_in_docs
    def deploy(
        self,
        init_tools: Callable,  # Callable[[Model], ToolRegistry]
        imports: List[Tuple[str, str]],
        base_agent_spec: Optional[CortexAgentApi.AgentSpec] = None,
        extra_packages: Optional[List[str]] = None,
    ) -> CortexAgentApi.AgentSpec:
        """Deploy the agent and all required Snowflake resources.

        Creates a stage (if ``manage_stage=True``), registers RAI tool stored
        procedures, and deploys the Cortex agent. On failure, attempts cleanup
        of partially deployed resources before raising.

        Parameters
        ----------
        init_tools : Callable
            Callable accepting a ``Model`` and returning a ``ToolRegistry``.
            The framework creates a ``Model`` (with the sproc's session) and
            passes it to this callable. The callable should initialize
            concepts/properties on the model and return a configured
            ``ToolRegistry``.
        imports : list of tuple of (str, str)
            Project code to upload to Snowflake. Use ``discover_imports()``
            for automatic recursive discovery. Format: list of
            ``(source_path, module_name)`` tuples.
        base_agent_spec : CortexAgentApi.AgentSpec, optional
            Base agent spec to extend with RAI tools. Default: ``None``
            (creates a new spec).
        extra_packages : list of str, optional
            Additional PyPI packages for the sproc environment.
            ``relationalai`` is included automatically.

        Returns
        -------
        CortexAgentApi.AgentSpec
            The deployed agent specification.

        Raises
        ------
        RAIException
            If deployment fails at any stage.
        """
        with span("cortex.deploy"):
            try:
                # Phase 1: Create Snowflake resources and build agent spec
                orchestration_config: Optional[CortexAgentApi.OrchestrationConfig] = None
                if base_agent_spec is None:
                    # Build orchestration config if budget parameters are specified
                    if self._config.budget_seconds is not None or self._config.budget_tokens is not None:
                        budget = CortexAgentApi.BudgetConfig(
                            seconds=self._config.budget_seconds,
                            tokens=self._config.budget_tokens
                        )
                        orchestration_config = CortexAgentApi.OrchestrationConfig(budget=budget)

                    base_agent_spec = CortexAgentApi.AgentSpec.from_dict({
                        "name": self._config.agent_name,
                        "models": {
                            "orchestration": self._config.llm,
                        }
                    })
                    assert base_agent_spec is not None

                    # Set orchestration config if we built one
                    if orchestration_config is not None:
                        base_agent_spec.orchestration = orchestration_config

                self._agent_spec = self._resources.create(
                    init_tools=init_tools,
                    imports=imports,
                    config=self._config,
                    agent_spec=base_agent_spec,
                    extra_packages=extra_packages,
                )

                # Phase 2: Deploy to Cortex
                with span("cortex.deploy.create_agent"):
                    self._api.create(
                        self._config.database,
                        self._config.schema,
                        self._agent_spec
                    )

                # Phase 3: Validate deployment
                with span("cortex.deploy.validate"):
                    result = self._api.describe(
                        self._config.database,
                        self._config.schema,
                        self._config.agent_name
                    )

                    if result.agent_spec is None:
                        exc("ValidationFailed", "Could not parse Agent description")

                return result.agent_spec

            except RAIException:
                raise
            except Exception as e:
                exc("DeploymentFailed", f"Deployment failed: {e}")

    @include_in_docs
    def update(
        self,
        init_tools: Callable,
        imports: List[Tuple[str, str]],
        extra_packages: Optional[List[str]] = None,
    ) -> CortexAgentApi.AgentSpec:
        """Update the agent's tools without recreating the agent.

        Recreates stored procedures with updated tool definitions and updates
        the agent configuration in Cortex. The agent name and identity are
        preserved; existing SI conversations are unaffected.

        Parameters
        ----------
        init_tools : Callable
            Updated callable accepting a ``Model`` and returning a
            ``ToolRegistry``. Same contract as ``deploy()``.
        imports : list of tuple of (str, str)
            Updated project code. Use ``discover_imports()``.
        extra_packages : list of str, optional
            Updated PyPI package list.

        Returns
        -------
        CortexAgentApi.AgentSpec
            Updated agent specification.

        Raises
        ------
        RAIException
            If agent is not deployed (call ``deploy()`` first).
        """
        with span("cortex.update"):
            if not self.status().fully_deployed():
                exc("NotDeployed", "Agent not deployed. Call deploy() first.")

            # Get current agent spec
            current = self._api.describe(
                self._config.database,
                self._config.schema,
                self._config.agent_name
            )

            # Recreate resources with updated tools
            updated_spec = self._resources.create(
                init_tools=init_tools,
                imports=imports,
                config=self._config,
                agent_spec=current.agent_spec,
                extra_packages=extra_packages,
            )

            # Update agent via API
            self._api.update(
                self._config.database,
                self._config.schema,
                self._config.agent_name,
                updated_spec
            )

            self._agent_spec = updated_spec
            return updated_spec

    @include_in_docs
    def chat(self) -> CortexAgentChat:
        """Get a programmatic chat interface for the deployed agent.

        Returns a ``CortexAgentChat`` that maintains conversation state across
        multiple ``send()`` calls. This is primarily useful for testing and
        automation -- SI users interact through the Snowflake Intelligence UI.

        Returns
        -------
        CortexAgentChat
            Chat interface instance.

        Raises
        ------
        RAIException
            If agent is not deployed (call ``deploy()`` first).
        """
        if not self.status().fully_deployed():
            exc("NotDeployed", "Agent not deployed. Call deploy() first.")

        return CortexAgentChat(
            client=self._http_client,
            db=self._config.database,
            schema=self._config.schema,
            name=self._config.agent_name,
        )

    @include_in_docs
    def status(self) -> DeploymentStatus:
        """Check deployment status of all agent components.

        Queries Snowflake and Cortex to verify existence of the agent, stage,
        and all stored procedures.

        Returns
        -------
        DeploymentStatus
            Use ``fully_deployed()``, ``partially_deployed()``, or ``clean()``
            for quick checks.
        """
        # Check Snowflake resources
        resource_status = self._resources.exists()

        # Check Cortex agent
        agent_exists = False
        try:
            self._api.describe(
                self._config.database,
                self._config.schema,
                self._config.agent_name
            )
            agent_exists = True
        except Exception:
            pass

        return DeploymentStatus(
            agent_exists=agent_exists,
            stage_exists=resource_status.stage_exists,
            sprocs=resource_status.sprocs,
        )

    @include_in_docs
    def cleanup(self, silent: bool = True) -> None:
        """Remove all agent resources from Snowflake and Cortex.

        Deletes the Cortex agent, all RAI tool stored procedures, and the
        stage (if ``manage_stage=True``). SI conversation history is
        permanently lost.

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

        Returns
        -------
        None
        """
        with span("cortex.cleanup"):
            # Delete Cortex agent first
            try:
                self._api.delete(
                    self._config.database,
                    self._config.schema,
                    self._config.agent_name,
                    silent=silent
                )
            except Exception:
                if not silent:
                    raise

            # Delete Snowflake resources
            self._resources.cleanup(silent=silent)

            # Close HTTP client to release connection pool resources
            try:
                self._http_client.close()
            except Exception:
                if not silent:
                    raise

            self._agent_spec = None

    @property
    def config(self) -> DeploymentConfig:
        """The ``DeploymentConfig`` used to initialize this manager."""
        return self._config

    @property
    def agent_spec(self) -> Optional[CortexAgentApi.AgentSpec]:
        """The deployed ``AgentSpec``, or ``None`` if not yet deployed."""
        return self._agent_spec

    @property
    def is_deployed(self) -> bool:
        """``True`` if fully deployed. Queries Snowflake."""
        return self.status().fully_deployed()
