import os
import re

server_path = r'e:\Git\antigravity-expleration\outlook-desktop-mcp\src\outlook_desktop_mcp\server.py'
formatting_path = r'e:\Git\antigravity-expleration\outlook-desktop-mcp\src\outlook_desktop_mcp\utils\formatting.py'

# 1. Update formatting.py
with open(formatting_path, 'r', encoding='utf-8') as f:
    text = f.read()

# Make format_* functions accept encoding="utf-8"
text = text.replace('def format_email_summary(item) -> dict:', 'def format_email_summary(item, encoding: str = "utf-8") -> dict:')
text = text.replace('def format_email_full(item, body_max_length: int = 5000) -> dict:', 'def format_email_full(item, body_max_length: int = 5000, encoding: str = "utf-8") -> dict:')
text = text.replace('result = format_email_summary(item)', 'result = format_email_summary(item, encoding=encoding)')

text = text.replace('def format_event_summary(item) -> dict:', 'def format_event_summary(item, encoding: str = "utf-8") -> dict:')
text = text.replace('def format_event_full(item, body_max_length: int = 5000) -> dict:', 'def format_event_full(item, body_max_length: int = 5000, encoding: str = "utf-8") -> dict:')
text = text.replace('result = format_event_summary(item)', 'result = format_event_summary(item, encoding=encoding)')

text = text.replace('def format_task_summary(item) -> dict:', 'def format_task_summary(item, encoding: str = "utf-8") -> dict:')
text = text.replace('def format_task_full(item, body_max_length: int = 5000) -> dict:', 'def format_task_full(item, body_max_length: int = 5000, encoding: str = "utf-8") -> dict:')
text = text.replace('result = format_task_summary(item)', 'result = format_task_summary(item, encoding=encoding)')

# Pass encoding to unquote in formatting
text = re.sub(
    r'urllib\.parse\.unquote\((item\.Subject)\)',
    r'urllib.parse.unquote(\1, encoding=encoding)',
    text
)

with open(formatting_path, 'w', encoding='utf-8') as f:
    f.write(text)

# 2. Update server.py
with open(server_path, 'r', encoding='utf-8') as f:
    text = f.read()

