# 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


def _resolve_region() -> str:
    """Resolve AWS region with priority: AWS_REGION > AWS_DEFAULT_REGION > profile/config > us-east-1.

    We check the env vars explicitly (rather than relying on boto3's own
    resolution) so the AWS_REGION > AWS_DEFAULT_REGION ordering is deterministic:
    when both are set, some boto3 versions return AWS_DEFAULT_REGION. Only when
    neither is set do we let boto3 resolve from the configured profile / ~/.aws/config,
    so a profile-only caller (AWS_PROFILE set, no env region) picks up that
    profile's region instead of silently defaulting to us-east-1.
    """
    env_region = os.environ.get('AWS_REGION') or os.environ.get('AWS_DEFAULT_REGION')
    if env_region:
        logger.debug(f'Region from AWS_REGION/AWS_DEFAULT_REGION env var: {env_region}')
        return env_region
    # Let boto3 resolve from AWS_PROFILE config or ~/.aws/config
    profile = os.environ.get('AWS_PROFILE')
    try:
        session = boto3.Session(profile_name=profile)
        if session.region_name:
            logger.debug(
                f'Region from AWS profile/config (profile={profile}): {session.region_name}'
            )
            return session.region_name
    except Exception as e:  # pragma: no cover - defensive; bad/missing profile
        logger.debug(f'Could not resolve region from boto3 session (profile={profile}): {e}')
    logger.debug(
        f'No region found in env or profile (profile={profile}), falling back to us-east-1'
    )
    return 'us-east-1'


AWS_REGION = _resolve_region()
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


def get_applicationsignals_client():
    """Return the module-level Application Signals client.

    Provided so callers (e.g. the service_events tools) can resolve the client lazily,
    which lets ``mock.patch`` of the module attribute propagate in tests.
    """
    return applicationsignals_client


def get_cloudwatch_client():
    """Return the module-level CloudWatch client (lazy accessor; see above)."""
    return cloudwatch_client


def get_logs_client():
    """Return the module-level CloudWatch Logs client (lazy accessor; see above)."""
    return logs_client
