"""
Patch the installed livekit-rtc SDK to support HKDF E2EE fields.

Run after pip install in the Docker build. Adds key_ring_size and
key_derivation_function fields to:
1. KeyProviderOptions dataclass (e2ee.py)
2. Proto conversion in Room.connect() (room.py)

These fields are added by the EC-compat Rust fork and are required
for Element Call E2EE compatibility.
"""

import os
import sys


def get_package_dir():
    """Find the installed livekit.rtc package directory."""
    import livekit.rtc as rtc
    return os.path.dirname(rtc.__file__)


def patch_e2ee(pkg_dir: str):
    """Add key_ring_size and key_derivation_function to KeyProviderOptions."""
    e2ee_path = os.path.join(pkg_dir, "e2ee.py")
    with open(e2ee_path) as f:
        content = f.read()

    if "key_ring_size" in content:
        print("e2ee.py already patched, skipping")
        return

    # Add new fields after failure_tolerance
    content = content.replace(
        "failure_tolerance: int = DEFAULT_FAILURE_TOLERANCE",
        "failure_tolerance: int = DEFAULT_FAILURE_TOLERANCE\n"
        "    key_ring_size: int = 16\n"
        "    key_derivation_function: int = 0  # 0=PBKDF2, 1=HKDF",
    )

    with open(e2ee_path, "w") as f:
        f.write(content)
    print(f"Patched {e2ee_path}: added key_ring_size, key_derivation_function")


def patch_room(pkg_dir: str):
    """Add key_ring_size and key_derivation_function to Room.connect() proto conversion."""
    room_path = os.path.join(pkg_dir, "room.py")
    with open(room_path) as f:
        content = f.read()

    if "key_ring_size" in content:
        print("room.py already patched, skipping")
        return

    # Patch the deprecated e2ee path (used by livekit-agents)
    old_e2ee = (
        "req.connect.options.e2ee.key_provider_options.ratchet_window_size = (\n"
        "                options.e2ee.key_provider_options.ratchet_window_size\n"
        "            )"
    )
    new_e2ee = (
        "req.connect.options.e2ee.key_provider_options.ratchet_window_size = (\n"
        "                options.e2ee.key_provider_options.ratchet_window_size\n"
        "            )\n"
        "            if hasattr(options.e2ee.key_provider_options, 'key_ring_size'):\n"
        "                req.connect.options.e2ee.key_provider_options.key_ring_size = (\n"
        "                    options.e2ee.key_provider_options.key_ring_size\n"
        "                )\n"
        "            if hasattr(options.e2ee.key_provider_options, 'key_derivation_function'):\n"
        "                req.connect.options.e2ee.key_provider_options.key_derivation_function = (\n"
        "                    options.e2ee.key_provider_options.key_derivation_function\n"
        "                )"
    )
    content = content.replace(old_e2ee, new_e2ee)

    # Patch the current encryption path too
    old_enc = (
        "req.connect.options.encryption.key_provider_options.ratchet_window_size = (\n"
        "                options.encryption.key_provider_options.ratchet_window_size\n"
        "            )"
    )
    new_enc = (
        "req.connect.options.encryption.key_provider_options.ratchet_window_size = (\n"
        "                options.encryption.key_provider_options.ratchet_window_size\n"
        "            )\n"
        "            if hasattr(options.encryption.key_provider_options, 'key_ring_size'):\n"
        "                req.connect.options.encryption.key_provider_options.key_ring_size = (\n"
        "                    options.encryption.key_provider_options.key_ring_size\n"
        "                )\n"
        "            if hasattr(options.encryption.key_provider_options, 'key_derivation_function'):\n"
        "                req.connect.options.encryption.key_provider_options.key_derivation_function = (\n"
        "                    options.encryption.key_provider_options.key_derivation_function\n"
        "                )"
    )
    content = content.replace(old_enc, new_enc)

    with open(room_path, "w") as f:
        f.write(content)
    print(f"Patched {room_path}: added key_ring_size, key_derivation_function to proto conversion")


def main():
    pkg_dir = get_package_dir()
    print(f"Patching livekit-rtc at {pkg_dir}")
    patch_e2ee(pkg_dir)
    patch_room(pkg_dir)
    print("SDK patching complete")


if __name__ == "__main__":
    main()