import re

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

# Add DNS tool handling to call_tool function
# Find the list_os handling and insert DNS handling after it
list_os_pattern = r"(elif request\.name == \"list_os\":\s*result = vultr_client\.list_os\(\)\s*)"
match = re.search(list_os_pattern, content)
if match:
    list_os_section = match.group(1)
    dns_handling = '''        elif request.name == "list_dns_domains":
            result = vultr_client.list_dns_domains()
        
        elif request.name == "create_dns_domain":
            domain = request.arguments.get("domain")
            ip = request.arguments.get("ip")
            if not domain or not ip:
                raise ValueError("domain and ip are required")
            result = vultr_client.create_dns_domain(domain, ip)
        
        elif request.name == "delete_dns_domain":
            domain = request.arguments.get("domain")
            if not domain:
                raise ValueError("domain is required")
            result = vultr_client.delete_dns_domain(domain)
        
        elif request.name == "get_dns_domain":
            domain = request.arguments.get("domain")
            if not domain:
                raise ValueError("domain is required")
            result = vultr_client.get_dns_domain(domain)
        
        elif request.name == "list_dns_records":
            domain = request.arguments.get("domain")
            if not domain:
                raise ValueError("domain is required")
            result = vultr_client.list_dns_records(domain)
        
        elif request.name == "create_dns_record":
            domain = request.arguments.get("domain")
            record_type = request.arguments.get("type")
            name = request.arguments.get("name")
            data = request.arguments.get("data")
            ttl = request.arguments.get("ttl", 300)
            priority = request.arguments.get("priority", 0)
            if not domain or not record_type or not name or not data:
                raise ValueError("domain, type, name, and data are required")
            result = vultr_client.create_dns_record(domain, record_type, name, data, ttl, priority)
        
        elif request.name == "update_dns_record":
            domain = request.arguments.get("domain")
            record_id = request.arguments.get("record_id")
            if not domain or not record_id:
                raise ValueError("domain and record_id are required")
            # Filter out None values
            update_args = {k: v for k, v in request.arguments.items() 
                          if k not in ["domain", "record_id"] and v is not None}
            result = vultr_client.update_dns_record(domain, record_id, **update_args)
        
        elif request.name == "delete_dns_record":
            domain = request.arguments.get("domain")
            record_id = request.arguments.get("record_id")
            if not domain or not record_id:
                raise ValueError("domain and record_id are required")
            result = vultr_client.delete_dns_record(domain, record_id)
        
'''
    # Insert DNS handling after list_os
    new_content = content.replace(list_os_section, list_os_section + dns_handling)
    content = new_content

# Add DNS resources to list_resources function
resources_pattern = r"(Resource\(\s*uri=\"vultr://os/list\",\s*name=\"Vultr Operating Systems\",\s*description=\"List of available Vultr operating systems\"\s*\)\s*,?\s*)"
match = re.search(resources_pattern, content, re.DOTALL)
if match:
    os_resource = match.group(1)
    dns_resources = '''        Resource(
            uri="vultr://dns/domains",
            name="Vultr DNS Domains",
            description="List of DNS domains"
        ),
        Resource(
            uri="vultr://dns/records",
            name="Vultr DNS Records",
            description="DNS records for a domain"
        ),
'''
    new_content = content.replace(os_resource, os_resource + dns_resources)
    content = new_content

# Add DNS prompts to list_prompts function
prompts_pattern = r"(Prompt\(\s*name=\"monitor_resources\",\s*description=\"Monitor Vultr resources\",\s*arguments=\[\s*\{.*?\}\s*\]\s*\)\s*,?\s*)"
match = re.search(prompts_pattern, content, re.DOTALL)
if match:
    monitor_prompt = match.group(1)
    dns_prompts = '''        Prompt(
            name="manage_dns",
            description="Manage DNS domains and records",
            arguments=[
                {"name": "action", "description": "Action to perform (create_domain, delete_domain, list_records, create_record, update_record, delete_record)", "required": True},
                {"name": "domain", "description": "Domain name", "required": True},
                {"name": "record_id", "description": "Record ID (for update/delete)", "required": False},
                {"name": "record_type", "description": "Record type (A, AAAA, CNAME, etc.)", "required": False},
                {"name": "name", "description": "Record name", "required": False},
                {"name": "data", "description": "Record data", "required": False},
                {"name": "ip", "description": "IP address (for create domain)", "required": False},
            ]
        ),
'''
    new_content = content.replace(monitor_prompt, monitor_prompt + dns_prompts)
    content = new_content

# Add DNS resource handling to get_resource function
get_resource_pattern = r"(elif resource_uri == \"vultr://os/list\":\s*result = vultr_client\.list_os\(\)\s*)"
match = re.search(get_resource_pattern, content)
if match:
    os_resource_handling = match.group(1)
    dns_resource_handling = '''        elif resource_uri == "vultr://dns/domains":
            result = vultr_client.list_dns_domains()
        elif resource_uri == "vultr://dns/records":
            # This requires a domain parameter, so we'll return instructions
            result = {"message": "Please use the list_dns_records tool with domain parameter to get DNS records"}
'''
    new_content = content.replace(os_resource_handling, os_resource_handling + dns_resource_handling)
    content = new_content

# Write the updated content back
with open('src/vultr_mcp/server.py', 'w') as f:
    f.write(content)

print("Server.py updated successfully with DNS functionality")