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


"""awslabs Timestream for InfluxDB MCP Server implementation."""

import boto3
import os
from awslabs.timestream_for_influxdb_mcp_server import __version__
from botocore.config import Config
from influxdb_client.client.influxdb_client import InfluxDBClient
from influxdb_client.client.write.point import Point
from influxdb_client.client.write_api import ASYNCHRONOUS, SYNCHRONOUS
from influxdb_client.domain.bucket_retention_rules import BucketRetentionRules
from influxdb_client.domain.write_precision import WritePrecision
from loguru import logger
from mcp.server.fastmcp import FastMCP
from pydantic import Field
from typing import Any, Dict, List, Optional
from urllib.parse import urlparse


USER_AGENT_EXTRA = f'md/awslabs#mcp#timestream-for-influxdb-mcp-server#{__version__}'
_config = Config(user_agent_extra=USER_AGENT_EXTRA)


# InfluxDB environment variables for data plane operations
INFLUXDB_TOKEN = os.environ.get('INFLUXDB_TOKEN')
INFLUXDB_URL = os.environ.get('INFLUXDB_URL')
INFLUXDB_ORG = os.environ.get('INFLUXDB_ORG')

# Define Field parameters as global variables to avoid duplication
# Common fields
REQUIRED_FIELD_DB_CLUSTER_ID = Field(
    ..., description='Service-generated unique identifier of the DB cluster.'
)

REQUIRED_FIELD_DB_INSTANCE_NAME = Field(
    ...,
    description='The name that uniquely identifies the DB instance. '
    'This name will also be a prefix included in the endpoint. '
    'DB instance names must be unique per customer and per region.',
)
REQUIRED_FIELD_DB_INSTANCE_TYPE = Field(
    ..., description='The Timestream for InfluxDB DB instance type to run InfluxDB on.'
)

OPTIONAL_FIELD_DB_INSTANCE_TYPE_CLUSTER_UPDATE = Field(
    None, description='Update the DB cluster to use the specified DB instance Type.'
)

REQUIRED_FIELD_PASSWORD = Field(
    ...,
    description='The password of the initial admin user created in InfluxDB. '
    'This password will allow you to access the InfluxDB UI to perform various administrative task '
    'and also use the InfluxDB CLI to create an operator token.',
)
REQUIRED_FIELD_ALLOCATED_STORAGE_GB = Field(
    ...,
    description='The amount of storage to allocate for your DB storage type in GiB (gibibytes).',
)
OPTIONAL_FIELD_ALLOCATED_STORAGE_GB_OPTIONAL = Field(
    None, description='The amount of storage to allocate for your DB storage type (in gibibytes).'
)
REQUIRED_FIELD_VPC_SECURITY_GROUP_IDS = Field(
    ..., description='A list of VPC security group IDs to associate with the DB.'
)

REQUIRED_FIELD_VPC_SUBNET_IDS = Field(
    ...,
    description='A list of VPC subnet IDs to associate with the DB. '
    'Provide at least two VPC subnet IDs in different Availability Zones when deploying with a Multi-AZ standby.',
)

OPTIONAL_FIELD_PUBLICLY_ACCESSIBLE = Field(
    True,
    description='Configures the DB with a public IP to facilitate access from outside the VPC.',
)

OPTIONAL_FIELD_TOOL_WRITE_MODE = Field(
    False,
    description='Tool is run in write mode and will be able to perform any create/update/delete operations. '
    'Default is read-only mode (False)',
)

OPTIONAL_FIELD_USERNAME = Field(
    None, description='The username of the initial admin user created in InfluxDB.'
)
OPTIONAL_FIELD_ORGANIZATION = Field(
    None,
    description='The name of the initial organization for the initial admin user in InfluxDB.'
    'An InfluxDB organization is a workspace for a group of users',
)
REQUIRED_FIELD_BUCKET = Field(..., description='The name of the initial InfluxDB bucket.')
OPTIONAL_FIELD_BUCKET = Field(None, description='The name of the initial InfluxDB bucket.')
OPTIONAL_FIELD_DB_STORAGE_TYPE = Field(
    None,
    description='The Timestream for InfluxDB DB storage type to read and write InfluxDB data.',
)
OPTIONAL_FIELD_DEPLOYMENT_TYPE_INSTANCE = Field(
    None,
    description='Specifies whether the DB instance will be deployed as a standalone instance or with a Multi-AZ standby for high availability.',
)
OPTIONAL_FIELD_NETWORK_TYPE = Field(
    None,
    description='Specifies whether the network type of the Timestream for InfluxDB cluster is IPv4 or DUAL.',
)

OPTIONAL_FIELD_PORT = Field(
    None, description='The port number on which InfluxDB accepts connections. Default: 8086'
)
OPTIONAL_FIELD_PORT_UPDATE = Field(
    None, description='Update the DB cluster to use the specified port.'
)

OPTIONAL_FIELD_FAILOVER_MODE = Field(
    None,
    description='Specifies the behavior of failure recovery when the primary node of the cluster fails.',
)
OPTIONAL_FIELD_FAILOVER_MODE_UPDATE = Field(
    None, description="Update the DB cluster's failover behavior."
)

OPTIONAL_FIELD_TAGS = Field(None, description='A list of tags to assign to the DB.')
OPTIONAL_FIELD_TAGS_PARAM_GROUP = Field(
    None, description='A list of key-value pairs to associate with the DB parameter group.'
)
OPTIONAL_FIELD_LOG_DELIVERY_CONFIGURATION = Field(
    None, description='Configuration for sending InfluxDB engine logs to a specified S3 bucket.'
)
OPTIONAL_FIELD_LOG_DELIVERY_CONFIGURATION_UPDATE = Field(
    None, description='The log delivery configuration to apply to the DB cluster.'
)

# Pagination fields
OPTIONAL_FIELD_NEXT_TOKEN = Field(
    None,
    description='The pagination token. To resume pagination, provide the next-token value as an argument of a subsequent API invocation.',
)

OPTIONAL_FIELD_MAX_RESULTS = Field(
    None,
    description='The maximum number of items to return in the output. If the total number of items available is more than the value specified, a nextToken is provided in the output.',
)

# Resource fields
REQUIRED_FIELD_RESOURCE_ARN = Field(
    ..., description='The Amazon Resource Name (ARN) of the tagged resource.'
)
REQUIRED_FIELD_TAG_KEYS = Field(..., description='The keys used to identify the tags to remove.')
REQUIRED_FIELD_TAGS_RESOURCE = Field(..., description='A list of key-value pairs as tags.')

