import json
import logging
from typing import Any

import dataiku
from dataiku.llm.agent_tools import BaseAgentTool

from backends.explorer.queries.explorer_data import run_llm_query
from backends.shared.params import RunLlmCypherParams
from graph.connections.neo4j.preset import Neo4jPresetConfig
from recipes.explore_graph.utils import (
    get_connection_info_registry_from_inputs,
    get_db_handler_from_snapshot,
    get_single_connection_info,
)


class GraphSearch(BaseAgentTool):
    def set_config(self, config, plugin_config):
        self.logger = logging.getLogger(__name__)
        self.config = config
        logging.basicConfig(
            level=logging.DEBUG if config.get("verbose_mode", False) else logging.INFO,
            format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
            force=True,
        )
        db_folder_id = config.get("db_folder_id")
        self.db_folder = dataiku.Folder(db_folder_id) if db_folder_id else None

        neo4j_preset = config.get("neo4j_preset")
        self.neo4j_configuration = Neo4jPresetConfig.from_dict(neo4j_preset) if neo4j_preset else None
        self.database_name = config.get("database_name")
        raw_database_type = config.get("database_type")
        self.database_type = "neo4j" if raw_database_type == "unmanaged_neo4j" else raw_database_type

        self.conn_info_registry = get_connection_info_registry_from_inputs(
            db_folder=self.db_folder,
            neo4j_configuration=self.neo4j_configuration,
            database_name=self.database_name,
            database_type=self.database_type,
        )
        self.conn_info = get_single_connection_info(
            db_folder=self.db_folder,
            neo4j_configuration=self.neo4j_configuration,
            database_name=self.database_name,
            database_type=self.database_type,
        )

        # Get llm connection
        client = dataiku.api_client()
        project = client.get_default_project()
        self.llm = project.get_llm(config.get("llm_id"))

        

    def get_descriptor(self, tool) -> dict[str, Any]:  # type: ignore
        return {
            "description": f"Searches a graph database with the following description. \n {self.config.get('content_description', 'no description provided')} \n Returns a JSON object",
            "inputSchema": {
                "$id": "https://dataiku.com/agents/tools/search/input",
                "title": "Input for graph search",
                "type": "object",
                "properties": {"query": {"type": "string", "description": "The query string"}},
                "required": ["query"],
            },
            "required": ["q"],
        }

    def invoke(self, input, trace) -> dict[str, Any]:  # type: ignore
        args = input["input"]
        query = args["query"]

        self.logger.info(f"Graph search query: {query}")

        db_handler = get_db_handler_from_snapshot(
            self.db_folder,
            self.conn_info.uid,
            self.neo4j_configuration,
            database_name=self.database_name,
            database_type=self.database_type,
        )

        result = run_llm_query(
            params=RunLlmCypherParams(graph_id=self.conn_info.uid, query=query),
            llm=self.llm,
            db_handler=db_handler,
            conn_info_registry=self.conn_info_registry,
            timeout_seconds=self.config.get("db_query_timeout_seconds", 60),
        )

        return {
            "output": json.dumps(result),
            "sources": [{"toolCallDescription": f"Graph search: {result['cypher_query']}", "items": []}],
        }
