# 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.
from app.core.registry import service_registry
from app.services.data_protection_rules.models.search_rule import (
    SearchDataProtectionRuleRequest,
    SearchDataProtectionRuleResponse,
)
from app.shared.exceptions.base import ExternalAPIError, ServiceError
from app.core.auth import get_access_token
from app.shared.utils.http_client import get_http_client
from app.services.constants import JSON_CONTENT_TYPE
from app.core.settings import settings, ENV_MODE_SAAS
from app.shared.logging import LOGGER, auto_context
from app.services.constants import SEARCH_PATH

METADATA_NAME = "metadata.name"
METADAT_DESCRIPTION = "metadata.description"

@service_registry.tool(
    name="data_protection_rule_search",
    description="""
    This tool searches all data protection rules to return data protection rules that match the given search query.
    Example: 'Find all data protection rules with Deny name.'
    In this case, data_protection_rule_search_query is 'Deny', and this tool returns all data protection rules that have Deny in their name or description.
    Example: 'Show me all data protection rules'
    In this case, data_protection_rule_search_query is '*'.
    """,
    tags={"search", "data_protection_rules"},
    meta={"version": "1.0", "service": "data_protection_rules"},
)
@auto_context
async def search_rules(
    request: SearchDataProtectionRuleRequest,
) -> SearchDataProtectionRuleResponse:
    LOGGER.info(
        f"In the data_protection_rule_search tool, searching data protection rules with query {request.data_protection_rule_search_query}."
    )
    token = await get_access_token()

    payload = get_dps_search_payload(
        data_protection_rule_search_query=request.data_protection_rule_search_query,
    )

    headers = {
        "Content-Type": JSON_CONTENT_TYPE,
        "Authorization": token,
    }
    client = get_http_client()

    try:
        response = await client.post(
            url=f"{settings.di_service_url}{SEARCH_PATH}?role=viewer&auth_scope=all",
            headers=headers,
            data=payload,
        )

        number_of_responses = response["size"]
        if number_of_responses == 0:
            LOGGER.info(
                "In the data_protection_rule_search tool, no data protection rules found."
            )
            return SearchDataProtectionRuleResponse(count=0, data_protection_rules=[])
        LOGGER.info(f"Found {number_of_responses} data protection rules.")
        data_protection_rules = []

        if settings.di_env_mode.upper() == ENV_MODE_SAAS:
            url_prefix = settings.di_service_url.replace("https://api.", "https://") + "/governance/rules/dataProtection/view/"
        else:
            url_prefix = settings.di_service_url + "/gov/rules/dataProtection/view/"
        for row in response["rows"]:
            data_protection_rules.append(
                {
                    "name": row.get("metadata", {}).get("name", ""),
                    "description": row.get("metadata", {}).get("description", ""),
                    "modified_on": row.get("metadata", {}).get("modified_on", ""),
                    "url": url_prefix + row.get("artifact_id", "")
                }
            )
    except ExternalAPIError as e:
        LOGGER.error(
            f"Failed to run data_protection_rule_search tool. Error while searching data protection rules: {str(e)}"
        )
        raise ExternalAPIError(
            f"Failed to run data_protection_rule_search tool. Error while searching data protection rules: {str(e)}"
        )
    except Exception as e:
        LOGGER.error(
            f"Failed to run data_protection_rule_search tool. Error while searching data protection rules: {str(e)}"
        )
        raise ServiceError(
            f"Failed to run data_protection_rule_search tool. Error while searching data protection rules: {str(e)}"
        )

    return SearchDataProtectionRuleResponse(count=number_of_responses, data_protection_rules=data_protection_rules)


@service_registry.tool(
    name="data_protection_rule_search",
    description="""
    This tool searches all data protection rules to return data protection rules that match the given search query.
    Example: 'Find all data protection rules with Deny name.'
    In this case, data_protection_rule_search_query is 'Deny', and this tool returns all data protection rules that have Deny in their name or description.
    Example: 'Show me all data protection rules'
    In this case, data_protection_rule_search_query is '*'.
    """,
    tags={"search", "data_protection_rules"},
    meta={"version": "1.0", "service": "data_protection_rules"},
)
@auto_context
async def wxo_search_rule(
    data_protection_rule_search_query: str
) -> SearchDataProtectionRuleResponse:
    """Watsonx Orchestrator compatible version that expands SearchDataProductsRequest object into individual parameters."""

    request = SearchDataProtectionRuleRequest(
        data_protection_rule_search_query=data_protection_rule_search_query,
    )

    # Call the original search_data_protection_rules function
    return await search_rules(request)

def get_dps_search_payload(data_protection_rule_search_query: str) -> dict:
    if data_protection_rule_search_query == "*":
        return {
            "size": 10000,
            "from": "0",
            "_source": [
                "artifact_id",
                METADATA_NAME,
                METADAT_DESCRIPTION,
                "metadata.modified_on"
            ],
            "query": {
                "bool": {
                    "must": [
                        {"match": {"provider_type_id": "dps"}},
                        {"match": {"metadata.artifact_type": "data_protection_rule"}}
                    ]
                }
            }
        }
    else:
        return {
        "size": 10000,
        "from": "0",
        "_source": [
            "artifact_id",
            METADATA_NAME,
            METADAT_DESCRIPTION,
            "metadata.modified_on"
        ],
        "query": {
            "bool": {
                "must": [
                    {"match": {"provider_type_id": "dps"}},
                    {"match": {"metadata.artifact_type": "data_protection_rule"}},
                    {
                            "gs_user_query": {
                            "search_string": data_protection_rule_search_query,
                            "semantic_search_enabled": True
                        }
                    }
                ]
            }
        }
    }
