import dataiku
from connect_to_snowflake import SnowflakeConnectionHelper
import requests
import logging

logger = logging.getLogger("snowflakecortexplugin")

def do(payload, config, plugin_config, inputs):
    if not config.get("snowflake_connection"):
        logger.warning("No snowflake_connection selected yet")
        return {"choices": []}
    if not config.get("agent_database"):
        logger.warning("No agent_database selected yet")
        return {"choices": []}
    if not config.get("agent_schema"):
        logger.warning("No agent_schema selected yet")
        return {"choices": []}

    snowflake_connector = SnowflakeConnectionHelper(config.get("snowflake_connection"))
    agent_database = config.get("agent_database")
    agent_schema = config.get("agent_schema")
    list_cortex_agents_endpoint = f"/api/v2/databases/{agent_database}/schemas/{agent_schema}/agents"
    list_cortex_agents_url = "https://" + \
        snowflake_connector.snowflake_host + list_cortex_agents_endpoint

    try:
        resp = requests.get(list_cortex_agents_url, headers=snowflake_connector.headers, timeout=100)

        if resp.status_code == 200:
            resp_json_list = resp.json()
            logger.info(f"Cortex Agents Available: {str(resp_json_list)}")
            cortex_agents_choices = [
                {"value": resp_json["name"], "label": resp_json["name"]} for resp_json in resp_json_list]
            return {"choices": cortex_agents_choices}
        else:
            response_error = f"HTTP Error: {resp['status']} - {resp.get('reason', 'Unknown reason')}. Response details: {resp}"
            return response_error

    except Exception as e:
        response_error = f"Error making request: {str(e)}"
        return response_error