# DB Parameter Group fiels
REQUIRED_FIELD_PARAMETER_GROUP_ID = Field(..., description='The id of the DB parameter group.')
REQUIRED_FIELD_PARAM_GROUP_NAME = Field(
    ...,
    description='The name of the DB parameter group. The name must be unique per customer and per region.',
)
OPTIONAL_FIELD_PARAM_GROUP_DESCRIPTION = Field(
    None, description='A description of the DB parameter group.'
)
OPTIONAL_FIELD_PARAMETERS = Field(
    None, description='A list of the parameters that comprise the DB parameter group.'
)
OPTIONAL_FIELD_DB_PARAMETER_GROUP_ID = Field(
    None, description='The id of the DB parameter group to assign to your DB.'
)
OPTIONAL_FIELD_DB_PARAMETER_GROUP_IDENTIFIER_UPDATE = Field(
    None, description='Update the DB cluster to use the specified DB parameter group.'
)

# DB Instance fields
REQUIRED_FIELD_DB_INSTANCE_IDENTIFIER = Field(..., description='The id of the DB instance.')

# Status fields
REQUIRED_FIELD_STATUS = Field(
    ..., description='The status to filter DB instances by (case-insensitive).'
)
REQUIRED_FIELD_STATUS_CLUSTER = Field(
    ..., description='The status to filter DB clusters by (case-insensitive).'
)

# InfluxDB fields
OPTIONAL_FIELD_URL = Field(
    None,
    description='The URL of the InfluxDB server. Falls back to INFLUXDB_URL env var if not provided.',
)
OPTIONAL_FIELD_TOKEN = Field(
    None,
    description='The authentication token. Falls back to INFLUXDB_TOKEN env var if not provided.',
)
REQUIRED_FIELD_BUCKET_INFLUX = Field(..., description='The destination bucket for writes.')
OPTIONAL_FIELD_ORG = Field(
    None,
    description='The organization name. Falls back to INFLUXDB_ORG env var if not provided.',
)
REQUIRED_FIELD_POINTS = Field(
    ...,
    description='List of data points to write. Each point should be a dictionary with measurement, tags, fields, and optional time.',
)
REQUIRED_FIELD_DATA_LINE_PROTOCOL = Field(
    ..., description='Data in InfluxDB Line Protocol format.'
)
OPTIONAL_FIELD_WRITE_PRECISION = Field(
    default='ns',
    description='The precision for the unix timestamps within the body line-protocol. One of: ns, us, ms, s (default is ns).',
)
OPTIONAL_FIELD_SYNC_MODE = Field(
    default='synchronous',
    description="The synchronization mode, either 'synchronous' or 'asynchronous'.",
)
OPTIONAL_FIELD_VERIFY_SSL = Field(
    True, description='Whether to verify SSL with https connections.'
)
REQUIRED_FIELD_QUERY = Field(..., description='The Flux query string.')

# Cluster name field
REQUIRED_FIELD_CLUSTER_NAME = Field(
    ...,
    description='The name that uniquely identifies the DB cluster when interacting with '
    'the Amazon Timestream for InfluxDB API and CLI commands. '
    'This name will also be a prefix included in the endpoint.',
)

mcp = FastMCP(
    'awslabs.timestream-for-influxdb-mcp-server',
    instructions="""
    This MCP server provides tools to interact with AWS Timestream for InfluxDB APIs.
    It allows you to create and manage databases, users, and perform other operations
    related to Timestream for InfluxDB service.
    """,
    dependencies=['loguru', 'boto3', 'influxdb-client'],
)


def get_timestream_influxdb_client():
    """Get the AWS Timestream for InfluxDB client."""
    aws_region: str = os.environ.get('AWS_REGION', 'us-east-1')
    aws_profile = os.environ.get('AWS_PROFILE')
    try:
        if aws_profile:
            logger.info(f'Using AWS profile for AWS Timestream Influx Client: {aws_profile}')
            client = boto3.Session(profile_name=aws_profile, region_name=aws_region).client(
                'timestream-influxdb', config=_config
            )
        else:
            client = boto3.Session(region_name=aws_region).client(
                'timestream-influxdb', config=_config
            )
    except Exception as e:
        logger.error(f'Error creating AWS Timestream for InfluxDB client: {str(e)}')
        raise

    return client


def resolve_influxdb_config(
    url: Optional[str] = None,
    token: Optional[str] = None,
    org: Optional[str] = None,
    require_org: bool = True,
) -> tuple:
    """Resolve InfluxDB configuration from parameters or environment variables.

    Args:
        url: The URL of the InfluxDB server (optional, falls back to INFLUXDB_URL env var).
        token: The authentication token (optional, falls back to INFLUXDB_TOKEN env var).
        org: The organization name (optional, falls back to INFLUXDB_ORG env var).
        require_org: Whether the organization is required (default True).

    Returns:
        A tuple of (resolved_url, resolved_token, resolved_org).

    Raises:
        ValueError: If required parameters are not provided.
    """
    resolved_url = url or INFLUXDB_URL
    resolved_token = token or INFLUXDB_TOKEN
    resolved_org = org or INFLUXDB_ORG

    if not resolved_url:
        raise ValueError(
            'URL must be provided either as parameter or via INFLUXDB_URL environment variable'
        )
    if not resolved_token:
        raise ValueError(
            'Token must be provided either as parameter or via INFLUXDB_TOKEN environment variable'
        )
    if require_org and not resolved_org:
        raise ValueError(
            'Organization must be provided either as parameter or via INFLUXDB_ORG environment variable'
        )

    return resolved_url, resolved_token, resolved_org


def get_influxdb_client(url, token, org=None, timeout=10000, verify_ssl: bool = True):
    """Get an InfluxDB client.

    Args:
        url: The URL of the InfluxDB server e.g. https://<host-name>:8086.
        token: The authentication token.
        org: The organization name.
        timeout: The timeout in milliseconds.
        verify_ssl: whether to verify SSL with https connections

    Returns:
        An InfluxDB client.

    Raises:
        ValueError: If the URL does not use HTTPS protocol or is not properly formatted.
    """
    try:
        parsed_url = urlparse(url)
        url_scheme = parsed_url.scheme
        if url_scheme != 'https' and url_scheme != 'http':
            raise ValueError('URL must use HTTP(S) protocol')
    except Exception as e:
        logger.error(f'Error parsing URL: {str(e)}')
        raise

    if not token:
        raise ValueError('Token must be provided')

    # Ensure org is not None when passed to InfluxDBClient
    org_param = org if org is not None else ''

    return InfluxDBClient(
        url=url, token=token, org=org_param, timeout=timeout, verify_ssl=verify_ssl
    )


