import re

# Read the file
with open('src/vultr_mcp/server.py', 'r') as f:
    lines = f.readlines()

# Find the line with 'return tools'
insert_line = -1
for i, line in enumerate(lines):
    if 'return tools' in line:
        insert_line = i
        break

if insert_line == -1:
    print("Error: Could not find 'return tools'")
    exit(1)

# Find the list_os tool (should be a few lines before return tools)
list_os_line = -1
for i in range(insert_line-1, max(0, insert_line-10), -1):
    if 'list_os' in lines[i]:
        list_os_line = i
        break

if list_os_line == -1:
    print("Error: Could not find 'list_os' tool")
    exit(1)

print(f"Found list_os at line {list_os_line+1}")
print(f"Found return tools at line {insert_line+1}")

# Find the end of the list_os tool definition
# Look for the closing parenthesis and comma
end_of_list_os = -1
for i in range(list_os_line, min(len(lines), list_os_line + 10)):
    if '),' in lines[i] or ')' in lines[i] and i < insert_line:
        end_of_list_os = i
        break

if end_of_list_os == -1:
    print("Error: Could not find end of list_os tool")
    exit(1)

print(f"End of list_os tool at line {end_of_list_os+1}")

# DNS tools to insert
dns_tools = '''        ),
        # DNS Management Tools
        Tool(
            name="list_dns_domains",
            description="List all DNS domains",
            inputSchema={"type": "object", "properties": {}}
        ),
        Tool(
            name="create_dns_domain",
            description="Create a new DNS domain",
            inputSchema={
                "type": "object",
                "properties": {
                    "domain": {
                        "type": "string",
                        "description": "Domain name (e.g., 'example.com')"
                    },
                    "ip": {
                        "type": "string",
                        "description": "IP address for the domain"
                    }
                },
                "required": ["domain", "ip"]
            }
        ),
        Tool(
            name="delete_dns_domain",
            description="Delete a DNS domain",
            inputSchema={
                "type": "object",
                "properties": {
                    "domain": {
                        "type": "string",
                        "description": "Domain name"
                    }
                },
                "required": ["domain"]
            }
        ),
        Tool(
            name="get_dns_domain",
            description="Get DNS domain details",
            inputSchema={
                "type": "object",
                "properties": {
                    "domain": {
                        "type": "string",
                        "description": "Domain name"
                    }
                },
                "required": ["domain"]
            }
        ),
        Tool(
            name="list_dns_records",
            description="List DNS records for a domain",
            inputSchema={
                "type": "object",
                "properties": {
                    "domain": {
                        "type": "string",
                        "description": "Domain name"
                    }
                },
                "required": ["domain"]
            }
        ),
        Tool(
            name="create_dns_record",
            description="Create a new DNS record",
            inputSchema={
                "type": "object",
                "properties": {
                    "domain": {
                        "type": "string",
                        "description": "Domain name"
                    },
                    "type": {
                        "type": "string",
                        "description": "Record type (A, AAAA, CNAME, MX, TXT, etc.)"
                    },
                    "name": {
                        "type": "string",
                        "description": "Record name (e.g., 'www', '@', 'mail')"
                    },
                    "data": {
                        "type": "string",
                        "description": "Record data (IP address, hostname, etc.)"
                    },
                    "ttl": {
                        "type": "integer",
                        "description": "Time to live in seconds",
                        "default": 300
                    },
                    "priority": {
                        "type": "integer",
                        "description": "Priority for MX records",
                        "default": 0
                    }
                },
                "required": ["domain", "type", "name", "data"]
            }
        ),
        Tool(
            name="update_dns_record",
            description="Update a DNS record",
            inputSchema={
                "type": "object",
                "properties": {
                    "domain": {
                        "type": "string",
                        "description": "Domain name"
                    },
                    "record_id": {
                        "type": "string",
                        "description": "Record ID"
                    },
                    "type": {
                        "type": "string",
                        "description": "Record type (A, AAAA, CNAME, MX, TXT, etc.)"
                    },
                    "name": {
                        "type": "string",
                        "description": "Record name"
                    },
                    "data": {
                        "type": "string",
                        "description": "Record data"
                    },
                    "ttl": {
                        "type": "integer",
                        "description": "Time to live in seconds"
                    },
                    "priority": {
                        "type": "integer",
                        "description": "Priority for MX records"
                    }
                },
                "required": ["domain", "record_id"]
            }
        ),
        Tool(
            name="delete_dns_record",
            description="Delete a DNS record",
            inputSchema={
                "type": "object",
                "properties": {
                    "domain": {
                        "type": "string",
                        "description": "Domain name"
                    },
                    "record_id": {
                        "type": "string",
                        "description": "Record ID"
                    }
                },
                "required": ["domain", "record_id"]
            }
        )'''

# Insert the DNS tools after the list_os tool
new_lines = lines[:end_of_list_os+1] + [dns_tools + '\n'] + lines[end_of_list_os+1:]

# Write the updated file
with open('src/vultr_mcp/server.py', 'w') as f:
    f.writelines(new_lines)

print("DNS tools have been successfully inserted into the tools list!")
print(f"Total lines in file: {len(new_lines)}")