# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""GlueInteractiveSessionsHandler for Data Processing MCP Server."""

from awslabs.aws_dataprocessing_mcp_server.models.glue_models import (
    CancelStatementData,
    CreateSessionData,
    DeleteSessionData,
    GetSessionData,
    GetStatementData,
    ListSessionsData,
    ListStatementsData,
    RunStatementData,
    StopSessionData,
)
from awslabs.aws_dataprocessing_mcp_server.utils.aws_helper import AwsHelper
from awslabs.aws_dataprocessing_mcp_server.utils.logging_helper import (
    LogLevel,
    log_with_request_id,
)
from botocore.exceptions import ClientError
from mcp.server.fastmcp import Context
from mcp.types import CallToolResult, TextContent
from pydantic import Field
from typing import Annotated, Any, Dict, List, Optional


class GlueInteractiveSessionsHandler:
    """Handler for Amazon Glue Interactive Sessions operations."""

    def __init__(self, mcp, allow_write: bool = False, allow_sensitive_data_access: bool = False):
        """Initialize the Glue Interactive Sessions handler.

        Args:
            mcp: The MCP server instance
            allow_write: Whether to enable write access (default: False)
            allow_sensitive_data_access: Whether to allow access to sensitive data (default: False)
        """
        self.mcp = mcp
        self.allow_write = allow_write
        self.allow_sensitive_data_access = allow_sensitive_data_access
        self.glue_client = AwsHelper.create_boto3_client('glue')

        # Register tools
        self.mcp.tool(name='manage_aws_glue_sessions')(self.manage_aws_glue_sessions)
        self.mcp.tool(name='manage_aws_glue_statements')(self.manage_aws_glue_statements)

    async def manage_aws_glue_sessions(
        self,
        ctx: Context,
        operation: Annotated[
            str,
            Field(
                description='Operation to perform: create-session, delete-session, get-session, list-sessions, stop-session. Choose "get-session" or "list-sessions" for read-only operations when write access is disabled.',
            ),
        ],
        session_id: Annotated[
            Optional[str],
            Field(
                description='ID of the session (required for delete-session, get-session, and stop-session operations).',
            ),
        ] = None,
        description: Annotated[
            Optional[str],
            Field(
                description='Description of the session (optional for create-session operation).',
            ),
        ] = None,
        role: Annotated[
            Optional[str],
            Field(
                description='IAM Role ARN (required for create-session operation).',
            ),
        ] = None,
        command: Annotated[
            Optional[Dict[str, str]],
            Field(
                description="Session command with Name (e.g., 'glueetl', 'gluestreaming') and optional PythonVersion (required for create-session operation).",
            ),
        ] = None,
        timeout: Annotated[
            Optional[int],
            Field(
                description='Number of minutes before session times out (optional for create-session operation).',
            ),
        ] = None,
        idle_timeout: Annotated[
            Optional[int],
            Field(
                description='Number of minutes when idle before session times out (optional for create-session operation).',
            ),
        ] = None,
        default_arguments: Annotated[
            Optional[Dict[str, str]],
            Field(
                description='Map of key-value pairs for session arguments (optional for create-session operation).',
            ),
        ] = None,
        connections: Annotated[
            Optional[Dict[str, List[str]]],
            Field(
                description='Connections to use for the session (optional for create-session operation).',
            ),
        ] = None,
        max_capacity: Annotated[
            Optional[float],
            Field(
                description='Number of Glue data processing units (DPUs) to allocate (optional for create-session operation).',
            ),
        ] = None,
        number_of_workers: Annotated[
            Optional[int],
            Field(
                description='Number of workers to use for the session (optional for create-session operation).',
            ),
        ] = None,
        worker_type: Annotated[
            Optional[str],
            Field(
                description='Type of predefined worker (G.1X, G.2X, G.4X, G.8X, Z.2X) (optional for create-session operation).',
            ),
        ] = None,
        security_configuration: Annotated[
            Optional[str],
            Field(
                description='Name of the SecurityConfiguration structure (optional for create-session operation).',
            ),
        ] = None,
        glue_version: Annotated[
            Optional[str],
            Field(
                description='Glue version to use (must be greater than 2.0) (optional for create-session operation).',
            ),
        ] = None,
        tags: Annotated[
            Optional[Dict[str, str]],
            Field(
                description='Map of key-value pairs (tags) for the session (optional for create-session operation).',
            ),
        ] = None,
        request_origin: Annotated[
            Optional[str],
            Field(
                description='Origin of the request (optional for all operations).',
            ),
        ] = None,
        max_results: Annotated[
            Optional[int],
            Field(
                description='Maximum number of results to return for list-sessions operation.',
            ),
        ] = None,
        next_token: Annotated[
            Optional[str],
            Field(
                description='Pagination token for list-sessions operation.',
            ),
        ] = None,
    ) -> CallToolResult:
        """Manage AWS Glue Interactive Sessions for running Spark and Ray workloads.

        This tool provides operations for creating and managing Glue Interactive Sessions, which
        enable interactive development and execution of Spark ETL scripts and Ray applications.
        Interactive sessions provide a responsive environment for data exploration, debugging,
        and iterative development.

        ## Requirements
        - The server must be run with the `--allow-write` flag for create-session, delete-session, and stop-session operations
        - Appropriate AWS permissions for Glue Interactive Session operations

        ## Operations
        - **create-session**: Create a new interactive session with specified configuration
        - **delete-session**: Delete an existing interactive session
        - **get-session**: Retrieve detailed information about a specific session
        - **list-sessions**: List all interactive sessions with optional filtering
        - **stop-session**: Stop a running interactive session

        ## Example
        ```python
        # Create a new Spark ETL session
        {
            'operation': 'create-session',
            'session_id': 'my-spark-session',
            'role': 'arn:aws:iam::123456789012:role/GlueInteractiveSessionRole',
            'command': {'Name': 'glueetl', 'PythonVersion': '3'},
            'glue_version': '3.0',
        }
        ```

        Args:
            ctx: MCP context
            operation: Operation to perform
            session_id: ID of the session
            description: Description of the session
            role: IAM Role ARN
            command: Session command configuration
            timeout: Number of minutes before session times out
            idle_timeout: Number of minutes when idle before session times out
            default_arguments: Map of key-value pairs for session arguments
            connections: Connections to use for the session
            max_capacity: Number of Glue DPUs to allocate
            number_of_workers: Number of workers to use
            worker_type: Type of predefined worker
            security_configuration: Name of the SecurityConfiguration structure
            glue_version: Glue version to use
            tags: Map of key-value pairs (tags) for the session
            request_origin: Origin of the request
            max_results: Maximum number of results to return
            next_token: Pagination token

        Returns:
            CallToolResult with operation status and data
        """
        try:
            if not self.allow_write and operation not in [
                'get-session',
                'list-sessions',
            ]:
                error_message = f'Operation {operation} is not allowed without write access'
                log_with_request_id(ctx, LogLevel.ERROR, error_message)

                return CallToolResult(
                    isError=True,
                    content=[TextContent(type='text', text=error_message)],
                )

            if operation == 'create-session':
                if not role or not command:
                    raise ValueError('role and command are required for create-session operation')

                # Prepare create session parameters
                create_params = {
                    'Id': session_id,
                    'Role': role,
                    'Command': command,
                }

                # Add optional parameters if provided
                if description:
                    create_params['Description'] = description
                if timeout:
                    create_params['Timeout'] = timeout
                if idle_timeout:
                    create_params['IdleTimeout'] = idle_timeout
                if default_arguments:
                    create_params['DefaultArguments'] = default_arguments
                if connections:
                    create_params['Connections'] = connections
                if max_capacity:
                    create_params['MaxCapacity'] = max_capacity
                if number_of_workers:
                    create_params['NumberOfWorkers'] = number_of_workers
                if worker_type:
                    create_params['WorkerType'] = worker_type
                if security_configuration:
                    create_params['SecurityConfiguration'] = security_configuration
                if glue_version:
                    create_params['GlueVersion'] = glue_version

                # Add MCP management tags
                resource_tags = AwsHelper.prepare_resource_tags('GlueSession')

                # Merge user-provided tags with MCP tags
                if tags and isinstance(tags, dict):
                    merged_tags = dict(tags)
                    merged_tags.update(resource_tags)
                    create_params['Tags'] = merged_tags
                else:
                    create_params['Tags'] = resource_tags

                if request_origin:
                    create_params['RequestOrigin'] = request_origin

                # Create the session
                response = self.glue_client.create_session(**create_params)

                success_message = (
                    f'Successfully created session {response.get("Session", {}).get("Id", "")}'
                )
                data = CreateSessionData(
                    session_id=response.get('Session', {}).get('Id', ''),
                    session=response.get('Session', {}),
                    operation='create-session',
                )

                return CallToolResult(
                    isError=False,
                    content=[
                        TextContent(type='text', text=success_message),
                        TextContent(type='text', text=data.model_dump_json()),
                    ],
                )

            elif operation == 'delete-session':
                if session_id is None:
                    raise ValueError('session_id is required for delete-session operation')

                # First check if the session is managed by MCP
                try:
                    # Get the session to check if it's managed by MCP
                    get_params = {'Id': session_id}
                    if request_origin:
                        get_params['RequestOrigin'] = request_origin

                    response = self.glue_client.get_session(**get_params)
                    session = response.get('Session', {})
                    tags = session.get('Tags', {})

                    # Construct the ARN for the session
                    region = AwsHelper.get_or_default_aws_region() or 'us-east-1'
                    account_id = AwsHelper.get_aws_account_id()
                    session_arn = f'arn:aws:glue:{region}:{account_id}:session/{session_id}'

                    # Check if the session is managed by MCP
                    if not AwsHelper.is_resource_mcp_managed(self.glue_client, session_arn, {}):
                        error_message = f'Cannot delete session {session_id} - it is not managed by the MCP server (missing required tags)'
                        log_with_request_id(ctx, LogLevel.ERROR, error_message)
                        return CallToolResult(
                            isError=True,
                            content=[TextContent(type='text', text=error_message)],
                        )
                except ClientError as e:
                    if e.response['Error']['Code'] == 'EntityNotFoundException':
                        error_message = f'Session {session_id} not found'
                        log_with_request_id(ctx, LogLevel.ERROR, error_message)
                        return CallToolResult(
                            isError=True,
                            content=[TextContent(type='text', text=error_message)],
                        )
                    else:
                        raise e

                # Prepare delete session parameters
                delete_params = {'Id': session_id}
                if request_origin:
                    delete_params['RequestOrigin'] = request_origin

                # Delete the session
                response = self.glue_client.delete_session(**delete_params)

                success_message = f'Successfully deleted session {session_id}'
                data = DeleteSessionData(
                    session_id=session_id,
                    operation='delete-session',
                )

                return CallToolResult(
                    isError=False,
                    content=[
                        TextContent(type='text', text=success_message),
                        TextContent(type='text', text=data.model_dump_json()),
                    ],
                )

            elif operation == 'get-session':
                if session_id is None:
                    raise ValueError('session_id is required for get-session operation')

                # Prepare get session parameters
                get_params = {'Id': session_id}
                if request_origin:
                    get_params['RequestOrigin'] = request_origin

                # Get the session
                response = self.glue_client.get_session(**get_params)

                success_message = f'Successfully retrieved session {session_id}'
                data = GetSessionData(
                    session_id=session_id,
                    session=response.get('Session', {}),
                    operation='get-session',
                )

                return CallToolResult(
                    isError=False,
                    content=[
                        TextContent(type='text', text=success_message),
                        TextContent(type='text', text=data.model_dump_json()),
                    ],
                )

            elif operation == 'list-sessions':
                # Prepare list sessions parameters
                params: Dict[str, Any] = {}
                if max_results is not None:
                    params['MaxResults'] = str(max_results)
                if next_token is not None:
                    params['NextToken'] = next_token
                if tags:
                    params['Tags'] = tags
                if request_origin:
                    params['RequestOrigin'] = request_origin

                # List sessions
                response = self.glue_client.list_sessions(**params)

                success_message = 'Successfully retrieved sessions'
                data = ListSessionsData(
                    sessions=response.get('Sessions', []),
                    ids=response.get('Ids', []),
                    next_token=response.get('NextToken'),
                    count=len(response.get('Sessions', [])),
                    operation='list-sessions',
                )

                return CallToolResult(
                    isError=False,
                    content=[
                        TextContent(type='text', text=success_message),
                        TextContent(type='text', text=data.model_dump_json()),
                    ],
                )

            elif operation == 'stop-session':
                if session_id is None:
                    raise ValueError('session_id is required for stop-session operation')

                # First check if the session is managed by MCP
                try:
                    # Get the session to check if it's managed by MCP
                    get_params = {'Id': session_id}
                    if request_origin:
                        get_params['RequestOrigin'] = request_origin

                    response = self.glue_client.get_session(**get_params)
                    session = response.get('Session', {})
                    tags = session.get('Tags', {})

                    # Construct the ARN for the session
                    region = AwsHelper.get_or_default_aws_region() or 'us-east-1'
                    account_id = AwsHelper.get_aws_account_id()
                    session_arn = f'arn:aws:glue:{region}:{account_id}:session/{session_id}'

                    # Check if the session is managed by MCP
                    if not AwsHelper.is_resource_mcp_managed(self.glue_client, session_arn, {}):
                        error_message = f'Cannot stop session {session_id} - it is not managed by the MCP server (missing required tags)'
                        log_with_request_id(ctx, LogLevel.ERROR, error_message)
                        return CallToolResult(
                            isError=True,
                            content=[TextContent(type='text', text=error_message)],
                        )
                except ClientError as e:
                    if e.response['Error']['Code'] == 'EntityNotFoundException':
                        error_message = f'Session {session_id} not found'
                        log_with_request_id(ctx, LogLevel.ERROR, error_message)
                        return CallToolResult(
                            isError=True,
                            content=[TextContent(type='text', text=error_message)],
                        )
                    else:
                        raise e

                # Prepare stop session parameters
                stop_params = {'Id': session_id}
                if request_origin:
                    stop_params['RequestOrigin'] = request_origin

                # Stop the session
                response = self.glue_client.stop_session(**stop_params)

                success_message = f'Successfully stopped session {session_id}'
                data = StopSessionData(
                    session_id=session_id,
                    operation='stop-session',
                )

                return CallToolResult(
                    isError=False,
                    content=[
                        TextContent(type='text', text=success_message),
                        TextContent(type='text', text=data.model_dump_json()),
                    ],
                )

            else:
                error_message = f'Invalid operation: {operation}. Must be one of: create-session, delete-session, get-session, list-sessions, stop-session'
                log_with_request_id(ctx, LogLevel.ERROR, error_message)
                return CallToolResult(
                    isError=True,
                    content=[TextContent(type='text', text=error_message)],
                )

        except ValueError as e:
            log_with_request_id(ctx, LogLevel.ERROR, f'Parameter validation error: {str(e)}')
            raise
        except Exception as e:
            error_message = f'Error in manage_aws_glue_sessions: {str(e)}'
            log_with_request_id(ctx, LogLevel.ERROR, error_message)
            return CallToolResult(
                isError=True,
                content=[TextContent(type='text', text=error_message)],
            )

    async def manage_aws_glue_statements(
        self,
        ctx: Context,
        operation: Annotated[
            str,
            Field(
                description='Operation to perform: run-statement, cancel-statement, get-statement, list-statements. Choose "get-statement" or "list-statements" for read-only operations when write access is disabled.',
            ),
        ],
        session_id: Annotated[
            str,
            Field(
                description='ID of the session (required for all operations).',
            ),
        ],
        statement_id: Annotated[
            Optional[int],
            Field(
                description='ID of the statement (required for cancel-statement and get-statement operations).',
            ),
        ] = None,
        code: Annotated[
            Optional[str],
            Field(
                description='Code to execute for run-statement operation (up to 68000 characters).',
            ),
        ] = None,
        request_origin: Annotated[
            Optional[str],
            Field(
                description='Origin of the request (optional for all operations).',
            ),
        ] = None,
        max_results: Annotated[
            Optional[int],
            Field(
                description='Maximum number of results to return for list-statements operation.',
            ),
        ] = None,
        next_token: Annotated[
            Optional[str],
            Field(
                description='Pagination token for list-statements operation.',
            ),
        ] = None,
    ) -> CallToolResult:
        r"""Manage AWS Glue Interactive Session Statements for executing code and retrieving results.

        This tool provides operations for executing code, canceling running statements, and retrieving
        results within Glue Interactive Sessions. It enables interactive data processing, exploration,
        and analysis using Spark or Ray in AWS Glue.

        ## Requirements
        - The server must be run with the `--allow-write` flag for run-statement and cancel-statement operations
        - Appropriate AWS permissions for Glue Interactive Session Statement operations
        - A valid session ID is required for all operations

        ## Operations
        - **run-statement**: Execute code in an interactive session and get a statement ID
        - **cancel-statement**: Cancel a running statement by ID
        - **get-statement**: Retrieve detailed information and results of a specific statement
        - **list-statements**: List all statements in a session with their status

        ## Example
        ```python
        # Run a PySpark statement in a session
        {
            'operation': 'run-statement',
            'session_id': 'my-spark-session',
            'code': "df = spark.read.csv('s3://my-bucket/data.csv', header=True)\ndf.show(5)",
        }
        ```

        Args:
            ctx: MCP context
            operation: Operation to perform
            session_id: ID of the session
            statement_id: ID of the statement
            code: Code to execute for run-statement operation
            request_origin: Origin of the request
            max_results: Maximum number of results to return
            next_token: Pagination token

        Returns:
            CallToolResult with operation status and data
        """
        try:
            if not self.allow_write and operation not in [
                'get-statement',
                'list-statements',
            ]:
                error_message = f'Operation {operation} is not allowed without write access'
                log_with_request_id(ctx, LogLevel.ERROR, error_message)

                return CallToolResult(
                    isError=True,
                    content=[TextContent(type='text', text=error_message)],
                )

            if operation == 'run-statement':
                if code is None:
                    raise ValueError('code is required for run-statement operation')

                # Prepare run statement parameters
                run_params = {
                    'SessionId': session_id,
                    'Code': code,
                }
                if request_origin:
                    run_params['RequestOrigin'] = request_origin

                # Run the statement
                response = self.glue_client.run_statement(**run_params)

                success_message = f'Successfully ran statement in session {session_id}'
                data = RunStatementData(
                    session_id=session_id,
                    statement_id=response.get('Id', 0),
                    operation='run-statement',
                )

                return CallToolResult(
                    isError=False,
                    content=[
                        TextContent(type='text', text=success_message),
                        TextContent(type='text', text=data.model_dump_json()),
                    ],
                )

            elif operation == 'cancel-statement':
                if statement_id is None:
                    raise ValueError('statement_id is required for cancel-statement operation')

                # Prepare cancel statement parameters
                cancel_params = {
                    'SessionId': session_id,
                    'Id': statement_id,
                }
                if request_origin:
                    cancel_params['RequestOrigin'] = request_origin

                # Cancel the statement
                self.glue_client.cancel_statement(**cancel_params)

                success_message = (
                    f'Successfully canceled statement {statement_id} in session {session_id}'
                )
                data = CancelStatementData(
                    session_id=session_id,
                    statement_id=statement_id,
                    operation='cancel-statement',
                )

                return CallToolResult(
                    isError=False,
                    content=[
                        TextContent(type='text', text=success_message),
                        TextContent(type='text', text=data.model_dump_json()),
                    ],
                )

            elif operation == 'get-statement':
                if statement_id is None:
                    raise ValueError('statement_id is required for get-statement operation')

                # SECURITY: This operation returns statement execution output (customer data)
                # Require --allow-sensitive-data-access flag to prevent unauthorized data exposure
                if not self.allow_sensitive_data_access:
                    error_message = 'Operation get-statement returns execution output with customer data and requires --allow-sensitive-data-access flag'
                    log_with_request_id(ctx, LogLevel.ERROR, error_message)
                    return CallToolResult(
                        isError=True,
                        content=[TextContent(type='text', text=error_message)],
                    )

                # Prepare get statement parameters
                get_params = {
                    'SessionId': session_id,
                    'Id': statement_id,
                }
                if request_origin:
                    get_params['RequestOrigin'] = request_origin

                # Get the statement
                response = self.glue_client.get_statement(**get_params)

                success_message = (
                    f'Successfully retrieved statement {statement_id} in session {session_id}'
                )
                data = GetStatementData(
                    session_id=session_id,
                    statement_id=statement_id,
                    statement=response.get('Statement', {}),
                    operation='get-statement',
                )

                return CallToolResult(
                    isError=False,
                    content=[
                        TextContent(type='text', text=success_message),
                        TextContent(type='text', text=data.model_dump_json()),
                    ],
                )

            elif operation == 'list-statements':
                # Prepare list statements parameters
                params = {'SessionId': session_id}
                if max_results is not None:
                    params['MaxResults'] = str(max_results)
                if next_token is not None:
                    params['NextToken'] = next_token
                if request_origin:
                    params['RequestOrigin'] = request_origin

                # List statements
                response = self.glue_client.list_statements(**params)

                success_message = f'Successfully retrieved statements for session {session_id}'
                data = ListStatementsData(
                    session_id=session_id,
                    statements=response.get('Statements', []),
                    next_token=response.get('NextToken'),
                    count=len(response.get('Statements', [])),
                    operation='list-statements',
                )

                return CallToolResult(
                    isError=False,
                    content=[
                        TextContent(type='text', text=success_message),
                        TextContent(type='text', text=data.model_dump_json()),
                    ],
                )

            else:
                error_message = f'Invalid operation: {operation}. Must be one of: run-statement, cancel-statement, get-statement, list-statements'
                log_with_request_id(ctx, LogLevel.ERROR, error_message)
                return CallToolResult(
                    isError=True,
                    content=[TextContent(type='text', text=error_message)],
                )

        except ValueError as e:
            log_with_request_id(ctx, LogLevel.ERROR, f'Parameter validation error: {str(e)}')
            raise
        except Exception as e:
            error_message = f'Error in manage_aws_glue_statements: {str(e)}'
            log_with_request_id(ctx, LogLevel.ERROR, error_message)
            return CallToolResult(
                isError=True,
                content=[TextContent(type='text', text=error_message)],
            )