@mcp.tool(
    name='CreateDbCluster', description='Create a new Timestream for InfluxDB database cluster.'
)
async def create_db_cluster(
    name: str = REQUIRED_FIELD_CLUSTER_NAME,
    db_instance_type: str = REQUIRED_FIELD_DB_INSTANCE_TYPE,
    password: str = REQUIRED_FIELD_PASSWORD,
    allocated_storage_gb: int = REQUIRED_FIELD_ALLOCATED_STORAGE_GB,
    vpc_security_group_ids: List[str] = REQUIRED_FIELD_VPC_SECURITY_GROUP_IDS,
    vpc_subnet_ids: List[str] = REQUIRED_FIELD_VPC_SUBNET_IDS,
    publicly_accessible: bool = OPTIONAL_FIELD_PUBLICLY_ACCESSIBLE,
    username: Optional[str] = OPTIONAL_FIELD_USERNAME,
    organization: Optional[str] = OPTIONAL_FIELD_ORGANIZATION,
    bucket: Optional[str] = OPTIONAL_FIELD_BUCKET,
    db_storage_type: Optional[str] = OPTIONAL_FIELD_DB_STORAGE_TYPE,
    deployment_type: Optional[str] = OPTIONAL_FIELD_DEPLOYMENT_TYPE_INSTANCE,
    networkType: Optional[str] = OPTIONAL_FIELD_NETWORK_TYPE,
    port: Optional[int] = OPTIONAL_FIELD_PORT,
    db_parameter_group_identifier: Optional[str] = OPTIONAL_FIELD_DB_PARAMETER_GROUP_ID,
    failover_mode: Optional[str] = OPTIONAL_FIELD_FAILOVER_MODE,
    tags: Optional[Dict[str, str]] = OPTIONAL_FIELD_TAGS,
    log_delivery_configuration: Optional[
        Dict[str, Any]
    ] = OPTIONAL_FIELD_LOG_DELIVERY_CONFIGURATION,
    tool_write_mode: bool = OPTIONAL_FIELD_TOOL_WRITE_MODE,
) -> Dict[str, Any]:
    """Create a new Timestream for InfluxDB database cluster.

    API reference: https://docs.aws.amazon.com/ts-influxdb/latest/ts-influxdb-api/API_CreateDbCluster.html

    Returns:
        Details of the created DB cluster.
    """
    if not tool_write_mode:
        raise Exception(
            'CreateDbCluster tool invocation not allowed when tool-write-mode is set to False'
        )

    ts_influx_client = get_timestream_influxdb_client()

    # Required parameters
    params = {
        'name': name,
        'dbInstanceType': db_instance_type,
        'password': password,
        'vpcSecurityGroupIds': vpc_security_group_ids,
        'vpcSubnetIds': vpc_subnet_ids,
        'allocatedStorage': allocated_storage_gb,
        'publiclyAccessible': publicly_accessible,
    }

    # Add optional parameters if provided
    if db_parameter_group_identifier:
        params['dbParameterGroupIdentifier'] = db_parameter_group_identifier
    if username:
        params['username'] = username
    if organization:
        params['organization'] = organization
    if bucket:
        params['bucket'] = bucket
    if port:
        params['port'] = port
    if db_storage_type:
        params['dbStorageType'] = db_storage_type
    if deployment_type:
        params['deploymentType'] = deployment_type
    if networkType:
        params['networkType'] = networkType
    if failover_mode:
        params['failoverMode'] = failover_mode
    if log_delivery_configuration:
        params['logDeliveryConfiguration'] = str(log_delivery_configuration)

    if tags:
        tag_list = [{'Key': k, 'Value': v} for k, v in tags.items()]
        params['tags'] = str(tag_list)

    try:
        response = ts_influx_client.create_db_cluster(**params)
        return response
    except Exception as e:
        logger.error(f'Error creating DB cluster: {str(e)}')
        raise e


@mcp.tool(
    name='CreateDbInstance', description='Create a new Timestream for InfluxDB database instance'
)
async def create_db_instance(
    db_instance_name: str = REQUIRED_FIELD_DB_INSTANCE_NAME,
    db_instance_type: str = REQUIRED_FIELD_DB_INSTANCE_TYPE,
    password: str = REQUIRED_FIELD_PASSWORD,
    allocated_storage_gb: int = REQUIRED_FIELD_ALLOCATED_STORAGE_GB,
    vpc_security_group_ids: List[str] = REQUIRED_FIELD_VPC_SECURITY_GROUP_IDS,
    vpc_subnet_ids: List[str] = REQUIRED_FIELD_VPC_SUBNET_IDS,
    publicly_accessible: bool = OPTIONAL_FIELD_PUBLICLY_ACCESSIBLE,
    username: Optional[str] = OPTIONAL_FIELD_USERNAME,
    organization: Optional[str] = OPTIONAL_FIELD_ORGANIZATION,
    bucket: Optional[str] = OPTIONAL_FIELD_BUCKET,
    db_storage_type: Optional[str] = OPTIONAL_FIELD_DB_STORAGE_TYPE,
    deployment_type: Optional[str] = OPTIONAL_FIELD_DEPLOYMENT_TYPE_INSTANCE,
    networkType: Optional[str] = OPTIONAL_FIELD_NETWORK_TYPE,
    port: Optional[int] = OPTIONAL_FIELD_PORT,
    db_parameter_group_id: Optional[str] = OPTIONAL_FIELD_DB_PARAMETER_GROUP_ID,
    tags: Optional[Dict[str, str]] = OPTIONAL_FIELD_TAGS,
    tool_write_mode: bool = OPTIONAL_FIELD_TOOL_WRITE_MODE,
) -> Dict[str, Any]:
    """Create a new Timestream for InfluxDB database instance.

    API reference: https://docs.aws.amazon.com/ts-influxdb/latest/ts-influxdb-api/API_CreateDbInstance.html#tsinfluxdb-CreateDbInstance-request-dbStorageType

    Returns:
        Details of the created DB instance.
    """
    if not tool_write_mode:
        raise Exception(
            'CreateDbInstance tool invocation not allowed when tool-write-mode is set to False'
        )

    ts_influx_client = get_timestream_influxdb_client()

    # Required parameters
    params = {
        'name': db_instance_name,
        'dbInstanceType': db_instance_type,
        'password': password,
        'vpcSecurityGroupIds': vpc_security_group_ids,
        'vpcSubnetIds': vpc_subnet_ids,
        'allocatedStorage': allocated_storage_gb,
        'publiclyAccessible': publicly_accessible,
    }

    # Add optional parameters if provided
    if db_parameter_group_id:
        params['dbParameterGroupIdentifier'] = db_parameter_group_id
    if username:
        params['username'] = username
    if organization:
        params['organization'] = organization
    if bucket:
        params['bucket'] = bucket
    if port:
        params['port'] = str(port)
    if username:
        params['username'] = username
    if db_storage_type:
        params['db_storage_type'] = db_storage_type
    if deployment_type:
        params['deployment_type'] = deployment_type
    if networkType:
        params['networkType'] = networkType

    if tags:
        tag_list = [{'Key': k, 'Value': v} for k, v in tags.items()]
        params['tags'] = str(tag_list)

    try:
        response = ts_influx_client.create_db_instance(**params)
        return response
    except Exception as e:
        logger.error(f'Error creating DB instance: {str(e)}')
        raise e