# Replace using string replacements wherever possible to avoid regex errors
replacements = [
    # list_emails
    ('    unread_only: bool = False,\n) -> str:', '    unread_only: bool = False,\n    encoding: str = "utf-8",\n) -> str:'),
    ('def _list(outlook, namespace, folder, count, unread_only):', 'def _list(outlook, namespace, folder, count, unread_only, encoding):'),
    ('return await bridge.call(_list, folder, count, unread_only)', 'return await bridge.call(_list, folder, count, unread_only, encoding)'),
    ('format_email_summary(items.Item(i + 1))', 'format_email_summary(items.Item(i + 1), encoding=encoding)'),
    
    # read_email
    ('    folder: str = "inbox",\n) -> str:', '    folder: str = "inbox",\n    encoding: str = "utf-8",\n) -> str:'),
    ('def _read(outlook, namespace, entry_id, subject_search, folder):', 'def _read(outlook, namespace, entry_id, subject_search, folder, encoding):'),
    ('return await bridge.call(_read, entry_id, subject_search, folder)', 'return await bridge.call(_read, entry_id, subject_search, folder, encoding)'),
    ('format_email_full(item)', 'format_email_full(item, encoding=encoding)'),
    ('format_email_full(items.Item(1))', 'format_email_full(items.Item(1), encoding=encoding)'),

    # search_emails
    ('async def search_emails(\n    query: str,\n    folder: str = "inbox",\n    count: int = 10,\n) -> str:', 'async def search_emails(\n    query: str,\n    folder: str = "inbox",\n    count: int = 10,\n    encoding: str = "utf-8",\n) -> str:'),
    ('def _search(outlook, namespace, query, folder, count):', 'def _search(outlook, namespace, query, folder, count, encoding):'),
    ('return await bridge.call(_search, query, folder, count)', 'return await bridge.call(_search, query, folder, count, encoding)'),

    # list_events
    ('async def list_events(\n    start_date: str = "",\n    end_date: str = "",\n    count: int = 20,\n) -> str:', 'async def list_events(\n    start_date: str = "",\n    end_date: str = "",\n    count: int = 20,\n    encoding: str = "utf-8",\n) -> str:'),
    ('def _list(outlook, namespace, start_date, end_date, count):', 'def _list(outlook, namespace, start_date, end_date, count, encoding):'),
    ('return await bridge.call(_list, start_date, end_date, count)', 'return await bridge.call(_list, start_date, end_date, count, encoding)'),
    ('format_event_summary(item)', 'format_event_summary(item, encoding=encoding)'),

    # get_event
    ('async def get_event(entry_id: str) -> str:', 'async def get_event(entry_id: str, encoding: str = "utf-8") -> str:'),
    ('def _get(outlook, namespace, entry_id):', 'def _get(outlook, namespace, entry_id, encoding):'),
    ('return await bridge.call(_get, entry_id)', 'return await bridge.call(_get, entry_id, encoding)'),
    ('format_event_full(item)', 'format_event_full(item, encoding=encoding)'),

    # search_events (calendar)
    ('async def search_events(\n    query: str,\n    start_date: str = "",\n    end_date: str = "",\n    count: int = 10,\n) -> str:', 'async def search_events(\n    query: str,\n    start_date: str = "",\n    end_date: str = "",\n    count: int = 10,\n    encoding: str = "utf-8",\n) -> str:'),
    ('def _search(outlook, namespace, query, start_date, end_date, count):', 'def _search(outlook, namespace, query, start_date, end_date, count, encoding):'),
    ('return await bridge.call(_search, query, start_date, end_date, count)', 'return await bridge.call(_search, query, start_date, end_date, count, encoding)'),

    # list_tasks
    ('async def list_tasks(\n    include_completed: bool = False,\n    count: int = 20,\n) -> str:', 'async def list_tasks(\n    include_completed: bool = False,\n    count: int = 20,\n    encoding: str = "utf-8",\n) -> str:'),
    ('def _list(outlook, namespace, include_completed, count):', 'def _list(outlook, namespace, include_completed, count, encoding):'),
    ('return await bridge.call(_list, include_completed, count)', 'return await bridge.call(_list, include_completed, count, encoding)'),
    ('format_task_summary(items.Item(i + 1))', 'format_task_summary(items.Item(i + 1), encoding=encoding)'),

    # get_task
    ('async def get_task(entry_id: str) -> str:', 'async def get_task(entry_id: str, encoding: str = "utf-8") -> str:'),
    # Note: def _get(outlook, namespace, entry_id): and return await bridge.call(_get, entry_id) was already matched by get_event so we must be precise if they match exactly.
    # Ah, let's fix that. get_event and get_task both use def _get(outlook, namespace, entry_id):
    # Actually, they are in their respective functions. replace() will replace all instances globally.
    # But get_task has the EXACT same _get as get_event.
    # So the global replacement already caught it.
    
    # list_attachments
    ('async def list_attachments(entry_id: str) -> str:', 'async def list_attachments(entry_id: str, encoding: str = "utf-8") -> str:'),
    # Again, def _list(outlook, namespace, entry_id): might be replaced already if it matched exactly
    ('urllib.parse.unquote(att.FileName)', 'urllib.parse.unquote(att.FileName, encoding=encoding)'),

    # save_attachment
    ('async def save_attachment(\n    entry_id: str,\n    attachment_index: int = 1,\n    save_directory: str = "",\n) -> str:', 'async def save_attachment(\n    entry_id: str,\n    attachment_index: int = 1,\n    save_directory: str = "",\n    encoding: str = "utf-8",\n) -> str:'),
    ('def _save(outlook, namespace, entry_id, attachment_index, save_directory):', 'def _save(outlook, namespace, entry_id, attachment_index, save_directory, encoding):'),
    ('return await bridge.call(_save, entry_id, attachment_index, save_directory)', 'return await bridge.call(_save, entry_id, attachment_index, save_directory, encoding)'),
    
    # update format_task_full(item)
    ('format_task_full(item)', 'format_task_full(item, encoding=encoding)'),
    
    # In list_attachments: return await bridge.call(_list, entry_id) -> already replaced by get_event / read_email logic? Let\'s be careful.
]

for old, new in replacements:
    if old in text:
        text = text.replace(old, new)

# Special handling for duplicate function signatures getting double-replaced
# wait, if I did `text.replace` it replaces ALL of them at once.
# "def _get(outlook, namespace, entry_id):" -> "def _get(outlook, namespace, entry_id, encoding):"
# This will catch BOTH get_event and get_task, which is PERFECT.
# "return await bridge.call(_get, entry_id)" -> "return await bridge.call(_get, entry_id, encoding)"
# This also catches BOTH, perfect.

# Let's fix list_attachments `def _list(outlook, namespace, entry_id):`
# It's unique enough (only list_attachments has that exact signature)
if 'def _list(outlook, namespace, entry_id):' in text:
    text = text.replace('def _list(outlook, namespace, entry_id):', 'def _list(outlook, namespace, entry_id, encoding):')
if 'return await bridge.call(_list, entry_id)' in text:
    text = text.replace('return await bridge.call(_list, entry_id)', 'return await bridge.call(_list, entry_id, encoding)')

with open(server_path, 'w', encoding='utf-8') as f:
    f.write(text)

print("Done Refactoring Encodings!")
