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

"""CloudWatch Application Signals MCP Server - AWS client initialization."""

import boto3
import os
from . import __version__
from botocore.config import Config
from loguru import logger


# Get AWS region from environment variable or use default
AWS_REGION = os.environ.get('AWS_REGION', 'us-east-1')
logger.debug(f'Using AWS region: {AWS_REGION}')


def _initialize_aws_clients():
    """Initialize AWS clients with proper configuration."""
    # Add caller suffix if MCP_RUN_FROM is set
    mcp_source = os.environ.get('MCP_RUN_FROM')
    user_agent_suffix = f'/{mcp_source}' if mcp_source else ''

    config = Config(
        user_agent_extra=f'awslabs.cloudwatch-applicationsignals-mcp-server/{__version__}{user_agent_suffix}'
    )

    # Get endpoint URLs from environment variables
    applicationsignals_endpoint = os.environ.get('MCP_APPLICATIONSIGNALS_ENDPOINT')
    logs_endpoint = os.environ.get('MCP_LOGS_ENDPOINT')
    cloudwatch_endpoint = os.environ.get('MCP_CLOUDWATCH_ENDPOINT')
    xray_endpoint = os.environ.get('MCP_XRAY_ENDPOINT')
    synthetics_endpoint = os.environ.get('MCP_SYNTHETICS_ENDPOINT')
    rum_endpoint = os.environ.get('MCP_RUM_ENDPOINT')

    # Log endpoint overrides
    if applicationsignals_endpoint:
        logger.debug(f'Using Application Signals endpoint override: {applicationsignals_endpoint}')
    if logs_endpoint:
        logger.debug(f'Using CloudWatch Logs endpoint override: {logs_endpoint}')
    if cloudwatch_endpoint:
        logger.debug(f'Using CloudWatch endpoint override: {cloudwatch_endpoint}')
    if xray_endpoint:
        logger.debug(f'Using X-Ray endpoint override: {xray_endpoint}')
    if synthetics_endpoint:
        logger.debug(f'Using Synthetics endpoint override: {synthetics_endpoint}')
    if rum_endpoint:
        logger.debug(f'Using RUM endpoint override: {rum_endpoint}')

    # Check for AWS_PROFILE environment variable
    if aws_profile := os.environ.get('AWS_PROFILE'):
        logger.debug(f'Using AWS profile: {aws_profile}')
        session = boto3.Session(profile_name=aws_profile, region_name=AWS_REGION)
        logs = session.client('logs', config=config, endpoint_url=logs_endpoint)
        applicationsignals = session.client(
            'application-signals',
            region_name=AWS_REGION,
            config=config,
            endpoint_url=applicationsignals_endpoint,
        )
        cloudwatch = session.client('cloudwatch', config=config, endpoint_url=cloudwatch_endpoint)
        xray = session.client('xray', config=config, endpoint_url=xray_endpoint)
        synthetics = session.client('synthetics', config=config, endpoint_url=synthetics_endpoint)
        rum = session.client('rum', config=config, endpoint_url=rum_endpoint)
        s3 = session.client('s3', config=config)
        iam = session.client('iam', config=config)
        lambda_client = session.client('lambda', config=config)
        sts = session.client('sts', config=config)
    else:
        logs = boto3.client(
            'logs', region_name=AWS_REGION, config=config, endpoint_url=logs_endpoint
        )
        applicationsignals = boto3.client(
            'application-signals',
            region_name=AWS_REGION,
            config=config,
            endpoint_url=applicationsignals_endpoint,
        )
        cloudwatch = boto3.client(
            'cloudwatch', region_name=AWS_REGION, config=config, endpoint_url=cloudwatch_endpoint
        )
        xray = boto3.client(
            'xray', region_name=AWS_REGION, config=config, endpoint_url=xray_endpoint
        )
        synthetics = boto3.client(
            'synthetics', region_name=AWS_REGION, config=config, endpoint_url=synthetics_endpoint
        )
        rum = boto3.client('rum', region_name=AWS_REGION, config=config, endpoint_url=rum_endpoint)
        s3 = boto3.client('s3', region_name=AWS_REGION, config=config)
        iam = boto3.client('iam', region_name=AWS_REGION, config=config)
        lambda_client = boto3.client('lambda', region_name=AWS_REGION, config=config)
        sts = boto3.client('sts', region_name=AWS_REGION, config=config)

    logger.debug('AWS clients initialized successfully')
    return logs, applicationsignals, cloudwatch, xray, synthetics, s3, iam, lambda_client, sts, rum


# Initialize clients at module level
try:
    (
        logs_client,
        applicationsignals_client,
        cloudwatch_client,
        xray_client,
        synthetics_client,
        s3_client,
        iam_client,
        lambda_client,
        sts_client,
        rum_client,
    ) = _initialize_aws_clients()
except Exception as e:
    logger.error(f'Failed to initialize AWS clients: {str(e)}')
    raise
