#!/usr/bin/env python3
"""
Headed X Engagement - runs with virtual display (Xvfb)
Uses persistent browser profile from VNC login session.

Usage:
  # First time: Start VNC and log in
  docker exec -it canna bash /workspace/scripts/start_vnc_browser.sh
  
  # Then run engagement (can be cron'd)
  docker exec -e DISPLAY=:99 canna python3 /workspace/scripts/headed_engagement.py
"""

import asyncio
import json
import os
import random
import sys
import urllib.request
from datetime import datetime
from pathlib import Path

# Ensure display is set
if "DISPLAY" not in os.environ:
    os.environ["DISPLAY"] = ":99"

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

PROFILE_DIR = Path("/workspace/browser_profile/x_headed")
COOKIES_PATH = Path("/workspace/scripts/x_cookies.json")
LOG_PATH = Path("/workspace/PROGRESS-LOG.md")

VOICE_PROFILE = """You are a Florida medical cannabis patient in your late 20s. Sassy, funny, self-deprecating. 
Lowercase twitter speak, 1-2 emojis max. Never sound like a brand. React to what they ACTUALLY said.

IMPORTANT - You know FL cannabis products:
- Flower = smokable bud (strains like Runtz, GMO, etc)
- Piatella/Rosin/Live Rosin/Hash = solventless concentrates you DAB (vaporize), NOT eat
- Shatter/Crumble/Wax = BHO concentrates you dab
- Vapes/Carts = vaporizer cartridges
- Edibles = food products (gummies, chocolates)
- RSO = Rick Simpson Oil, oral/sublingual

If they show a concentrate (piatella, rosin, hash), talk about DABBING it, not eating. 
Reference effects like "melted", "couch locked", "blasted", "zooted"."""

TARGET_ACCOUNTS = [
    "Trulieve", "MUV_FL", "TheFloweryFL", "GrowHealthyFL",
    "SunburnCannabis", "curaleaffl", "JungleBoysFlorida"
]


def log_action(action: str, details: str = ""):
    ts = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    print(f"[{ts}] {action}: {details}", flush=True)
    try:
        content = LOG_PATH.read_text() if LOG_PATH.exists() else "# PROGRESS-LOG.md\n"
        today = datetime.now().strftime("%Y-%m-%d")
        if f"### {today}" not in content:
            content += f"\n### {today}\n| Time | Action | Details |\n|------|--------|--------|\n"
        content += f"| {ts} | {action} | {details} |\n"
        LOG_PATH.write_text(content)
    except:
        pass


def load_cookies():
    """Load cookies from JSON if profile doesn't have them."""
    if not COOKIES_PATH.exists():
        return []
    raw = json.loads(COOKIES_PATH.read_text())
    cookies = []
    for name in ["auth_token", "ct0", "twid", "kdt", "guest_id", "personalization_id"]:
        if name in raw and raw[name]:
            cookies.append({
                "name": name, "value": raw[name], "domain": ".x.com", "path": "/",
                "secure": True, "sameSite": "None" if name != "ct0" else "Lax"
            })
    return cookies


def generate_reply(tweet_text: str, username: str) -> str:
    """Generate reply via LLM."""
    if not tweet_text or len(tweet_text.strip()) < 10:
        return random.choice(["👀", "say less", "on my way"])
    
    prompt = f"""{VOICE_PROFILE}

Tweet from @{username}: "{tweet_text}"

Write ONE reply (under 150 chars). Just the reply:"""

    try:
        payload = json.dumps({
            "model": "claude-sonnet-4-20250514",
            "max_tokens": 80,
            "messages": [{"role": "user", "content": prompt}]
        }).encode()
        
        req = urllib.request.Request(
            "https://api.z.ai/api/anthropic/v1/messages",
            data=payload,
            headers={
                "Content-Type": "application/json",
                "x-api-key": "3887e2ed1a0a42ada0f9dc04461d3dce.O9dPzq1s2KQB2uMm",
                "anthropic-version": "2023-06-01"
            }
        )
        with urllib.request.urlopen(req, timeout=20) as resp:
            data = json.loads(resp.read().decode())
            reply = data["content"][0]["text"].strip().strip('"\'')
            if 5 < len(reply) <= 280:
                return reply
    except Exception as e:
        print(f"  LLM error: {e}", flush=True)
    
    return random.choice(["my wallet didn't need this", "say less 🏃‍♂️", "📝"])


