#!/usr/bin/env python3
"""
Inject X cookies into Patchright browser profile.
Run after getting fresh cookies from user's browser.

Usage:
  docker exec canna python3 /workspace/scripts/inject_cookies.py
"""

import asyncio
import json
from pathlib import Path

try:
    from patchright.async_api import async_playwright
except ImportError:
    print("pip install patchright")
    exit(1)

COOKIES_PATH = Path("/workspace/scripts/x_cookies.json")
PROFILE_DIR = Path("/workspace/browser_profile/x_stealth")


async def inject_cookies():
    if not COOKIES_PATH.exists():
        print(f"❌ No cookies file at {COOKIES_PATH}")
        return False
    
    PROFILE_DIR.mkdir(parents=True, exist_ok=True)
    
    raw_cookies = json.loads(COOKIES_PATH.read_text())
    
    # Convert to browser format
    browser_cookies = []
    cookie_map = {
        "auth_token": {"httpOnly": True, "secure": True, "sameSite": "None"},
        "ct0": {"httpOnly": False, "secure": True, "sameSite": "Lax"},
        "twid": {"httpOnly": False, "secure": True, "sameSite": "None"},
        "kdt": {"httpOnly": True, "secure": True, "sameSite": "Lax"},
        "guest_id": {"httpOnly": False, "secure": True, "sameSite": "None"},
        "guest_id_ads": {"httpOnly": False, "secure": True, "sameSite": "None"},
        "guest_id_marketing": {"httpOnly": False, "secure": True, "sameSite": "None"},
        "personalization_id": {"httpOnly": False, "secure": True, "sameSite": "None"},
    }
    
    for name, value in raw_cookies.items():
        if name in cookie_map:
            cookie = {
                "name": name,
                "value": value,
                "domain": ".x.com",
                "path": "/",
                **cookie_map[name]
            }
            browser_cookies.append(cookie)
    
    print(f"🍪 Injecting {len(browser_cookies)} cookies...")
    
    async with async_playwright() as p:
        context = await p.chromium.launch_persistent_context(
            str(PROFILE_DIR),
            headless=True,
            args=["--no-sandbox"]
        )
        
        await context.add_cookies(browser_cookies)
        
        # Verify by loading X
        page = context.pages[0] if context.pages else await context.new_page()
        await page.goto("https://x.com/home", wait_until="domcontentloaded")
        await asyncio.sleep(3)
        
        if "login" in page.url.lower():
            print("❌ Cookies didn't work - still on login page")
            await context.close()
            return False
        
        print("✅ Cookies injected and verified!")
        print(f"   Profile saved to: {PROFILE_DIR}")
        await context.close()
        return True


if __name__ == "__main__":
    asyncio.run(inject_cookies())