from dataiku.llm.agent_tools import BaseAgentTool
import dataiku
import logging


logger = logging.getLogger(__name__)


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

    def get_descriptor(self, tool):
        return {
            "description": "Look up orders by status in the example orders dataset.",
            "inputSchema": {
                "$id": "https://dataiku.com/agents/tools/order-lookup/input",
                "title": "Order lookup input",
                "type": "object",
                "properties": {
                    "filter": {
                        "type": "object",
                        "description": "Filter to apply to the orders dataset.",
                        "properties": {
                            "column": {
                                "type": "string",
                                "enum": ["Status"],
                                "description": "Column to filter. This example tool only accepts Status.",
                            },
                            "operator": {
                                "type": "string",
                                "enum": ["EQUALS"],
                                "description": "Comparison operator. This example tool only supports equality.",
                            },
                            "value": {
                                "type": "string",
                                "description": "Status value to match, for example PROCESSING.",
                            },
                        },
                        "required": ["column", "operator", "value"],
                        "additionalProperties": False,
                    }
                },
                "required": ["filter"],
                "additionalProperties": False,
            },
        }

    def invoke(self, input, trace):
        args = input.get("input", {})
        filter_spec = args.get("filter", {})

        column = filter_spec.get("column")
        operator = filter_spec.get("operator")
        value = filter_spec.get("value")

        if column != "Status" or operator != "EQUALS" or value is None:
            return {
                "output": {
                    "rows": [],
                    "message": "This example tool only supports Status EQUALS <value> filters.",
                }
            }

        requested_status = str(value).upper()
        rows = []

        dataset = dataiku.Dataset("orders")
        for row in dataset.iter_rows():
            if str(row.get("Status", "")).upper() == requested_status:
                rows.append(
                    {
                        "OrderID": row.get("OrderID"),
                        "CustomerName": row.get("CustomerName"),
                        "Status": row.get("Status"),
                        "Total": row.get("Total"),
                        "Region": row.get("Region"),
                    }
                )

        return {
            "output": {
                "rows": rows,
                "message": f"Found {len(rows)} order(s) with Status={requested_status}.",
            }
        }