async def human_type(page, text: str):
    """Type like a human."""
    for char in text:
        await page.keyboard.type(char, delay=random.randint(35, 90))
        if random.random() < 0.02:
            await asyncio.sleep(random.uniform(0.2, 0.6))


async def dismiss_overlays(page):
    """Aggressively dismiss X's popups."""
    # Cookie consent
    for sel in ['[data-testid="twc-cc-mask"]', '[data-testid="mask"]']:
        try:
            el = page.locator(sel)
            if await el.is_visible(timeout=300):
                await page.keyboard.press("Escape")
                await asyncio.sleep(0.3)
        except:
            pass
    
    # Close buttons
    for sel in ['[aria-label="Close"]', 'button:has-text("Got it")', 'button:has-text("Not now")']:
        try:
            el = page.locator(sel).first
            if await el.is_visible(timeout=300):
                await el.click()
                await asyncio.sleep(0.3)
        except:
            pass
    
    # Final escape
    await page.keyboard.press("Escape")
    await asyncio.sleep(0.2)


async def ensure_xvfb():
    """Start Xvfb if not running."""
    import subprocess
    display = os.environ.get("DISPLAY", ":99")
    display_num = display.replace(":", "")
    
    # Check if Xvfb is running
    result = subprocess.run(["pgrep", "-f", f"Xvfb :{display_num}"], capture_output=True)
    if result.returncode != 0:
        print(f"🖥️  Starting Xvfb on {display}...", flush=True)
        subprocess.Popen(
            ["Xvfb", f":{display_num}", "-screen", "0", "1920x1080x24"],
            stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL
        )
        await asyncio.sleep(2)