@mcp.tool(
    name='LsInstancesOfCluster',
    description='List all Timestream for InfluxDB instances belonging to a specific DB cluster.',
)
async def list_db_instances_for_cluster(
    db_cluster_id: str = REQUIRED_FIELD_DB_CLUSTER_ID,
    next_token: Optional[str] = OPTIONAL_FIELD_NEXT_TOKEN,
    max_results: Optional[int] = OPTIONAL_FIELD_MAX_RESULTS,
) -> Dict[str, Any]:
    """Returns a list of Timestream for InfluxDB DB instances belonging to a specific cluster.

    API reference: https://docs.aws.amazon.com/ts-influxdb/latest/ts-influxdb-api/API_ListDbInstancesForCluster.html

    Returns:
        A list of Timestream for InfluxDB instance summaries belonging to the cluster.
    """
    ts_influx_client = get_timestream_influxdb_client()

    params = {'dbClusterId': db_cluster_id}

    if next_token:
        params['nextToken'] = next_token
    if max_results:
        params['maxResults'] = str(max_results)

    try:
        response = ts_influx_client.list_db_instances_for_cluster(**params)
        return response
    except Exception as e:
        logger.error(f'Error listing DB instances for cluster: {str(e)}')
        raise e


@mcp.tool(name='ListDbInstances', description='List all Timestream for InfluxDB DB instances')
async def list_db_instances(
    next_token: Optional[str] = OPTIONAL_FIELD_NEXT_TOKEN,
    max_results: Optional[int] = OPTIONAL_FIELD_MAX_RESULTS,
) -> Dict[str, Any]:
    """Returns a list of Timestream for InfluxDB DB instances.

    API reference: https://docs.aws.amazon.com/ts-influxdb/latest/ts-influxdb-api/API_ListDbInstances.html

    Returns:
        A list of Timestream for InfluxDB DB instance summaries.
    """
    ts_influx_client = get_timestream_influxdb_client()

    params = {}
    if next_token:
        params['nextToken'] = next_token
    if max_results:
        params['maxResults'] = str(max_results)

    try:
        response = ts_influx_client.list_db_instances(**params)
        return response
    except Exception as e:
        logger.error(f'Error listing DB instances: {str(e)}')
        raise e


@mcp.tool(name='ListDbClusters', description='List all Timestream for InfluxDB DB clusters.')
async def list_db_clusters(
    next_token: Optional[str] = OPTIONAL_FIELD_NEXT_TOKEN,
    max_results: Optional[int] = OPTIONAL_FIELD_MAX_RESULTS,
) -> Dict[str, Any]:
    """Returns a list of Timestream for InfluxDB DB clusters.

    API reference: https://docs.aws.amazon.com/ts-influxdb/latest/ts-influxdb-api/API_ListDbClusters.html

    Returns:
        A list of Timestream for InfluxDB cluster summaries.
    """
    ts_influx_client = get_timestream_influxdb_client()

    params = {}
    if next_token:
        params['nextToken'] = next_token
    if max_results:
        params['maxResults'] = str(max_results)

    try:
        response = ts_influx_client.list_db_clusters(**params)
        return response
    except Exception as e:
        logger.error(f'Error listing DB clusters: {str(e)}')
        raise e


@mcp.tool(
    name='GetDbParameterGroup',
    description='Get a Timestream for InfluxDB DB parameter group details for a db_parameter_group_id',
)
async def get_db_parameter_group(
    identifier: str = REQUIRED_FIELD_PARAMETER_GROUP_ID,
) -> Dict[str, Any]:
    """Returns a Timestream for InfluxDB DB parameter group.

    API reference: https://docs.aws.amazon.com/ts-influxdb/latest/ts-influxdb-api/API_GetDbParameterGroup.html

    Returns:
        Details of the DB parameter group.
    """
    ts_influx_client = get_timestream_influxdb_client()

    try:
        response = ts_influx_client.get_db_parameter_group(identifier=identifier)
        return response
    except Exception as e:
        logger.error(f'Error getting DB parameter group: {str(e)}')
        raise e


@mcp.tool(
    name='GetDbInstance',
    description='Returns a Timestream for InfluxDB DB instance details by the instance-identifier',
)
async def get_db_instance(
    identifier: str = REQUIRED_FIELD_DB_INSTANCE_IDENTIFIER,
) -> Dict[str, Any]:
    """Returns a Timestream for InfluxDB DB instance.

    API reference: https://docs.aws.amazon.com/ts-influxdb/latest/ts-influxdb-api/API_GetDbInstance.html

    Returns:
        Details of the DB instance.
    """
    ts_influx_client = get_timestream_influxdb_client()

    try:
        response = ts_influx_client.get_db_instance(identifier=identifier)
        return response
    except Exception as e:
        logger.error(f'Error getting DB instance: {str(e)}')
        raise e


@mcp.tool(
    name='GetDbCluster',
    description='Returns a Timestream for InfluxDB DB cluster details by the db_cluster_id',
)
async def get_db_cluster(
    db_cluster_id: str = REQUIRED_FIELD_DB_CLUSTER_ID,
) -> Dict[str, Any]:
    """Retrieves information about a Timestream for InfluxDB cluster.

    API reference: https://docs.aws.amazon.com/ts-influxdb/latest/ts-influxdb-api/API_GetDbCluster.html

    Returns:
        Details of the DB cluster.
    """
    ts_influx_client = get_timestream_influxdb_client()

    try:
        response = ts_influx_client.get_db_cluster(dbClusterId=db_cluster_id)
        return response
    except Exception as e:
        logger.error(f'Error getting DB cluster: {str(e)}')
        raise e


@mcp.tool(
    name='DeleteDbInstance',
    description='Deletes a Timestream for InfluxDB DB instance by the instance-identifier',
)
async def delete_db_instance(
    identifier: str = REQUIRED_FIELD_DB_INSTANCE_IDENTIFIER,
    tool_write_mode: bool = OPTIONAL_FIELD_TOOL_WRITE_MODE,
) -> Dict[str, Any]:
    """Deletes a Timestream for InfluxDB DB instance.

    API reference: https://docs.aws.amazon.com/ts-influxdb/latest/ts-influxdb-api/API_DeleteDbInstance.html

    Returns:
        Details of the deleted DB instance.
    """
    if not tool_write_mode:
        raise Exception(
            'DeleteDbInstance tool invocation not allowed when tool-write-mode is set to False'
        )

    ts_influx_client = get_timestream_influxdb_client()

    try:
        response = ts_influx_client.delete_db_instance(identifier=identifier)
        return response
    except Exception as e:
        logger.error(f'Error deleting DB instance: {str(e)}')
        raise e


@mcp.tool(
    name='DeleteDbCluster',
    description='Deletes a Timestream for InfluxDB cluster by the db_cluster_id',
)
async def delete_db_cluster(
    db_cluster_id: str = REQUIRED_FIELD_DB_CLUSTER_ID,
    tool_write_mode: bool = OPTIONAL_FIELD_TOOL_WRITE_MODE,
) -> Dict[str, Any]:
    """Deletes a Timestream for InfluxDB cluster.

    API reference: https://docs.aws.amazon.com/ts-influxdb/latest/ts-influxdb-api/API_DeleteDbCluster.html

    Returns:
        Details of the deleted DB cluster.
    """
    if not tool_write_mode:
        raise Exception(
            'DeleteDbCluster tool invocation not allowed when tool-write-mode is set to False'
        )

    ts_influx_client = get_timestream_influxdb_client()

    try:
        response = ts_influx_client.delete_db_cluster(dbClusterId=db_cluster_id)
        return response
    except Exception as e:
        logger.error(f'Error deleting DB cluster: {str(e)}')
        raise e


