from dataiku.llm.agent_tools import BaseAgentTool
from snowflake.snowpark.exceptions import SnowparkSQLException
import json
import requests
import logging
import datetime
from connect_to_snowflake import SnowflakeConnectionHelper
from snowflake_tools_utils import make_dates_json_safe

logger = logging.getLogger("snowflakecortexanalysttool")

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

    def get_descriptor(self, tool):
        return {
            "description": "Queries Snowflake tables using a Cortex Analyst service in order to answer a business question. Returns the Cortex Analyst service response, which may include text, SQL query suggestions, the result of a query, or an error message.",
            "inputSchema": {
                "$id": "https://dataiku.com/agents/tools/cortexanalyst/input",
                "title": "Input for the Cortex Analyst tool",
                "type": "object",
                "properties": {
                    "query": {
                        "type": "string",
                        "description": "The Cortex Analyst input message string"
                    }
                },
                "required": ["query"]
            }
        }

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

    def get_analyst_response(
        self,
        messages,
        semantic_view_or_semantic_model_file,
        semantic_view,
        semantic_model_file_path,
        snowflake_host,
        headers
    ):
        """
        Send input message to the Cortex Analyst API and returns a response.

        Args:
            messages (List[Dict]): The input messages (ideally conversation history per Cortex Analyst API).
            semantic_view_or_semantic_model_file (String): Either "semantic_view" or "semantic_model_file"
            semantic_view (String): The Cortex Analyst semantic view (e.g. DATABASE.SCHEMA.SEMANTIC_VIEW) in Snowflake. Should be created beforehand.
            semantic_model_file_path (String): The path to the Cortex Analyst semantic model .yaml file (e.g. @DATABASE.SCHEMA.STAGE/SEMANTIC_MODEL_FILE.yaml) in Snowflake. Should be created beforehand.
            snowflake_host (String): The Snowflake account ID e.g. XY12345
            snowflake_token (String): A Snowflake token obtained through a Snowflake connection request

        Returns:
            Optional[Dict]: The response from the Cortex Analyst API.
        """
        # Prepare the request body with the user's prompt
        if semantic_view_or_semantic_model_file == "semantic_view" and semantic_view:
            request_body = {
                "messages": messages,
                "semantic_view": semantic_view,
            }

        elif semantic_view_or_semantic_model_file == "semantic_model_file" and semantic_model_file_path:
            request_body = {
                "messages": messages,
                "semantic_model_file": semantic_model_file_path,
            }

        request_url = f"https://{snowflake_host}/api/v2/cortex/analyst/message"
        logger.info(
            f"Sending request to Snowflake Cortex Analyst on endpoint {request_url}")

        # Send a POST request to the Cortex Analyst API endpoint
        response = requests.post(
            url=request_url,
            json=request_body,
            headers=headers
        )

        # Content is a string with serialized JSON object
        parsed_content = json.loads(response.text)

        # Check if the response is successful
        if response.status_code < 400:
            # Return the content of the response as a JSON object
            logger.info(f"Got successful reponse: {response.status_code}")
            return parsed_content, None
        else:
            raise Exception(f"""Cortex Analyst request failed:
HTTP code: {response.status_code}
Error text: {response.text}""")

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

        # Log config to trace
        trace.span["name"] = "SNOWFLAKE_CORTEX_ANALYST_TOOL_CALL"
        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)
        semantic_view_or_semantic_model_file = self.config.get(
            "semantic_view_or_semantic_model_file", None)
        semantic_model_file_path = self.config.get(
            "semantic_model_file_path", None)
        semantic_view = self.config.get("semantic_view", None)
        records_returned_limit = self.config.get("records_returned_limit", -1)

        # Connect to Snowflake and get connection info
        snowflake_connector = SnowflakeConnectionHelper(snowflake_connection_name)

        # Create a new user message
        new_user_message = {
            "role": "user",
            "content": [{"type": "text", "text": query}],
        }

        # Send the user message to the Cortex Analyst service, and get a response
        response, error_msg = self.get_analyst_response([new_user_message],
                                                        semantic_view_or_semantic_model_file,
                                                        semantic_view,
                                                        semantic_model_file_path,
                                                        snowflake_connector.snowflake_host,
                                                        snowflake_connector.headers)

        # If there's an error message, return it
        if error_msg is not None:
            return {
                "output": error_msg
            }

        # Get the Cortex Analyst response message
        result_content = response['message']['content']
        logger.info(f"Cortex Analyst response: {result_content}")

        # Create a new list to send back to the agent
        result_output = []
        source_items = []
        artifacts = []
        # Iterate through text, suggestions, and SQL code returned by the Cortex Analyst Service
        # Add the text and suggestions to the output. If SQL code, try to run it. If the code succeeds, add the query result to the output.
        for content in result_content:
            if content['type'] == "text":
                result_output.append({"text": content['text']})
            elif content["type"] == "suggestions":
                result_output.append({"suggestions": content["suggestions"]})
            elif content['type'] == 'sql':
                source_items.append({
                    "type": "GENERATED_SQL_QUERY",
                    "title": "Cortex Analyst SQL Query",
                    "performedQuery": content['statement']
                })
                logger.info(
                    f"Added SQL source: {content['statement'][:100]}...")
                try:
                    result_output.append({"generated_sql_query": content['statement']})
                    logger.info(f"Running SQL code: {content['statement']}")

                    # Run the query, return a Snowpark Dataframe.
                    # Note: The Snowpark Dataframe "limit" method seems to not like the ";" auto-generated at the end of Cortex Analyst queries, so we strip it
                    snowpark_df = snowflake_connector.snowpark_session.sql(
                        content['statement'].rstrip(';'))

                    # If records returned limit parameter set (to a positive value), apply it
                    if records_returned_limit > 0:
                        logger.info(
                            f"Applying limit of {records_returned_limit} to returned records")
                        df = snowpark_df.limit(
                            records_returned_limit).to_pandas()
                    else:
                        df = snowpark_df.to_pandas()

                    columns = df.columns.tolist()

                    # This accounts for datetime objects in columns
                    df = df.copy()
                    for col in df.select_dtypes(include=["datetime64[ns]", "datetimetz"]):
                        df[col] = df[col].astype(str)
                    data = df.to_numpy().tolist()

                    if df.empty:
                        logger.info("Query returned no data")
                        result_output.append({"query_results": "No data"})
                    else:
                        df_string = df.to_string()
                        logger.info(f"Query returned data: {df_string}")
                        result_output.append({"query_results": make_dates_json_safe({ "columns": columns, "data": data })})

                        source_items_payload = make_dates_json_safe({
                            "type": "RECORDS",
                            "title": "Cortex Analyst SQL Results",
                            "records": {
                                "columns": columns,
                                "data": data
                            }
                        })
                        source_items.append(source_items_payload)

                        artifacts_payload = make_dates_json_safe({
                            "name": "Records",
                            "description": "Cortex Analyst SQL Results",
                            "parts": [{
                                "type": "RECORDS",
                                "records": {
                                    "columns": columns,
                                    "data": data
                                }
                            }]
                        })
                        artifacts.append(artifacts_payload)

                except SnowparkSQLException as e:
                    logger.info(
                        f"Could not execute generated SQL query. Error: {e}")
                    result_output.append(
                        f"Could not execute generated SQL query. Error: {e}")
            else:
                pass

        # Return the Cortex Analyst response text, suggestions, and SQL query results
        return {
            "output": result_output,
            "sources": [{
                "toolCallDescription": "Snowflake Cortex Analyst Tool",
                "items": source_items
            }],
            "artifacts": artifacts
        }
