# Copyright [2025] [IBM]
# Licensed under the Apache License, Version 2.0 (http://www.apache.org/licenses/LICENSE-2.0)
# See the LICENSE file in the project root for license information.

# This file has been modified with the assistance of IBM Bob AI tool

from typing import Union, Literal

from app.core.registry import service_registry
from app.services.data_product.models.search_data_products import (
    SearchDataProductsRequest,
    SearchDataProductsResponse,
)
from app.services.data_product.utils.common_utils import get_dph_catalog_id_for_user
from app.shared.utils.tool_helper_service import tool_helper_service
from app.shared.logging import LOGGER, auto_context


@service_registry.tool(
    name="data_product_search_data_products",
    description="""
    This tool searches all data products (both released and draft) to return data products that match the given search_query.
    Example: 'Find all data products that match the name Environment.'
    In this case, product_search_query is 'Environment', and this tool returns all data products (released and draft) that have Environment in their name.
    Example: 'Find all data products in the Audit domain.'
    In this case, product_search_query is '*', search_filter_type is 'Domain' and search_filter_value is 'Audit'.
    """,
    tags={"search", "data_product"},
    meta={"version": "1.0", "service": "data_product"},
)
@auto_context
async def search_data_products(
    request: SearchDataProductsRequest,
) -> SearchDataProductsResponse:
    LOGGER.info(
        f"In the data_product_search_data_products tool, searching data products with query {request.product_search_query}, filter type {request.search_filter_type} and filter value {request.search_filter_value}."
    )
    DPH_CATALOG_ID = await get_dph_catalog_id_for_user()

    if request.search_filter_type.lower() == "domain":
        should_query = [
            {
                "match": {
                    "entity.data_product_version.domain_name": {
                        "query": request.search_filter_value,
                        "operator": "and",
                        "fuzziness": "AUTO",
                        "prefix_length": 1,
                        "max_expansions": 50,
                    }
                }
            }
        ]
    else:
        search_fields = ["metadata.name", "metadata.description", "metadata.tags"]
        parts_out_search_fields = [
            "entity.data_product_version.parts_out.name",
            "entity.data_product_version.parts_out.description",
            "entity.data_product_version.parts_out.column_names",
            "entity.data_product_version.parts_out.terms",
            "entity.data_product_version.parts_out.column_terms",
        ]
        should_query = [
            {
                "gs_user_query": {
                    "search_string": request.product_search_query,
                    "search_fields": search_fields,
                    "nlq_analyzer_enabled": True,
                    "semantic_search_enabled": True,
                }
            },
            {
                "nested": {
                    "path": "custom_attributes",
                    "query": {
                        "gs_user_query": {
                            "search_string": request.product_search_query,
                            "nested": True,
                        }
                    },
                }
            },
            {
                "nested": {
                    "path": "entity.data_product_version.parts_out",
                    "query": {
                        "gs_user_query": {
                            "search_string": request.product_search_query,
                            "nested": True,
                            "search_fields": parts_out_search_fields,
                            "nlq_analyzer_enabled": True,
                            "semantic_search_enabled": True,
                        }
                    },
                }
            },
        ]

    search_payload = {
        "_source": [
            "artifact_id",
            "last_updated_at",
            "metadata.name",
            "metadata.description",
            "metadata.tags",
            "metadata.created_on",
            "entity.assets.catalog_id",
            "entity.data_product_version",
            "custom_attributes",
        ],
        "query": {
            "bool": {
                "should": should_query,
                "minimum_should_match": 1,
                "filter": [
                    {"terms": {"metadata.artifact_type": ["ibm_data_product_version"]}},
                    {"terms": {"entity.data_product_version.state": ["available","draft"]}},
                    {"terms": {"entity.assets.catalog_id": [DPH_CATALOG_ID]}},
                ],
            }
        },
        "size": 20,
        "aggregations": {
            "product_id": {
                "terms": {"field": "entity.data_product_version.product_id"}
            },
            "state": {"terms": {"field": "entity.data_product_version.state"}},
        },
    }

    response = await tool_helper_service.execute_post_request(
        url=f"{tool_helper_service.base_url}/v3/search?role=viewer&auth_scope=ibm_data_product_catalog",
        json=search_payload,
        tool_name="data_product_search_data_products",
    )

    number_of_responses = response.get("size", len(response.get("rows", [])))
    if number_of_responses == 0:
        LOGGER.info(
            "In the data_product_search_data_products tool, no data products found."
        )
        return SearchDataProductsResponse(count=0, data_products=[])
    
    LOGGER.info(f"Found {number_of_responses} data products.")
    products = []
    for row in response["rows"]:
        metadata = row.get("metadata", {})
        entity = row.get("entity", {})
        products.append(
            {
                "data_product_id": entity.get("data_product_version", {}).get("product_id", ""),
                "name": metadata.get("name", ""),
                "description": metadata.get("description", ""),
                "created_on": metadata.get("created_on", ""),
                "domain": entity.get("data_product_version", {}).get("domain_name", ""),
                "data_asset_items": [
                    {"name": parts_out.get("name", ""), "description": parts_out.get("description", "")}
                    for parts_out in entity.get("data_product_version", {}).get("parts_out", [])
                ]
            }
        )

    return SearchDataProductsResponse(count=number_of_responses, data_products=products)


@service_registry.tool(
    name="data_product_search_data_products",
    description="""
    This tool searches all data products and data product drafts to return data products and data product drafts that match the given search_query.
    Example: 'Find all data products and data product drafts that match the name Environment.'
    In this case, product_search_query is 'Environment', and this tool returns all data products and data product drafts that have Environment in their name.
    Example: 'Find all data products and data product drafts in the Sustainability domain.'
    In this case, product_search_query is '*', search_filter_type is 'Domain' and search_filter_value is 'Sustainability'.
    
    Args:
        product_search_query (Literal["*"], str): The search query to search for data products and drafts. If the user wants to search for data products or drafts with a specific name, this is the name to search for. If user wants to search for all data products or drafts, this value should be "*".
        search_filter_type (Literal["None", "Domain"]): Specify what to filter by. It can be one of the following: None, Domain. If the user wants to filter by domain, then this value should be Domain otherwise None.
        search_filter_value (str): The value to filter by. For example, if search_filter_type is Domain, then this is the domain name to filter by. Use 'None' when search_filter_type is 'None'.
    """,
    tags={"search", "data_product"},
    meta={"version": "1.0", "service": "data_product"},
)
@auto_context
async def wxo_search_data_products(
    product_search_query: Union[Literal["*"], str], search_filter_type: Literal["None", "Domain"], search_filter_value: str
) -> SearchDataProductsResponse:
    """Watsonx Orchestrator compatible version that expands SearchDataProductsRequest object into individual parameters."""

    request = SearchDataProductsRequest(
        product_search_query=product_search_query,
        search_filter_type=search_filter_type,
        search_filter_value=search_filter_value,
    )

    # Call the original search_data_products function
    return await search_data_products(request)