@mcp.tool(
    name='ListDbParamGroups', description='List all Timestream for InfluxDB DB parameter groups.'
)
async def list_db_parameter_groups(
    next_token: Optional[str] = OPTIONAL_FIELD_NEXT_TOKEN,
    max_results: Optional[int] = OPTIONAL_FIELD_MAX_RESULTS,
) -> Dict[str, Any]:
    """Returns a list of Timestream for InfluxDB DB parameter groups.

    API reference: https://docs.aws.amazon.com/ts-influxdb/latest/ts-influxdb-api/API_ListDbParameterGroups.html

    Returns:
        A list of Timestream for InfluxDB DB parameter group summaries.
    """
    ts_influx_client = get_timestream_influxdb_client()

    params = {}
    if next_token:
        params['nextToken'] = next_token
    if max_results:
        params['maxResults'] = str(max_results)

    try:
        response = ts_influx_client.list_db_parameter_groups(**params)
        return response
    except Exception as e:
        logger.error(f'Error listing DB parameter groups: {str(e)}')
        raise e


@mcp.tool(name='ListTagsForResource', description='A list of tags applied to the resource.')
async def list_tags_for_resource(
    resource_arn: str = REQUIRED_FIELD_RESOURCE_ARN,
) -> Dict[str, Any]:
    """A list of tags applied to the resource.

    API reference: https://docs.aws.amazon.com/ts-influxdb/latest/ts-influxdb-api/API_ListTagsForResource.html

    Returns:
        A list of tags used to categorize and track resources.
    """
    ts_influx_client = get_timestream_influxdb_client()

    try:
        response = ts_influx_client.list_tags_for_resource(resourceArn=resource_arn)
        return response
    except Exception as e:
        logger.error(f'Error listing tags for resource: {str(e)}')
        raise e


@mcp.tool(
    name='TagResource',
    description='Tags are composed of a Key/Value pairs. Apply them to Timestream for InfluxDB resource.',
)
async def tag_resource(
    resource_arn: str = REQUIRED_FIELD_RESOURCE_ARN,
    tags: Dict[str, str] = REQUIRED_FIELD_TAGS_RESOURCE,
    tool_write_mode: bool = OPTIONAL_FIELD_TOOL_WRITE_MODE,
) -> Dict[str, Any]:
    """Tags are composed of a Key/Value pairs. You can use tags to categorize and track your Timestream for InfluxDB resources.

    API reference: https://docs.aws.amazon.com/ts-influxdb/latest/ts-influxdb-api/API_TagResource.html

    Returns:
        Status of the tag operation.
    """
    if not tool_write_mode:
        raise Exception(
            'TagResource tool invocation not allowed when tool-write-mode is set to False'
        )

    ts_influx_client = get_timestream_influxdb_client()

    # Convert tags dictionary to list of Key/Value pairs
    tag_list = [{'Key': k, 'Value': v} for k, v in tags.items()]

    try:
        response = ts_influx_client.tag_resource(resourceArn=resource_arn, tags=tag_list)
        return response
    except Exception as e:
        logger.error(f'Error tagging resource: {str(e)}')
        raise e


@mcp.tool(
    name='UntagResource',
    description='Removes the tags, identified by the keys, from the specified resource.',
)
async def untag_resource(
    resource_arn: str = REQUIRED_FIELD_RESOURCE_ARN,
    tag_keys: List[str] = REQUIRED_FIELD_TAG_KEYS,
    tool_write_mode: bool = OPTIONAL_FIELD_TOOL_WRITE_MODE,
) -> Dict[str, Any]:
    """Removes the tag from the specified resource.

    API reference: https://docs.aws.amazon.com/ts-influxdb/latest/ts-influxdb-api/API_UntagResource.html

    Returns:
        Status of the untag operation.
    """
    if not tool_write_mode:
        raise Exception(
            'UntagResource tool invocation not allowed when tool-write-mode is set to False'
        )

    ts_influx_client = get_timestream_influxdb_client()

    try:
        response = ts_influx_client.untag_resource(resourceArn=resource_arn, tagKeys=tag_keys)
        return response
    except Exception as e:
        logger.error(f'Error untagging resource: {str(e)}')
        raise e


@mcp.tool(name='UpdateDbCluster', description='Updates a Timestream for InfluxDB cluster.')
async def update_db_cluster(
    db_cluster_id: str = REQUIRED_FIELD_DB_CLUSTER_ID,
    db_instance_type: Optional[str] = OPTIONAL_FIELD_DB_INSTANCE_TYPE_CLUSTER_UPDATE,
    db_parameter_group_identifier: Optional[
        str
    ] = OPTIONAL_FIELD_DB_PARAMETER_GROUP_IDENTIFIER_UPDATE,
    port: Optional[int] = OPTIONAL_FIELD_PORT_UPDATE,
    failover_mode: Optional[str] = OPTIONAL_FIELD_FAILOVER_MODE_UPDATE,
    log_delivery_configuration: Optional[
        Dict[str, Any]
    ] = OPTIONAL_FIELD_LOG_DELIVERY_CONFIGURATION_UPDATE,
    tool_write_mode: bool = OPTIONAL_FIELD_TOOL_WRITE_MODE,
) -> Dict[str, Any]:
    """Updates a Timestream for InfluxDB cluster.

    API reference: https://docs.aws.amazon.com/ts-influxdb/latest/ts-influxdb-api/API_UpdateDbCluster.html

    Returns:
        Details of the updated DB cluster.
    """
    if not tool_write_mode:
        raise Exception(
            'UpdateDbCluster tool invocation not allowed when tool-write-mode is set to False'
        )

    ts_influx_client = get_timestream_influxdb_client()

    # Required parameters
    params = {'dbClusterId': db_cluster_id}

    # Add optional parameters if provided
    if db_instance_type:
        params['dbInstanceType'] = db_instance_type
    if db_parameter_group_identifier:
        params['dbParameterGroupIdentifier'] = db_parameter_group_identifier
    if port:
        params['port'] = str(port)
    if failover_mode:
        params['failoverMode'] = failover_mode
    if log_delivery_configuration:
        params['logDeliveryConfiguration'] = str(log_delivery_configuration)

    try:
        response = ts_influx_client.update_db_cluster(**params)
        return response
    except Exception as e:
        logger.error(f'Error updating DB cluster: {str(e)}')
        raise e


