"""
The ``agent`` module provides a client for the Snowflake Cortex Agents REST API.
"""
import dataclasses
from dataclasses import dataclass
from typing import List, Optional, Dict, Any, Tuple
import json
import copy

import httpx

from relationalai.util.docutils import include_in_docs
from relationalai.util.error import Diagnostic, RAIException, exc

__include_in_docs__ = True


@include_in_docs
class APIResponseError(RAIException):
    """Custom exception for API response errors."""
    def __init__(self, message):
        diag = Diagnostic("APIResponseError", str(message), severity="error")
        super().__init__(diag)
        self.message = message


@include_in_docs
class CortexAgentApi:
    """Client for interacting with the Snowflake Cortex Agents REST API.

    Provides methods for creating, updating, describing, listing, deleting,
    and running Cortex agents. Requires the ``snowflake.cortex_user``
    database role. All requests timeout after 15 minutes.

    Parameters
    ----------
    client : httpx.Client
        Configured ``httpx.Client`` with authentication headers.
    """
    def __init__(self, client: httpx.Client):
        self._c = client

    @dataclass
    class Profile:
        """Agent profile configuration for visual display properties.

        Attributes
        ----------
        avatar : str or None
            URL or identifier for the agent's avatar image.
        color : str or None
            Color code for agent's visual theme.
        display_name : str or None
            Human-readable name shown in the UI.
        """
        avatar: Optional[str] = None
        color: Optional[str] = None
        display_name: Optional[str] = None

    @dataclass
    class Instructions:
        """Agent instructions for guiding behavior and responses.

        Attributes
        ----------
        response : str or None
            Instructions for how the agent should format and deliver responses.
        orchestration : str or None
            Instructions for how the agent should orchestrate tool usage and
            workflow.
        """
        response: Optional[str] = None
        orchestration: Optional[str] = None

    @dataclass
    class BudgetConfig:
        """Budget constraints for agent execution.

        If more than one constraint is specified, whichever is first hit will
        end the request.

        Attributes
        ----------
        seconds : int or None
            Time budget in seconds for agent execution.
        tokens : int or None
            Token budget for model consumption.
        """
        seconds: Optional[int] = None
        tokens: Optional[int] = None

    @dataclass
    class OrchestrationConfig:
        """Configuration for agent orchestration behavior.

        Attributes
        ----------
        budget : CortexAgentApi.BudgetConfig or None
            Budget constraints for the agent.
        """
        budget: Optional['CortexAgentApi.BudgetConfig'] = None

    @dataclass
    class Models:
        """Model configuration for agent orchestration.

        Specifies which LLM model the agent should use for orchestrating tool
        usage and generating responses.

        Supported orchestration model values as of Jan 2026:

        - ``claude-haiku-4-5``     (Claude Haiku 4.5)
        - ``claude-sonnet-4-5``    (Claude Sonnet 4.5)
        - ``claude-4-sonnet``      (Claude Sonnet 4)
        - ``claude-3-7-sonnet``    (Claude Sonnet 3.7)
        - ``claude-3-5-sonnet``    (Claude Sonnet 3.5)
        - ``openai-gpt-5``         (GPT 5)
        - ``openai-gpt-4.1``       (GPT 4.1)

        Attributes
        ----------
        orchestration : str or None
            The model identifier for agent orchestration.
        """
        orchestration: Optional[str] = None

    @dataclass
    class ToolSpec:
        """Specification for a tool that the agent can use.

        Tools enable agents to perform actions like searching, analyzing data,
        or calling custom functions.

        Attributes
        ----------
        type : str
            Tool type (e.g., ``"function"``, ``"cortex_analyst"``,
            ``"cortex_search"``, ``"web_search"``).
        name : str
            Unique identifier for the tool.
        description : str or None
            Human-readable description of what the tool does.
        input_schema : dict or None
            JSON schema defining the tool's input parameters.
        """
        type: str
        name: str
        description: Optional[str] = None
        input_schema: Optional[Dict[str, Any]] = None

    @dataclass
    class Tool:
        """Container for a tool specification.

        Attributes
        ----------
        tool_spec : CortexAgentApi.ToolSpec
            The tool specification.
        """
        tool_spec: 'CortexAgentApi.ToolSpec'

    @dataclass
    class ExecutionEnvironment:
        """Configuration for server-side tool execution environment.

        Defines the Snowflake warehouse and timeout settings for executing
        tools that require compute resources.

        Attributes
        ----------
        query_timeout : int or None
            Maximum execution time in seconds for queries.
        type : str or None
            Environment type identifier.
        warehouse : str or None
            Name of the Snowflake warehouse to use for execution.
        """
        query_timeout: Optional[int] = None
        type: Optional[str] = None
        warehouse: Optional[str] = None

    @dataclass
    class ToolResourceSpec:
        """Resource configuration for a specific tool.

        Defines the data sources, execution environment, and other resources
        that a tool needs to operate.

        Attributes
        ----------
        execution_environment : CortexAgentApi.ExecutionEnvironment or None
            Warehouse configuration for server-side execution.
        semantic_model_file : str or None
            Path to semantic model file for Cortex Analyst.
        semantic_view : str or None
            Fully qualified name of semantic view.
        identifier : str or None
            Resource identifier.
        name : str or None
            Resource name.
        type : str or None
            Resource type.
        """
        execution_environment: Optional['CortexAgentApi.ExecutionEnvironment'] = None
        semantic_model_file: Optional[str] = None
        semantic_view: Optional[str] = None
        identifier: Optional[str] = None
        name: Optional[str] = None
        type: Optional[str] = None

    @dataclass
    class AgentSpec:
        """Complete specification for a Cortex agent.

        Defines all configuration aspects of an agent including its identity,
        behavior, capabilities, and resource access.

        Attributes
        ----------
        name : str or None
            Agent identifier (must be unique within schema).
        profile : CortexAgentApi.Profile or None
            Visual and display configuration.
        models : CortexAgentApi.Models or None
            Model selection for orchestration.
        instructions : CortexAgentApi.Instructions or None
            Behavioral guidance for the agent.
        tools : list of CortexAgentApi.Tool or None
            Tools available to the agent.
        tool_resources : dict or None
            Resource configurations keyed by tool name.
        orchestration : CortexAgentApi.OrchestrationConfig or None
            Orchestration configuration including budget constraints.
        """
        name: Optional[str] = None
        profile: Optional['CortexAgentApi.Profile'] = None
        models: Optional['CortexAgentApi.Models'] = None
        instructions: Optional['CortexAgentApi.Instructions'] = None
        tools: Optional[List['CortexAgentApi.Tool']] = None
        tool_resources: Optional[Dict[str, 'CortexAgentApi.ToolResourceSpec']] = None
        orchestration: Optional['CortexAgentApi.OrchestrationConfig'] = None

        def set_semantic_view(self, tool_name: str, view_fqn: str) -> 'CortexAgentApi.AgentSpec':
            """Set the semantic view for a specific tool.

            Parameters
            ----------
            tool_name : str
                Name of the tool to configure.
            view_fqn : str
                Fully qualified name of the semantic view.

            Returns
            -------
            CortexAgentApi.AgentSpec
                Self for method chaining.
            """
            assert self.tool_resources is not None, "tool_resources must be initialized"
            assert tool_name in self.tool_resources, f"tool_name '{tool_name}' not in tool_resources"
            self.tool_resources[tool_name].semantic_view = view_fqn
            return self

        def copy(self) -> 'CortexAgentApi.AgentSpec':
            """Create a deep copy of this agent specification.

            Returns
            -------
            CortexAgentApi.AgentSpec
                Independent copy of this AgentSpec.
            """
            return copy.deepcopy(self)

        def merge(self, other: 'CortexAgentApi.AgentSpec') -> 'CortexAgentApi.AgentSpec':
            """Merge another AgentSpec into this one, creating a new combined specification.

            Non-None fields from the other spec override or extend this spec's
            fields. Tools and tool_resources are combined rather than replaced.

            Parameters
            ----------
            other : CortexAgentApi.AgentSpec
                AgentSpec to merge into this one.

            Returns
            -------
            CortexAgentApi.AgentSpec
                New AgentSpec with merged configuration.
            """
            result = self.copy()
            if other.name is not None:
                result.name = other.name
            if other.profile is not None:
                result.profile = other.profile
            if other.models is not None:
                result.models = other.models
            if other.instructions is not None:
                result.instructions = other.instructions
            if other.orchestration is not None:
                result.orchestration = other.orchestration
            if other.tools is not None:
                if result.tools is None:
                    result.tools = copy.deepcopy(other.tools)
                else:
                    result.tools = result.tools + copy.deepcopy(other.tools)
            if other.tool_resources is not None:
                if result.tool_resources is None:
                    result.tool_resources = copy.deepcopy(other.tool_resources)
                else:
                    result.tool_resources = {**result.tool_resources, **copy.deepcopy(other.tool_resources)}
            return result

        def to_dict(self) -> dict:
            """Convert this AgentSpec to a dictionary.

            Returns
            -------
            dict
                Dictionary representation suitable for JSON serialization.
            """
            return dataclasses.asdict(self)

        @staticmethod
        def from_dict(spec_data: dict) -> Optional['CortexAgentApi.AgentSpec']:
            """Parse an AgentSpec from a dictionary.

            Parameters
            ----------
            spec_data : dict
                Dictionary containing agent specification data.

            Returns
            -------
            CortexAgentApi.AgentSpec or None
                Parsed AgentSpec instance, or ``None`` if parsing fails.
            """
            # Parse name
            name = spec_data.get("name")

            # Parse profile
            profile = None
            if spec_data.get("profile"):
                profile = CortexAgentApi.Profile(
                    avatar=spec_data["profile"].get("avatar"),
                    color=spec_data["profile"].get("color"),
                    display_name=spec_data["profile"].get("display_name")
                )

            # Parse models
            models = None
            if spec_data.get("models"):
                models = CortexAgentApi.Models(
                    orchestration=spec_data["models"].get("orchestration")
                )

            # Parse instructions
            instructions = None
            if spec_data.get("instructions"):
                instructions = CortexAgentApi.Instructions(
                    response=spec_data["instructions"].get("response"),
                    orchestration=spec_data["instructions"].get("orchestration")
                )

            # Parse tools
            tools = None
            if spec_data.get("tools"):
                tools = []
                for tool_data in spec_data["tools"]:
                    tool_spec_data = tool_data.get("tool_spec", {})
                    tool_spec = CortexAgentApi.ToolSpec(
                        type=tool_spec_data["type"],
                        name=tool_spec_data["name"],
                        description=tool_spec_data.get("description"),
                        input_schema=tool_spec_data.get("input_schema")
                    )
                    tool = CortexAgentApi.Tool(tool_spec=tool_spec)
                    tools.append(tool)

            # Parse tool resources
            tool_resources = None
            if spec_data.get("tool_resources"):
                tool_resources = {}
                for resource_name, resource_data in spec_data["tool_resources"].items():
                    execution_env = None
                    if resource_data.get("execution_environment"):
                        env_data = resource_data["execution_environment"]
                        execution_env = CortexAgentApi.ExecutionEnvironment(
                            query_timeout=env_data.get("query_timeout"),
                            type=env_data.get("type"),
                            warehouse=env_data.get("warehouse")
                        )

                    resource_spec = CortexAgentApi.ToolResourceSpec(
                        execution_environment=execution_env,
                        semantic_model_file=resource_data.get("semantic_model_file"),
                        semantic_view=resource_data.get("semantic_view"),
                        identifier=resource_data.get("identifier"),
                        name=resource_data.get("name"),
                        type=resource_data.get("type")
                    )
                    tool_resources[resource_name] = resource_spec

            # Parse orchestration
            orchestration = None
            if spec_data.get("orchestration"):
                budget = None
                if spec_data["orchestration"].get("budget"):
                    budget_data = spec_data["orchestration"]["budget"]
                    budget = CortexAgentApi.BudgetConfig(
                        seconds=budget_data.get("seconds"),
                        tokens=budget_data.get("tokens")
                    )
                orchestration = CortexAgentApi.OrchestrationConfig(budget=budget)

            return CortexAgentApi.AgentSpec(
                name=name,
                profile=profile,
                models=models,
                instructions=instructions,
                tools=tools,
                tool_resources=tool_resources,
                orchestration=orchestration
            )

    @include_in_docs
    def create(self, database: str, schema: str, request: AgentSpec, create_mode: str = "orReplace") -> None:
        """Create an agent.

        Parameters
        ----------
        database : str
            Database name.
        schema : str
            Schema name.
        request : AgentSpec
            Agent specification.
        create_mode : str
            One of ``"orReplace"``, ``"errorIfExists"``, or ``"ifNotExists"``.
            Default: ``"orReplace"``.

        Returns
        -------
        None
        """
        valid_modes = ["orReplace", "errorIfExists", "ifNotExists"]
        if create_mode not in valid_modes:
            exc(
                "InvalidCreateMode",
                f"create_mode must be one of {valid_modes}, got: {create_mode}",
            )

        response = self._c.post(
            url=f"/api/v2/databases/{database}/schemas/{schema}/agents?createMode={create_mode}",
            json=dataclasses.asdict(request),
        )
        # {"status": "..."}
        response.raise_for_status()

    @include_in_docs
    def update(self, database: str, schema: str, name: str, request: AgentSpec) -> None:
        """Update an existing agent's configuration.

        Modifies the properties of an existing agent. Fields not included in
        the request are left unchanged.

        Parameters
        ----------
        database : str
            Database name.
        schema : str
            Schema name.
        name : str
            Agent name to update.
        request : AgentSpec
            Updated agent specification.

        Returns
        -------
        None

        Raises
        ------
        httpx.HTTPStatusError
            If the request fails.
        """
        response = self._c.put(
            url=f"/api/v2/databases/{database}/schemas/{schema}/agents/{name}",
            json=dataclasses.asdict(request),
        )
        response.raise_for_status()

    @dataclass
    class DescribeAgentResponse:
        """Response containing comprehensive agent details.

        Attributes
        ----------
        agent_spec_dict : dict
            Raw dictionary of agent specification.
        created_on : str
            Timestamp of agent creation.
        database_name : str
            Database containing the agent.
        name : str
            Agent name.
        owner : str
            User or role that owns the agent.
        schema_name : str
            Schema containing the agent.
        agent_spec : CortexAgentApi.AgentSpec or None
            Parsed agent specification.
        """
        agent_spec_dict: Dict[str, Any]
        created_on: str
        database_name: str
        name: str
        owner: str
        schema_name: str
        agent_spec: Optional['CortexAgentApi.AgentSpec'] = None

    @include_in_docs
    def describe(self, database: str, schema: str, name: str) -> DescribeAgentResponse:
        """Retrieve comprehensive details about a specific agent.

        Returns complete agent metadata including configuration, ownership,
        and creation information.

        Parameters
        ----------
        database : str
            Database name.
        schema : str
            Schema name.
        name : str
            Agent name.

        Returns
        -------
        DescribeAgentResponse
            Full agent details.

        Raises
        ------
        httpx.HTTPStatusError
            If the request fails or agent doesn't exist.
        """
        response = self._c.get(
            url=f"/api/v2/databases/{database}/schemas/{schema}/agents/{name}",
        )
        response.raise_for_status()
        data = response.json()

        agent_spec_dict = json.loads(data['agent_spec'])
        parsed_spec = CortexAgentApi.AgentSpec.from_dict(agent_spec_dict)
        if parsed_spec is not None and not parsed_spec.name:
            parsed_spec.name = name

        return CortexAgentApi.DescribeAgentResponse(
            agent_spec_dict=agent_spec_dict,
            created_on=data['created_on'],
            database_name=data['database_name'],
            name=data['name'],
            owner=data['owner'],
            schema_name=data['schema_name'],
            agent_spec=parsed_spec
        )

    @dataclass
    class ListAgentsResponseInner:
        """Summary information for a single agent in a list response.

        Attributes
        ----------
        name : str
            Agent name.
        comment : str or None
            Optional comment/description.
        created_on : str
            Timestamp of agent creation.
        database_name : str
            Database containing the agent.
        owner : str
            User or role that owns the agent.
        profile : CortexAgentApi.Profile or None
            Agent's visual configuration.
        schema_name : str
            Schema containing the agent.
        """
        name: str
        comment: Optional[str]
        created_on: str
        database_name: str
        owner: str
        profile: Optional['CortexAgentApi.Profile']
        schema_name: str

    @dataclass
    class ListAgentsResponse:
        """Response containing a list of agent summaries.

        Attributes
        ----------
        agents : list of CortexAgentApi.ListAgentsResponseInner
            Agent summary objects.
        """
        agents: List['CortexAgentApi.ListAgentsResponseInner']

    @include_in_docs
    def list(self, database: str, schema: str) -> ListAgentsResponse:
        """List all agents in a schema.

        Returns summary information for all agents in the specified database
        and schema.

        Parameters
        ----------
        database : str
            Database name.
        schema : str
            Schema name.

        Returns
        -------
        ListAgentsResponse
            Agent summaries.

        Raises
        ------
        httpx.HTTPStatusError
            If the request fails.
        """
        response = self._c.get(
            url=f"/api/v2/databases/{database}/schemas/{schema}/agents",
        )
        response.raise_for_status()
        data = response.json()
        agents = []
        for agent_data in data:
            profile = None
            if agent_data.get('profile'):
                profile = CortexAgentApi.Profile(**agent_data['profile'])

            agent = CortexAgentApi.ListAgentsResponseInner(
                name=agent_data['name'],
                comment=agent_data.get('comment'),
                created_on=agent_data['created_on'],
                database_name=agent_data['database_name'],
                owner=agent_data['owner'],
                profile=profile,
                schema_name=agent_data['schema_name']
            )
            agents.append(agent)
        return CortexAgentApi.ListAgentsResponse(agents=agents)

    @include_in_docs
    def delete(self, database: str, schema: str, name: str, silent=True):
        """Delete an agent.

        Permanently removes an agent from the database. The behavior when the
        agent doesn't exist depends on the ``silent`` parameter.

        Parameters
        ----------
        database : str
            Database name.
        schema : str
            Schema name.
        name : str
            Agent name to delete.
        silent : bool
            If ``True``, suppress errors when agent doesn't exist.
            Default: ``True``.

        Returns
        -------
        httpx.Response
            The HTTP response object.

        Raises
        ------
        httpx.HTTPStatusError
            If ``silent=False`` and the request fails.
        """
        response = self._c.delete(
            url=f"/api/v2/databases/{database}/schemas/{schema}/agents/{name}",
        )
        if not silent:
            response.raise_for_status()
        return response

    @include_in_docs
    def run(self, database: str, schema: str, name: str, thread_id: int, parent_message_id: int, prompt: str) -> Tuple[dict, int]:
        """Execute an agent with a user prompt in a conversation thread.

        Sends a message to the agent and streams back the response. Processes
        server-sent events to extract the agent's response and updated message
        ID.

        Parameters
        ----------
        database : str
            Database name.
        schema : str
            Schema name.
        name : str
            Agent name to run.
        thread_id : int
            Conversation thread identifier.
        parent_message_id : int
            ID of the previous message in the thread.
        prompt : str
            User's message/prompt for the agent.

        Returns
        -------
        tuple of (dict, int)
            ``(response_dict, new_parent_message_id)`` where
            ``response_dict`` contains the agent's response with ``"content"``
            and ``"role"`` fields, and ``new_parent_message_id`` is the message
            ID for the next interaction.

        Raises
        ------
        APIResponseError
            If the agent returns an error.
        httpx.HTTPStatusError
            If the HTTP request fails.
        AssertionError
            If the response doesn't contain expected content.
        """
        error_dict = None
        with self._c.stream(
                "POST",
                url=f"/api/v2/databases/{database}/schemas/{schema}/agents/{name}:run",
                json={
                    "thread_id": thread_id,
                    "parent_message_id": parent_message_id,
                    "messages": [{
                        "role": "user",
                        "content": [{
                            "type": "text",
                            "text": prompt
                        }]
                    }],
                },
        ) as response:
            if not response.is_success:
                response.raise_for_status()
            data = None
            event = None
            for line in response.iter_lines():
                # print(line)
                if line.startswith("event: "):
                    event = line
                elif line.startswith("data: "):
                    if event == "event: error":
                        error_dict = json.loads(line[6:])
                    if event == "event: response":
                        data = json.loads(line[6:])
                        break
                    if event == "event: metadata":
                        # data: {"metadata": {"message_id": 37856360561, "role": "assistant"}}
                        data = json.loads(line[6:])
                        if data["metadata"]["role"] == "assistant":
                            parent_message_id = data["metadata"]["message_id"]
        if error_dict:
            raise APIResponseError(error_dict)
        assert data is not None, "Response data should not be None"
        assert "content" in data and data["content"] is not None and len(data["content"]) > 0, "Response must have content"
        data["role"] = "assistant"
        return data, parent_message_id
