import os

def patch_langmem_for_python310():
    # Path where langmem is installed inside the Docker container on HF Spaces
    base_path = "/usr/local/lib/python3.10/site-packages/langmem/knowledge"

    extraction_path = os.path.join(base_path, "extraction.py")
    tools_path = os.path.join(base_path, "tools.py")

    # === PATCH extraction.py ===
    if os.path.exists(extraction_path):
        with open(extraction_path, "r", encoding="utf-8") as file:
            content = file.read()

        if "typing.NotRequired" in content:
            content = content.replace("typing.NotRequired", "NotRequired")
            if "from typing_extensions import NotRequired" not in content:
                content = content.replace("import uuid", "import uuid\nfrom typing_extensions import NotRequired")

            with open(extraction_path, "w", encoding="utf-8") as file:
                file.write(content)
            print("✅ Patched extraction.py")
        else:
            print("✅ extraction.py already patched")

    # === PATCH tools.py ===
    if os.path.exists(tools_path):
        with open(tools_path, "r", encoding="utf-8") as file:
            content = file.read()

        if "typing.Literal[*actions_permitted]" in content:
            fixed = '''typing.Optional[
        tuple[typing.Literal["create", "update", "delete"], ...]
    ] = ("create", "update", "delete"),'''
            content = content.replace("typing.Literal[*actions_permitted]", fixed)

            with open(tools_path, "w", encoding="utf-8") as file:
                file.write(content)
            print("✅ Patched tools.py")
        else:
            print("✅ tools.py already patched")

patch_langmem_for_python310()