@mcp.tool(name='UpdateDbInstance', description='Updates a Timestream for InfluxDB DB instance.')
async def update_db_instance(
    identifier: str = REQUIRED_FIELD_DB_INSTANCE_IDENTIFIER,
    db_instance_type: Optional[str] = OPTIONAL_FIELD_DB_INSTANCE_TYPE_CLUSTER_UPDATE,
    db_parameter_group_identifier: Optional[str] = OPTIONAL_FIELD_DB_PARAMETER_GROUP_ID,
    port: Optional[int] = OPTIONAL_FIELD_PORT,
    allocated_storage_gb: Optional[int] = OPTIONAL_FIELD_ALLOCATED_STORAGE_GB_OPTIONAL,
    db_storage_type: Optional[str] = OPTIONAL_FIELD_DB_STORAGE_TYPE,
    deployment_type: Optional[str] = OPTIONAL_FIELD_DEPLOYMENT_TYPE_INSTANCE,
    log_delivery_configuration: Optional[
        Dict[str, Any]
    ] = OPTIONAL_FIELD_LOG_DELIVERY_CONFIGURATION,
    tool_write_mode: bool = OPTIONAL_FIELD_TOOL_WRITE_MODE,
) -> Dict[str, Any]:
    """Updates a Timestream for InfluxDB DB instance.

    API reference: https://docs.aws.amazon.com/ts-influxdb/latest/ts-influxdb-api/API_UpdateDbInstance.html

    Returns:
        Details of the updated DB instance.
    """
    if not tool_write_mode:
        raise Exception(
            'UpdateDbInstance tool invocation not allowed when tool-write-mode is set to False'
        )

    ts_influx_client = get_timestream_influxdb_client()

    # Required parameters
    params = {'identifier': identifier}

    # Add optional parameters if provided
    if db_instance_type:
        params['dbInstanceType'] = db_instance_type
    if db_parameter_group_identifier:
        params['dbParameterGroupIdentifier'] = db_parameter_group_identifier
    if port:
        params['port'] = str(port)
    if allocated_storage_gb:
        params['allocatedStorage'] = str(allocated_storage_gb)
    if db_storage_type:
        params['dbStorageType'] = db_storage_type
    if deployment_type:
        params['deploymentType'] = deployment_type
    if log_delivery_configuration:
        params['logDeliveryConfiguration'] = str(log_delivery_configuration)

    try:
        response = ts_influx_client.update_db_instance(**params)
        return response
    except Exception as e:
        logger.error(f'Error updating DB instance: {str(e)}')
        raise e


@mcp.tool(
    name='LsInstancesByStatus',
    description='Returns a list of Timestream for InfluxDB DB instances filtered by status (case-insensitive).',
)
async def list_db_instances_by_status(
    status: str = REQUIRED_FIELD_STATUS,
) -> Dict[str, Any]:
    """Returns a list of Timestream for InfluxDB DB instances filtered by status (case-insensitive).

    This tool paginates through all DB instances and filters them by the provided status
    in a case-insensitive manner.

    Returns:
        A list of Timestream for InfluxDB DB instance summaries matching the specified status.
    """
    ts_influx_client = get_timestream_influxdb_client()

    # Convert status to lowercase for case-insensitive comparison
    status_lower = status.lower()

    # Initialize variables for pagination
    next_token = None
    filtered_instances = []

    try:
        # Paginate through all instances
        while True:
            # Prepare parameters for the API call
            params = {}
            if next_token:
                params['nextToken'] = next_token

            # Call the ListDbInstances API
            response = ts_influx_client.list_db_instances(**params)

            # Filter instances by status (case-insensitive)
            if 'items' in response:
                for instance in response['items']:
                    if (
                        'status' in instance
                        and instance['status'] is not None
                        and instance['status'].lower() == status_lower
                    ):
                        filtered_instances.append(instance)

            # Check if there are more results to fetch
            if 'nextToken' in response and response['nextToken']:
                next_token = response['nextToken']
            else:
                # No more results to fetch
                break

        # Prepare the response
        result = {'items': filtered_instances, 'count': len(filtered_instances)}

        return result
    except Exception as e:
        logger.error(f'Error listing DB instances by status: {str(e)}')
        raise e


@mcp.tool(
    name='ListClustersByStatus',
    description='Returns a list of Timestream for InfluxDB DB clusters filtered by status (case-insensitive).',
)
async def list_db_clusters_by_status(
    status: str = REQUIRED_FIELD_STATUS_CLUSTER,
) -> Dict[str, Any]:
    """Returns a list of Timestream for InfluxDB DB clusters filtered by status (case-insensitive).

    This tool paginates through all DB clusters and filters them by the provided status
    in a case-insensitive manner.

    Returns:
        A list of Timestream for InfluxDB DB cluster summaries matching the specified status.
    """
    ts_influx_client = get_timestream_influxdb_client()

    # Convert status to lowercase for case-insensitive comparison
    status_lower = status.lower()

    # Initialize variables for pagination
    next_token = None
    filtered_clusters = []

    try:
        # Paginate through all clusters
        while True:
            # Prepare parameters for the API call
            params = {}
            if next_token:
                params['nextToken'] = next_token

            # Call the ListDbClusters API
            response = ts_influx_client.list_db_clusters(**params)

            # Filter clusters by status (case-insensitive)
            if 'items' in response:
                for cluster in response['items']:
                    if (
                        'status' in cluster
                        and cluster['status'] is not None
                        and cluster['status'].lower() == status_lower
                    ):
                        filtered_clusters.append(cluster)

            # Check if there are more results to fetch
            if 'nextToken' in response and response['nextToken']:
                next_token = response['nextToken']
            else:
                # No more results to fetch
                break

        # Prepare the response
        result = {'items': filtered_clusters, 'count': len(filtered_clusters)}

        return result
    except Exception as e:
        logger.error(f'Error listing DB clusters by status: {str(e)}')
        raise e


@mcp.tool(
    name='CreateDbParamGroup',
    description='Creates a new Timestream for InfluxDB DB parameter group to associate with DB instances.',
)
async def create_db_parameter_group(
    name: str = REQUIRED_FIELD_PARAM_GROUP_NAME,
    tool_write_mode: bool = OPTIONAL_FIELD_TOOL_WRITE_MODE,
    description: Optional[str] = OPTIONAL_FIELD_PARAM_GROUP_DESCRIPTION,
    parameters: Optional[Dict[str, Any]] = OPTIONAL_FIELD_PARAMETERS,
    tags: Optional[Dict[str, str]] = OPTIONAL_FIELD_TAGS,
) -> Dict[str, Any]:
    """Creates a new Timestream for InfluxDB DB parameter group to associate with DB instances.

    API reference: https://docs.aws.amazon.com/ts-influxdb/latest/ts-influxdb-api/API_CreateDbParameterGroup.html

    Returns:
        Details of the created DB parameter group.
    """
    if not tool_write_mode:
        raise Exception(
            'CreateDbParamGroup tool invocation not allowed when tool-write-mode is set to False'
        )

    ts_influx_client = get_timestream_influxdb_client()

    # Required parameters
    params = {'name': name}

    # Add optional parameters if provided
    if description:
        params['description'] = description
    if parameters:
        params['parameters'] = str(parameters)
    if tags:
        tag_list = [{'Key': k, 'Value': v} for k, v in tags.items()]
        params['tags'] = str(tag_list)

    try:
        response = ts_influx_client.create_db_parameter_group(**params)
        return response
    except Exception as e:
        logger.error(f'Error creating DB parameter group: {str(e)}')
        raise e


