# 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.

"""Startup Script Generator for AWS Serverless MCP Server.

Automatically generates appropriate startup scripts for different runtimes
to work with Lambda Web Adapter.
"""

import os
import re
import shlex
import stat
from loguru import logger
from typing import Dict, Optional


class EntryPointNotFoundError(Exception):
    """Custom error class for entry point not found errors."""

    def __init__(self, entry_point: str, built_artifacts_path: str):
        """Initializes the EntryPointNotFoundError with the specified entry point and built artifacts path.

        Args:
            entry_point (str): The name of the entry point file that was not found.
            built_artifacts_path (str): The path to the directory containing built artifacts.

        Raises:
            EntryPointNotFoundError: If the specified entry point file does not exist in the built artifacts path.
        """
        self.entry_point = entry_point
        self.built_artifacts_path = built_artifacts_path
        message = f'Entry point file not found: {os.path.join(built_artifacts_path, entry_point)}'
        super().__init__(message)
        self.name = 'EntryPointNotFoundError'


async def generate_startup_script(
    runtime: str,
    entry_point: str,
    built_artifacts_path: str,
    startup_script_name: Optional[str] = None,
    additional_env: Optional[Dict[str, str]] = None,
) -> str:
    """Generate a startup script based on runtime and entry point. This script starts up your web server so that beings listening for requests.

    Args:
        runtime: Lambda runtime (e.g., nodejs22.x, python3.13)
        entry_point: Application entry point
        built_artifacts_path: Path to the built artifacts
        startup_script_name: Name of the startup script (default: 'bootstrap')
        additional_env: Additional environment variables

    Returns:
        str: Path to the generated startup script

    Raises:
        EntryPointNotFoundError: If the entry point file doesn't exist
    """
    startup_script_name = startup_script_name or get_default_startup_script_name(runtime)
    script_path = os.path.join(built_artifacts_path, startup_script_name)

    logger.info(f'Generating startup script for runtime: {runtime}, entry point: {entry_point}')

    # Validate entry_point BEFORE checking file existence
    validate_entry_point(entry_point)

    # Validate environment variables if provided
    if additional_env:
        validate_environment_variables(additional_env)

    # Check for path traversal attacks
    entry_point_path = os.path.join(built_artifacts_path, entry_point)
    resolved_entry_point = os.path.realpath(entry_point_path)
    resolved_artifacts_path = os.path.realpath(built_artifacts_path)

    if not resolved_entry_point.startswith(resolved_artifacts_path + os.sep):
        raise ValueError(
            'Path traversal detected: entry_point resolves outside built_artifacts_path'
        )

    # Check if entry point exists
    if not os.path.exists(entry_point_path):
        error = EntryPointNotFoundError(entry_point, built_artifacts_path)
        logger.error(error.args[0])

        # Provide helpful suggestions
        logger.info('Available files in the artifacts directory:')
        try:
            files = os.listdir(built_artifacts_path)
            if len(files) == 0:
                logger.info('  (directory is empty)')
            else:
                for file in files:
                    logger.info(f'  - {file}')
        except Exception:
            logger.error(f'Could not read directory: {built_artifacts_path}')

        raise error

    # Generate script content based on runtime
    script_content = generate_script_content(runtime, entry_point, additional_env)

    # Write script to file
    with open(script_path, 'w', encoding='utf-8') as f:
        f.write(script_content)

    # Make script executable
    os.chmod(
        script_path, os.stat(script_path).st_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH
    )

    logger.info(f'Startup script generated at: {script_path}')

    return startup_script_name


def get_default_startup_script_name(runtime: str) -> str:
    """Get default startup script name for a runtime.

    Args:
        runtime: Lambda runtime

    Returns:
        str: Default startup script name
    """
    # Lambda expects a file named "bootstrap" for custom runtimes
    return 'bootstrap'


def validate_entry_point(entry_point: str) -> None:
    """Validate entry point against allowlist regex to prevent command injection.

    Args:
        entry_point: Application entry point to validate

    Raises:
        InvalidEntryPointError: If entry point contains invalid characters
    """
    # Allowlist: only alphanumeric, dots, underscores, hyphens, and forward slashes
    if not re.match(r'^[a-zA-Z0-9._/-]+$', entry_point):
        raise ValueError(
            'Entry point contains invalid characters. Only alphanumeric characters, dots, underscores, hyphens, and forward slashes are allowed.'
        )


def validate_environment_variables(additional_env: Dict[str, str]) -> None:
    """Validate environment variable keys and values.

    Args:
        additional_env: Dictionary of environment variables to validate

    Raises:
        InvalidEnvironmentVariableError: If any key or value is invalid
    """
    for key, value in additional_env.items():
        # Environment variable keys should follow POSIX naming conventions
        if not re.match(r'^[a-zA-Z_][a-zA-Z0-9_]*$', key):
            raise ValueError(
                'Environment variable key must start with a letter or underscore and contain only alphanumeric characters and underscores'
            )

        # Check for null bytes in values which can cause issues
        if '\0' in value:
            raise ValueError('Environment variable value contains null bytes')


def generate_script_content(
    runtime: str, entry_point: str, additional_env: Optional[Dict[str, str]] = None
) -> str:
    """Generate script content based on runtime and entry point.

    Args:
        runtime: Lambda runtime
        entry_point: Application entry point
        additional_env: Additional environment variables

    Returns:
        str: Script content with properly escaped values
    """
    # Sanitize entry_point to prevent command injection
    safe_entry_point = shlex.quote(entry_point)

    # Generate environment variables setup with proper escaping
    env_setup = ''
    if additional_env:
        env_vars = []
        for key, value in additional_env.items():
            # Key is already validated, but value needs escaping
            safe_value = shlex.quote(value)
            env_vars.append(f'export {key}={safe_value}')
        env_setup = '\n'.join(env_vars) + '\n\n'

    if runtime.startswith('nodejs'):
        return f"""#!/bin/bash
{env_setup}# Start the application
exec node {safe_entry_point}
"""
    elif runtime.startswith('python'):
        return f"""#!/bin/bash
{env_setup}# Start the application
exec python {safe_entry_point}
"""
    elif runtime.startswith('java'):
        # Determine if it's a JAR file or a class
        is_jar = entry_point.lower().endswith('.jar')

        if is_jar:
            return f"""#!/bin/bash
{env_setup}# Start the application
exec java -jar {safe_entry_point}
"""
        else:
            return f"""#!/bin/bash
{env_setup}# Start the application
exec java {safe_entry_point}
"""
    elif runtime.startswith('dotnet'):
        return f"""#!/bin/bash
{env_setup}# Start the application
exec dotnet {safe_entry_point}
"""
    elif runtime.startswith('go'):
        return f"""#!/bin/bash
{env_setup}# Start the application
exec ./{safe_entry_point}
"""
    elif runtime.startswith('ruby'):
        return f"""#!/bin/bash
{env_setup}# Start the application
exec ruby {safe_entry_point}
"""
    else:
        # Generic script for unknown runtimes
        return f"""#!/bin/bash
{env_setup}# Start the application
exec {safe_entry_point}
"""
