from dataiku.llm.agent_tools import BaseAgentTool
from snowflake.core import Root
import json
import logging
from connect_to_snowflake import SnowflakeConnectionHelper

logger = logging.getLogger("snowflakecortexsearchtool")

class SnowflakeCortexSearchTool(BaseAgentTool):
    def set_config(self, config, plugin_config):
        self.config = config

    def get_descriptor(self, tool):
        return {
            "description": "Searches a Snowflake Cortex Search service. Returns an array of results. For each result, returns the content of all columns selected.",
            "inputSchema": {
                "$id": "https://dataiku.com/agents/tools/cortexsearch/input",
                "title": "Input for the cortex search tool",
                "type": "object",
                "properties": {
                    "query": {
                        "type": "string",
                        "description": "The Cortex Search query string"
                    }
                },
                "required": ["query"]
            }
        }

    def load_sample_query(self, tool):
        return {
            "query": "Your question for the Cortex Search tool"
        }

    def invoke(self, input, trace):
        # Get tool input arguments, inlcuding the user query 'query'
        args = input["input"]
        query = args["query"]

        # Log inputs and config to trace
        trace.span["name"] = "SNOWFLAKE_CORTEX_SEARCH_TOOL_CALL"
        for key, value in args.items():
            trace.inputs[key] = value
        trace.attributes["config"] = self.config

        # Get the tool parameters, set at the Dataiku project, tool level
        snowflake_connection_name = self.config.get("snowflake_connection", None)
        cortex_search_database = self.config.get("cortex_search_database", None)
        cortex_search_schema = self.config.get("cortex_search_schema", None)
        cortex_search_service = self.config.get("cortex_search_service", None)
        max_documents = self.config.get("max_documents", None)
        search_column = self.config.get("search_column", None)
        metadata_columns = self.config.get("metadata_columns", None)

        # Combine all search columns to send to Cortex Search
        all_search_columns = metadata_columns.copy()
        all_search_columns.append(search_column)

        # Connect to Snowflake
        helper = SnowflakeConnectionHelper(snowflake_connection_name) #, cortex_search_database, cortex_search_schema)
        root = Root(helper.snowpark_session)

        # Get the Cortex Search service
        cortex_search_service = root.databases[cortex_search_database].schemas[cortex_search_schema].cortex_search_services[cortex_search_service]

        # Search the Cortex Search service for the user query
        response = cortex_search_service.search(query, all_search_columns, limit=max_documents)
        logger.info(f"Cortex Search response: {response.json()}")

        # Get the results
        results = json.loads(response.json())['results']

        # Load the results into source items that will citable by Dataiku Agent Connect
        source_items = []
        for result in results:
            title_list = []
            for metadata_column in metadata_columns:
                title_list.append(result[metadata_column])
            title = ", ".join(title_list)
            source_item = {
                "type": "SIMPLE_DOCUMENT",
                "title": title,
                "textSnippet": result[search_column]
            }
            source_items.append(source_item)

        # Log outputs to trace
        trace.outputs["output"] = results

        # Return the Cortex Search service response
        return {
            "output": results,
            "sources": [{
                "toolCallDescription": f"Snowflake Cortex Search tool",
                "items": source_items
            }]
        }