@mcp.tool(name='InfluxDBWritePoints', description='Write data points to InfluxDB endpoint.')
async def influxdb_write_points(
    url: Optional[str] = OPTIONAL_FIELD_URL,
    token: Optional[str] = OPTIONAL_FIELD_TOKEN,
    bucket: str = REQUIRED_FIELD_BUCKET_INFLUX,
    org: Optional[str] = OPTIONAL_FIELD_ORG,
    points: List[Dict[str, Any]] = REQUIRED_FIELD_POINTS,
    time_precision: str = OPTIONAL_FIELD_WRITE_PRECISION,
    sync_mode: Optional[str] = OPTIONAL_FIELD_SYNC_MODE,
    verify_ssl: bool = OPTIONAL_FIELD_VERIFY_SSL,
    tool_write_mode: bool = OPTIONAL_FIELD_TOOL_WRITE_MODE,
) -> Dict[str, Any]:
    """Write data points to InfluxDB.

        Example of points:
        [
           {
             "measurement": "my_measurement",
             "tags": {"location": "Prague"},
             "fields": {"temperature": 25.3}
             "time": "2025-06-06T19:00:00Z"
            }
        ]

    Returns:
        Status of the write operation.
    """
    if not tool_write_mode:
        raise Exception(
            'InfluxDBWritePoints tool invocation not allowed when tool-write-mode is set to False'
        )

    resolved_url, resolved_token, resolved_org = resolve_influxdb_config(url, token, org)

    try:
        client = get_influxdb_client(
            url=resolved_url, token=resolved_token, org=resolved_org, verify_ssl=verify_ssl
        )
        if sync_mode and sync_mode.lower() == 'synchronous':
            write_api = client.write_api(write_options=SYNCHRONOUS)
        else:
            write_api = client.write_api(write_options=ASYNCHRONOUS)

        # Convert dictionary points to Point objects
        influx_points = []
        for p in points:
            point = Point(p['measurement'])

            # Add tags
            if 'tags' in p:
                for tag_key, tag_value in p['tags'].items():
                    point = point.tag(tag_key, tag_value)

            # Add fields
            if 'fields' in p:
                for field_key, field_value in p['fields'].items():
                    point = point.field(field_key, field_value)

            # Add time if provided
            if 'time' in p:
                point = point.time(p['time'])

            influx_points.append(point)

        # Write points
        write_api.write(
            bucket=bucket,
            org=resolved_org,
            record=influx_points,
            write_precision=getattr(WritePrecision, time_precision.upper()),
            verify_ssl=verify_ssl,
        )

        # Close client
        client.close()

        return {
            'status': 'success',
            'message': f'Successfully wrote {len(points)} points to InfluxDB',
        }
    except Exception as e:
        logger.error(f'Error writing points to InfluxDB: {str(e)}')
        return {'status': 'error', 'message': str(e)}


@mcp.tool(name='InfluxDBWriteLP', description='Write data in Line Protocol format to InfluxDB.')
async def influxdb_write_line_protocol(
    url: Optional[str] = OPTIONAL_FIELD_URL,
    token: Optional[str] = OPTIONAL_FIELD_TOKEN,
    bucket: str = REQUIRED_FIELD_BUCKET_INFLUX,
    org: Optional[str] = OPTIONAL_FIELD_ORG,
    data_line_protocol: str = REQUIRED_FIELD_DATA_LINE_PROTOCOL,
    time_precision: str = OPTIONAL_FIELD_WRITE_PRECISION,
    sync_mode: str = OPTIONAL_FIELD_SYNC_MODE,
    verify_ssl: bool = OPTIONAL_FIELD_VERIFY_SSL,
    tool_write_mode: bool = OPTIONAL_FIELD_TOOL_WRITE_MODE,
) -> Dict[str, Any]:
    """Write data in Line Protocol format to InfluxDB.

    Returns:
        Status of the write operation.
    """
    if not tool_write_mode:
        raise Exception(
            'InfluxDBWriteLineProtocol tool invocation not allowed when tool-write-mode is set to False'
        )

    resolved_url, resolved_token, resolved_org = resolve_influxdb_config(url, token, org)

    try:
        client = get_influxdb_client(
            url=resolved_url, token=resolved_token, org=resolved_org, verify_ssl=verify_ssl
        )

        # Set write mode
        if sync_mode and sync_mode.lower() == 'synchronous':
            write_api = client.write_api(write_options=SYNCHRONOUS)
        else:
            write_api = client.write_api(write_options=ASYNCHRONOUS)

        # Write line protocol
        write_api.write(
            bucket=bucket,
            org=resolved_org,
            record=data_line_protocol,
            write_precision=getattr(WritePrecision, time_precision.upper()),
            verify_ssl=verify_ssl,
        )

        # Close client
        client.close()

        return {
            'status': 'success',
            'message': 'Successfully wrote line protocol data to InfluxDB',
        }
    except Exception as e:
        logger.error(f'Error writing line protocol to InfluxDB: {str(e)}')
        return {'status': 'error', 'message': str(e)}


@mcp.tool(name='InfluxDBQuery', description='Query data from InfluxDB using Flux query language.')
async def influxdb_query(
    url: Optional[str] = OPTIONAL_FIELD_URL,
    token: Optional[str] = OPTIONAL_FIELD_TOKEN,
    org: Optional[str] = OPTIONAL_FIELD_ORG,
    query: str = REQUIRED_FIELD_QUERY,
    verify_ssl: bool = OPTIONAL_FIELD_VERIFY_SSL,
) -> Dict[str, Any]:
    """Query data from InfluxDB using Flux query language.

    Returns:
        Query results in the specified format.
    """
    resolved_url, resolved_token, resolved_org = resolve_influxdb_config(url, token, org)

    try:
        client = get_influxdb_client(
            url=resolved_url, token=resolved_token, org=resolved_org, verify_ssl=verify_ssl
        )
        query_api = client.query_api()

        # Return as JSON
        tables = query_api.query(org=resolved_org, query=query)

        # Process the tables into a more usable format
        # System keys that are not tags
        system_keys = {
            'result',
            'table',
            '_start',
            '_stop',
            '_time',
            '_value',
            '_field',
            '_measurement',
        }

        result = []
        for table in tables:
            for record in table.records:
                # Extract tags by filtering out system keys
                tags = {
                    k: v
                    for k, v in record.values.items()
                    if k not in system_keys and v is not None
                }

                result.append(
                    {
                        'measurement': record.get_measurement(),
                        'field': record.get_field(),
                        'value': record.get_value(),
                        'time': record.get_time().isoformat() if record.get_time() else None,
                        'tags': tags,
                    }
                )

        client.close()
        return {'status': 'success', 'result': result, 'format': 'json'}

    except Exception as e:
        logger.error(f'Error querying InfluxDB: {str(e)}')
        return {'status': 'error', 'message': str(e)}


