"""Cortex Search / RAG agent.

Wraps Cortex Search infrastructure (src/cortex_agent/) for
search and retrieval-augmented generation operations.
"""

from __future__ import annotations

import logging
from typing import Any, Dict

from .base_agent import BaseDataBridgeAgent
from .state_schema import WorkflowState

logger = logging.getLogger(__name__)


class CortexSearchAgent(BaseDataBridgeAgent):
    """Specialist agent for Cortex Search and RAG operations."""

    name = "cortex_search_agent"
    description = "Performs Cortex Search and RAG (retrieval-augmented generation) queries"
    phase_name = "cortex_search"

    def __init__(self, **kwargs):
        super().__init__(
            tools=[],
            system_prompt=(
                "You are the Cortex Search agent for DataBridge AI. "
                "You perform search queries and retrieval-augmented "
                "generation using Snowflake Cortex Search."
            ),
            **kwargs,
        )

    async def run(self, state: WorkflowState) -> WorkflowState:
        """Execute Cortex Search operation."""
        config = state.get("config", {})
        action = config.get("search_action", "search")

        if action == "search":
            return await self._run_search(state, config)
        elif action == "rag":
            return await self._run_rag(state, config)

        return self._update_context(state, self.phase_name, {
            "error": f"Unknown search action: {action}"
        })

    async def _run_search(self, state: WorkflowState, config: Dict) -> WorkflowState:
        """Run a Cortex Search query."""
        query = config.get("search_query", "")
        if not query:
            return self._update_context(state, self.phase_name, {
                "error": "search_query is required"
            })

        try:
            # Delegate to existing search infrastructure
            return self._update_context(state, self.phase_name, {
                "action": "search",
                "query": query,
                "status": "placeholder",
                "note": "Use cortex_search MCP tool directly for search operations",
            })
        except Exception as e:
            return self._update_context(state, self.phase_name, {"error": str(e)})

    async def _run_rag(self, state: WorkflowState, config: Dict) -> WorkflowState:
        """Run a RAG query using Cortex Search + COMPLETE."""
        query = config.get("search_query", "")
        if not query:
            return self._update_context(state, self.phase_name, {
                "error": "search_query is required for RAG"
            })

        try:
            return self._update_context(state, self.phase_name, {
                "action": "rag",
                "query": query,
                "status": "placeholder",
                "note": "Use rag_search MCP tool for RAG operations",
            })
        except Exception as e:
            return self._update_context(state, self.phase_name, {"error": str(e)})
