# Copyright [2026] [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.

import sqlglot
from sqlglot import errors
from sqlglot.expressions import Select, With

from app.core.auth import get_bss_account_id
from app.core.registry import service_registry
from app.services.constants import REPORTING_BASE_ENDPOINT
from app.services.reporting.models.sql_query_execution import (
    SqlQueryExecutionRequest,
    SqlQueryExecutionResponse,
    SqlResult,
)
from app.shared.exceptions.base import ExternalAPIError, ServiceError
from app.shared.logging.generate_context import auto_context
from app.shared.logging.utils import LOGGER
from app.shared.utils.tool_helper_service import tool_helper_service


def _validate_sql_query(sql_query: str) -> str:
    """
    Validate the SQL query input.

    Args:
        sql_query (str): The input SQL query.

    The validation enforces the following rules:
        - The SQL query must not be empty.
        - The query must not exceed 100,000 characters.
        - The query must not contain semicolons (to prevent multiple statements).
        - The query must start with "SELECT" or "WITH" (case-insensitive).

    Returns:
        str: The trimmed SQL query if valid.

    Raises:
        ServiceError: If the query is invalid.
    """
    if not sql_query:
        raise ServiceError("INVALID_QUERY: Query cannot be empty")

    trimmed = sql_query.strip()

    # Length safeguard
    if len(trimmed) > 100_000:
        LOGGER.info("INVALID_QUERY: SQL too long")
        raise ServiceError("INVALID_QUERY: SQL too long")

    try:
        # Parse the query using SQLGlot
        parsed = sqlglot.parse(trimmed)
    except errors.ParseError as e:
        LOGGER.info(f"INVALID_QUERY: Failed to parse SQL: {e}")
        raise ServiceError("INVALID_QUERY: SQL syntax error")

    # Check for multiple statements
    if len(parsed) != 1:
        LOGGER.info("INVALID_QUERY: Multiple SQL statements not allowed.")
        raise ServiceError("INVALID_QUERY: Multiple SQL statements not allowed.")

    # Check the type of the statement
    root_expr = parsed[0]

    if not isinstance(root_expr, (Select, With)):
        LOGGER.info("INVALID_QUERY: Only SELECT or WITH queries are supported.")
        raise ServiceError(
            "INVALID_QUERY: Only SELECT or WITH queries are supported."
        )
    return trimmed


# Placeholder constants for quote replacement
SINGLE_QUOTE_REPLACEMENT = "<!--__SPLACEHOLDER__-->"
DOUBLE_QUOTE_REPLACEMENT = "<!--__DPLACEHOLDER__-->"


def _unescape_sql_quotes(sql_query: str) -> str:
    """
    Restore quotes in SQL strings by replacing the placeholder back.

    Args:
        sql_query (str): The SQL query with placeholder strings.

    Returns:
        str: The SQL query with restored quotes.
    """
    return sql_query.replace(SINGLE_QUOTE_REPLACEMENT, "'").replace(
        DOUBLE_QUOTE_REPLACEMENT, '"'
    )


@service_registry.tool(
    name="reporting_sql_query_execution",
    description="Execute a user-provided SQL SELECT query against a tenant-specific reporting database. This tool enforces read-only access (only SELECT statements) and includes basic query validation for safety.",
)
@auto_context
async def sql_query_execution(
    request: SqlQueryExecutionRequest,
) -> SqlQueryExecutionResponse:
    """
    Execute a SQL SELECT query against a tenant-specific database.

    This tool executes a user-provided SQL SELECT query against a tenant-specific database.
    It is designed to work with queries generated by a Text-to-SQL LLM and enforces read-only
    access (only SELECT statements) with basic query validation for safety.

    Args:
        request (SqlQueryExecutionRequest): The request containing the SQL query to execute.

    Returns:
        SqlQueryExecutionResponse: The result of the SQL execution, containing either:
            - On success: status="success", sql_result with rows and executed query
            - On failure: status="failed", message with error details

    Raises:
        ServiceError: If the SQL query validation fails.
        ExternalAPIError: If the API call to execute the query fails.
    """
    # Unescape quotes in the SQL query
    sanitized_sql_query = _unescape_sql_quotes(request.sql_query)

    LOGGER.info(f"Validating SQL query: {sanitized_sql_query}")

    # Validate the SQL query
    sanitized_sql_query = _validate_sql_query(sanitized_sql_query)

    LOGGER.info(f"Executing SQL query: {sanitized_sql_query}")

    # Prepare the payload
    payload = {"sql": sanitized_sql_query}

    # Get tenant ID
    tenant_id = await get_bss_account_id()
    execute_sql_url = (
        f"{tool_helper_service.base_url}{REPORTING_BASE_ENDPOINT}/{tenant_id}/ikcquery/executesql"
    )

    try:
        # Execute the SQL query
        response = await tool_helper_service.execute_post_request(
            url=execute_sql_url, json=payload
        )
    except Exception as e:
        LOGGER.error(f"Failed to execute SQL query: {e}")
        raise ExternalAPIError(
            f"Failed to execute SQL query: {str(e)}",
            service="reporting",
            tool="sql_query_execution",
        )

    # Validate response structure
    if (
        not response
        or "sql_result" not in response
        or "rows" not in response.get("sql_result", {})
    ):
        LOGGER.error(
            f"Invalid response format: missing 'sql_result' or 'rows'. Response: {response}"
        )
        return SqlQueryExecutionResponse(
            status="failed",
            message=f"Invalid response format. Please verify the SQL query and data source: {response}",
        )

    LOGGER.info("Successfully executed SQL query")

    # Return successful response
    return SqlQueryExecutionResponse(
        status="success",
        sql_result=SqlResult(
            rows=response.get("sql_result", {}).get("rows", []),
            sql_query=request.sql_query,
        ),
    )


@service_registry.tool(
    name="reporting_sql_query_execution",
    description="Execute a user-provided SQL SELECT query against a tenant-specific reporting database. This tool enforces read-only access (only SELECT statements) and includes basic query validation for safety.",
)
@auto_context
async def wxo_sql_query_execution(sql_query: str) -> SqlQueryExecutionResponse:
    """
    Watsonx Orchestrator compatible version that expands SqlQueryExecutionRequest object into individual parameters.

    Args:
        sql_query (str): The SQL SELECT query to execute. Must be a syntactically valid and read-only query.

    Returns:
        SqlQueryExecutionResponse: The result of the SQL execution.
    """
    req = SqlQueryExecutionRequest(sql_query=sql_query)

    # Call the original sql_query_execution function
    return await sql_query_execution(req)