@mcp.tool(name='InfluxDBListBuckets', description='List all buckets in InfluxDB.')
async def influxdb_list_buckets(
    url: Optional[str] = OPTIONAL_FIELD_URL,
    token: Optional[str] = OPTIONAL_FIELD_TOKEN,
    org: Optional[str] = OPTIONAL_FIELD_ORG,
    verify_ssl: bool = OPTIONAL_FIELD_VERIFY_SSL,
) -> Dict[str, Any]:
    """List all buckets in InfluxDB.

    Returns:
        List of buckets with their details.
    """
    resolved_url, resolved_token, resolved_org = resolve_influxdb_config(url, token, org)

    try:
        client = get_influxdb_client(
            url=resolved_url, token=resolved_token, org=resolved_org, verify_ssl=verify_ssl
        )
        buckets_api = client.buckets_api()
        buckets = buckets_api.find_buckets().buckets

        result = []
        for bucket in buckets:
            result.append(
                {
                    'id': bucket.id,
                    'name': bucket.name,
                    'org_id': bucket.org_id,
                    'retention_period': bucket.retention_rules[0].every_seconds
                    if bucket.retention_rules
                    else None,
                    'created_at': bucket.created_at.isoformat() if bucket.created_at else None,
                    'updated_at': bucket.updated_at.isoformat() if bucket.updated_at else None,
                }
            )

        client.close()
        return {'status': 'success', 'buckets': result}

    except Exception as e:
        logger.error(f'Error listing InfluxDB buckets: {str(e)}')
        return {'status': 'error', 'message': str(e)}


@mcp.tool(name='InfluxDBCreateBucket', description='Create a new bucket in InfluxDB.')
async def influxdb_create_bucket(
    bucket_name: str = Field(..., description='The name of the bucket to create.'),
    url: Optional[str] = OPTIONAL_FIELD_URL,
    token: Optional[str] = OPTIONAL_FIELD_TOKEN,
    org: Optional[str] = OPTIONAL_FIELD_ORG,
    retention_seconds: Optional[int] = Field(
        None, description='Retention period in seconds. 0 or None means infinite retention.'
    ),
    description: Optional[str] = Field(None, description='Description of the bucket.'),
    verify_ssl: bool = OPTIONAL_FIELD_VERIFY_SSL,
    tool_write_mode: bool = OPTIONAL_FIELD_TOOL_WRITE_MODE,
) -> Dict[str, Any]:
    """Create a new bucket in InfluxDB.

    Returns:
        Details of the created bucket.
    """
    if not tool_write_mode:
        raise Exception(
            'InfluxDBCreateBucket tool invocation not allowed when tool-write-mode is set to False'
        )

    resolved_url, resolved_token, resolved_org = resolve_influxdb_config(url, token, org)

    try:
        client = get_influxdb_client(
            url=resolved_url, token=resolved_token, org=resolved_org, verify_ssl=verify_ssl
        )
        buckets_api = client.buckets_api()

        # Get org ID from org name
        orgs_api = client.organizations_api()
        orgs = orgs_api.find_organizations(org=resolved_org)
        if not orgs:
            raise ValueError(f'Organization "{resolved_org}" not found')
        org_id = orgs[0].id

        retention_rules = []
        if retention_seconds and retention_seconds > 0:
            retention_rules.append(
                BucketRetentionRules(type='expire', every_seconds=retention_seconds)
            )

        bucket = buckets_api.create_bucket(
            bucket_name=bucket_name,
            org_id=org_id,
            retention_rules=retention_rules,
            description=description,
        )

        client.close()
        return {
            'status': 'success',
            'bucket': {
                'id': bucket.id,
                'name': bucket.name,
                'org_id': bucket.org_id,
                'retention_period': bucket.retention_rules[0].every_seconds
                if bucket.retention_rules
                else None,
                'created_at': bucket.created_at.isoformat() if bucket.created_at else None,
            },
        }

    except Exception as e:
        logger.error(f'Error creating InfluxDB bucket: {str(e)}')
        return {'status': 'error', 'message': str(e)}


@mcp.tool(name='InfluxDBListOrgs', description='List all organizations in InfluxDB.')
async def influxdb_list_orgs(
    url: Optional[str] = OPTIONAL_FIELD_URL,
    token: Optional[str] = OPTIONAL_FIELD_TOKEN,
    verify_ssl: bool = OPTIONAL_FIELD_VERIFY_SSL,
) -> Dict[str, Any]:
    """List all organizations in InfluxDB.

    Returns:
        List of organizations with their details.
    """
    resolved_url, resolved_token, _ = resolve_influxdb_config(
        url, token, org=None, require_org=False
    )

    try:
        client = get_influxdb_client(url=resolved_url, token=resolved_token, verify_ssl=verify_ssl)
        orgs_api = client.organizations_api()
        orgs = orgs_api.find_organizations()

        result = []
        for org in orgs:
            result.append(
                {
                    'id': org.id,
                    'name': org.name,
                    'description': org.description,
                    'created_at': org.created_at.isoformat() if org.created_at else None,
                    'updated_at': org.updated_at.isoformat() if org.updated_at else None,
                }
            )

        client.close()
        return {'status': 'success', 'organizations': result}

    except Exception as e:
        logger.error(f'Error listing InfluxDB organizations: {str(e)}')
        return {'status': 'error', 'message': str(e)}


@mcp.tool(name='InfluxDBCreateOrg', description='Create a new organization in InfluxDB.')
async def influxdb_create_org(
    org_name: str = Field(..., description='The name of the organization to create.'),
    url: Optional[str] = OPTIONAL_FIELD_URL,
    token: Optional[str] = OPTIONAL_FIELD_TOKEN,
    verify_ssl: bool = OPTIONAL_FIELD_VERIFY_SSL,
    tool_write_mode: bool = OPTIONAL_FIELD_TOOL_WRITE_MODE,
) -> Dict[str, Any]:
    """Create a new organization in InfluxDB.

    Returns:
        Details of the created organization.
    """
    if not tool_write_mode:
        raise Exception(
            'InfluxDBCreateOrg tool invocation not allowed when tool-write-mode is set to False'
        )

    resolved_url, resolved_token, _ = resolve_influxdb_config(
        url, token, org=None, require_org=False
    )

    try:
        client = get_influxdb_client(url=resolved_url, token=resolved_token, verify_ssl=verify_ssl)
        orgs_api = client.organizations_api()

        org = orgs_api.create_organization(name=org_name)

        client.close()
        return {
            'status': 'success',
            'organization': {
                'id': org.id,
                'name': org.name,
                'created_at': org.created_at.isoformat() if org.created_at else None,
            },
        }

    except Exception as e:
        logger.error(f'Error creating InfluxDB organization: {str(e)}')
        return {'status': 'error', 'message': str(e)}


def main():
    """Main entry point for the MCP server application."""
    logger.info('Starting Timestream for InfluxDB MCP Server')
    mcp.run()


if __name__ == '__main__':
    main()
