import re

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

# Find the tools list section
# Look for the list_os tool definition and insert DNS tools before the return
list_os_tool_pattern = r"(Tool\(\s*name=\"list_os\",\s*description=\"List available Vultr operating systems\",\s*inputSchema=\{"type": "object", "properties": \{\}\s*\}\s*\)\s*,?\s*)"
match = re.search(list_os_tool_pattern, content, re.DOTALL)

if match:
    list_os_tool = match.group(1)
    
    # DNS tools definitions (from lines 512-651)
    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 DNS tools after list_os tool
    new_content = content.replace(list_os_tool, list_os_tool + dns_tools)
    
    # Remove the duplicate DNS tools definitions (lines 512-651)
    # Find the start of duplicate DNS tools
    duplicate_start = new_content.find('# DNS Management Tools')
    if duplicate_start != -1:
        # Find the end of the duplicate section (look for next Tool definition or end of file)
        duplicate_end = duplicate_start
        lines = new_content[duplicate_start:].split('\n')
        in_tool = False
        for i, line in enumerate(lines):
            if 'Tool(' in line:
                in_tool = True
            if in_tool and line.strip().endswith(')'):
                in_tool = False
            if i > 100 and not in_tool and line.strip() and not line.strip().startswith('Tool('):
                duplicate_end = duplicate_start + sum(len(l) + 1 for l in lines[:i])
                break
        
        if duplicate_end > duplicate_start:
            new_content = new_content[:duplicate_start] + new_content[duplicate_end:]
    
    # Write the fixed content
    with open('src/vultr_mcp/server.py', 'w') as f:
        f.write(new_content)
    
    print("DNS tools have been moved to the correct location in list_tools function")
else:
    print("Could not find list_os tool definition")