import os

layout_path = 'src/app/layout.tsx'
with open(layout_path, 'r') as f:
    layout_content = f.read()

drm_component = """
'use client';
import { useEffect } from 'react';

export function DRMWrapper({ children }: { children: React.ReactNode }) {
  useEffect(() => {
    const handleContextMenu = (e: MouseEvent) => {
      // Prevent right click on images
      if (e.target instanceof HTMLImageElement) {
        e.preventDefault();
      }
    };
    const handleDragStart = (e: DragEvent) => {
      if (e.target instanceof HTMLImageElement) {
        e.preventDefault();
      }
    };
    const handleKeyDown = (e: KeyboardEvent) => {
      if ((e.ctrlKey || e.metaKey) && (e.key === 's' || e.key === 'p')) {
        e.preventDefault();
      }
    };
    
    document.addEventListener('contextmenu', handleContextMenu);
    document.addEventListener('dragstart', handleDragStart);
    document.addEventListener('keydown', handleKeyDown);
    
    return () => {
      document.removeEventListener('contextmenu', handleContextMenu);
      document.removeEventListener('dragstart', handleDragStart);
      document.removeEventListener('keydown', handleKeyDown);
    };
  }, []);

  return <div className="h-full">{children}</div>;
}
"""

with open('src/components/DRMWrapper.tsx', 'w') as f:
    f.write(drm_component)

if "DRMWrapper" not in layout_content:
    layout_content = layout_content.replace('import { ThemeProvider', 'import { DRMWrapper } from "@/components/DRMWrapper";\nimport { ThemeProvider')
    layout_content = layout_content.replace('<body>', '<body>\n        <DRMWrapper>')
    layout_content = layout_content.replace('</body>', '        </DRMWrapper>\n      </body>')
    with open(layout_path, 'w') as f:
        f.write(layout_content)

print("DRM added.")