async def run_engagement(reply_count: int = 2, like_count: int = 5):
    """Run headed engagement."""
    print("🥷 Headed X Engagement", flush=True)
    print(f"   DISPLAY={os.environ.get('DISPLAY', 'not set')}", flush=True)
    print("=" * 50, flush=True)
    
    await ensure_xvfb()
    
    PROFILE_DIR.mkdir(parents=True, exist_ok=True)
    stats = {"likes": 0, "replies": 0}
    
    async with async_playwright() as p:
        # Launch HEADED browser (not headless)
        context = await p.chromium.launch_persistent_context(
            str(PROFILE_DIR),
            headless=False,  # HEADED MODE
            viewport={"width": 1366, "height": 768},
            args=[
                "--no-sandbox",
                "--disable-dev-shm-usage",
                "--disable-blink-features=AutomationControlled",
            ]
        )
        
        # Inject cookies as backup
        cookies = load_cookies()
        if cookies:
            print(f"🍪 Injecting {len(cookies)} backup cookies...", flush=True)
            await context.add_cookies(cookies)
        
        page = context.pages[0] if context.pages else await context.new_page()
        
        try:
            print("📱 Loading X...", flush=True)
            await page.goto("https://x.com/home", wait_until="domcontentloaded")
            await asyncio.sleep(4)
            
            # Dismiss any initial popups
            await dismiss_overlays(page)
            
            if "login" in page.url.lower():
                print("❌ Not logged in. Run start_vnc_browser.sh first.", flush=True)
                await context.close()
                return stats
            
            log_action("SESSION", "Logged in, starting headed engagement")
            
            # Warmup
            print("🔥 Warming up...", flush=True)
            for _ in range(4):
                await page.mouse.wheel(0, random.randint(200, 400))
                await asyncio.sleep(random.uniform(1, 2.5))
                await dismiss_overlays(page)
            
            # Engage
            targets = random.sample(TARGET_ACCOUNTS, min(3, len(TARGET_ACCOUNTS)))
            
            for username in targets:
                if stats["likes"] >= like_count and stats["replies"] >= reply_count:
                    break
                
                print(f"\n👤 @{username}...", flush=True)
                
                try:
                    await page.goto(f"https://x.com/{username}", wait_until="domcontentloaded")
                    await asyncio.sleep(random.uniform(3, 5))
                    await dismiss_overlays(page)
                    
                    # Scroll
                    for _ in range(2):
                        await page.mouse.wheel(0, random.randint(150, 300))
                        await asyncio.sleep(random.uniform(0.8, 1.5))
                    
                    await dismiss_overlays(page)
                    
                    tweets = await page.locator('[data-testid="tweet"]').all()
                    
                    for i, tweet in enumerate(tweets[:3]):
                        if stats["likes"] >= like_count and stats["replies"] >= reply_count:
                            break
                        
                        try:
                            await dismiss_overlays(page)
                            
                            # Like
                            if stats["likes"] < like_count:
                                like_btn = tweet.locator('[data-testid="like"]')
                                if await like_btn.is_visible():
                                    try:
                                        await like_btn.click(timeout=5000)
                                        stats["likes"] += 1
                                        log_action("LIKE", f"@{username}")
                                        await asyncio.sleep(random.uniform(1.5, 3))
                                    except:
                                        await dismiss_overlays(page)
                            
                            # Reply (first tweet)
                            if stats["replies"] < reply_count and i == 0:
                                print(f"  💬 Attempting reply...", flush=True)
                                tweet_text = ""
                                try:
                                    txt = tweet.locator('[data-testid="tweetText"]')
                                    if await txt.count() > 0:
                                        tweet_text = await txt.inner_text()
                                        print(f"     Tweet: {tweet_text[:50]}...", flush=True)
                                except Exception as e:
                                    print(f"     No tweet text: {e}", flush=True)
                                
                                reply_btn = tweet.locator('[data-testid="reply"]')
                                btn_visible = await reply_btn.is_visible()
                                print(f"     Reply btn visible: {btn_visible}", flush=True)
                                if btn_visible:
                                    # Try clicking tweet first to open detail view
                                    try:
                                        tweet_link = tweet.locator('a[href*="/status/"]').first
                                        if await tweet_link.is_visible():
                                            await tweet_link.click()
                                            print(f"     Opened tweet detail view...", flush=True)
                                            await asyncio.sleep(random.uniform(2, 3))
                                            await dismiss_overlays(page)
                                    except:
                                        pass
                                    
                                    # Now look for reply input on detail page
                                    print(f"     FULL TWEET TEXT: [{tweet_text}]", flush=True)  # Log full text!
                                    reply_text = generate_reply(tweet_text, username)
                                    print(f"     Generated: {reply_text}", flush=True)  # Log full reply!
                                    
                                    # Try inline reply box first
                                    textarea = page.locator('[data-testid="tweetTextarea_0"]')
                                    textarea_visible = False
                                    try:
                                        textarea_visible = await textarea.is_visible(timeout=2000)
                                    except:
                                        pass
                                    
                                    # If no inline textarea, try keyboard shortcut 'r' for reply
                                    if not textarea_visible:
                                        print(f"     Trying keyboard shortcut 'r'...", flush=True)
                                        await page.keyboard.press("r")
                                        await asyncio.sleep(random.uniform(2, 3))
                                        await dismiss_overlays(page)
                                        try:
                                            textarea_visible = await textarea.is_visible(timeout=3000)
                                        except:
                                            pass
                                    
                                    # Still no textarea? Try clicking reply button
                                    if not textarea_visible:
                                        reply_btn2 = page.locator('[data-testid="reply"]').first
                                        if await reply_btn2.is_visible():
                                            await reply_btn2.click()
                                            print(f"     Clicked reply btn...", flush=True)
                                            await asyncio.sleep(random.uniform(2, 3))
                                            await dismiss_overlays(page)
                                            try:
                                                textarea_visible = await textarea.is_visible(timeout=3000)
                                            except:
                                                pass
                                    
                                    print(f"     Textarea visible: {textarea_visible}", flush=True)
                                    if textarea_visible:
                                        # Focus and clear
                                        await textarea.click()
                                        await asyncio.sleep(0.3)
                                        
                                        # Type directly into the element
                                        print(f"     Typing reply...", flush=True)
                                        await textarea.fill("")  # Clear first
                                        await textarea.type(reply_text, delay=random.randint(40, 80))
                                        await asyncio.sleep(random.uniform(1, 2))
                                        
                                        # Take screenshot to debug
                                        ss = PROFILE_DIR.parent / f"reply_attempt_{datetime.now().strftime('%H%M%S')}.png"
                                        await page.screenshot(path=str(ss))
                                        print(f"     Screenshot: {ss.name}", flush=True)
                                        
                                        # Find Reply button - try multiple selectors
                                        submit = page.locator('[data-testid="tweetButton"], [data-testid="tweetButtonInline"], button:has-text("Reply")').first
                                        btn_visible = False
                                        try:
                                            btn_visible = await submit.is_visible(timeout=2000)
                                        except:
                                            pass
                                        print(f"     Submit btn visible: {btn_visible}", flush=True)
                                        
                                        if btn_visible:
                                            # Use Ctrl+Enter to submit (keyboard shortcut)
                                            print(f"     Submitting with Ctrl+Enter...", flush=True)
                                            await page.keyboard.press("Control+Enter")
                                            await asyncio.sleep(3)
                                            
                                            # Screenshot AFTER posting attempt
                                            ss2 = PROFILE_DIR.parent / f"after_post_{datetime.now().strftime('%H%M%S')}.png"
                                            await page.screenshot(path=str(ss2))
                                            print(f"     Post screenshot: {ss2.name}", flush=True)
                                            
                                            # Check for success toast "Your post was sent"
                                            toast = page.locator('text="Your post was sent"')
                                            toast_visible = False
                                            try:
                                                toast_visible = await toast.is_visible(timeout=2000)
                                            except:
                                                pass
                                            
                                            # Also check if textarea was cleared (text is gone)
                                            textarea_cleared = False
                                            try:
                                                current_text = await textarea.input_value()
                                                textarea_cleared = not current_text or current_text.strip() == ""
                                            except:
                                                textarea_cleared = True  # Element gone = success
                                            
                                            print(f"     Toast visible: {toast_visible}, textarea cleared: {textarea_cleared}", flush=True)
                                            
                                            if toast_visible or textarea_cleared:
                                                stats["replies"] += 1
                                                log_action("REPLY", f"@{username}: {reply_text[:40]}...")
                                                print(f"  ✅ {reply_text[:50]}...", flush=True)
                                            else:
                                                print(f"  ⚠️ Reply may not have posted", flush=True)
                                                await page.keyboard.press("Escape")
                                                await asyncio.sleep(0.5)
                                        else:
                                            print(f"  ⚠️ Submit btn not found, trying click anyway...", flush=True)
                                            try:
                                                await submit.click(force=True, timeout=3000)
                                                await asyncio.sleep(3)
                                                stats["replies"] += 1
                                                log_action("REPLY", f"@{username}: {reply_text[:40]}...")
                                                print(f"  ✅ Force clicked!", flush=True)
                                            except Exception as e:
                                                print(f"  ⚠️ Force click failed: {str(e)[:40]}", flush=True)
                                    else:
                                        await page.keyboard.press("Escape")
                                        
                        except Exception as e:
                            print(f"  Error: {str(e)[:60]}", flush=True)
                            await dismiss_overlays(page)
                    
                    await asyncio.sleep(random.uniform(2, 4))
                    
                except Exception as e:
                    print(f"  Error @{username}: {str(e)[:60]}", flush=True)
            
            print("\n" + "=" * 50, flush=True)
            print(f"📊 DONE - Likes: {stats['likes']}, Replies: {stats['replies']}", flush=True)
            log_action("HEADED_DONE", f"L:{stats['likes']} R:{stats['replies']}")
            
        except Exception as e:
            print(f"Error: {e}", flush=True)
            ss = PROFILE_DIR.parent / f"error_{datetime.now().strftime('%H%M%S')}.png"
            await page.screenshot(path=str(ss))
            print(f"Screenshot: {ss}", flush=True)
        finally:
            await context.close()
    
    return stats


async def main():
    if "--dry-run" in sys.argv:
        print("DRY RUN - targets:", TARGET_ACCOUNTS[:3])
        return
    
    # Default values
    reply_count = 2
    like_count = 5
    
    # Presets
    if "--heavy" in sys.argv:
        reply_count, like_count = 4, 10
    elif "--light" in sys.argv:
        reply_count, like_count = 1, 3
    
    # Explicit overrides: --replies N --likes N
    for i, arg in enumerate(sys.argv):
        if arg == "--replies" and i + 1 < len(sys.argv):
            reply_count = int(sys.argv[i + 1])
        elif arg == "--likes" and i + 1 < len(sys.argv):
            like_count = int(sys.argv[i + 1])
    
    print(f"Config: {reply_count} replies, {like_count} likes", flush=True)
    await run_engagement(reply_count, like_count)


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