import sys
import json
import re
import time
import base64
from selenium import webdriver
from selenium.webdriver.chrome.options import Options

def parse_kmb(val_str):
    if not val_str: return 0
    val_str = val_str.upper().replace(',', '').strip()
    multiplier = 1
    if val_str.endswith('K'):
        multiplier = 1000
        val_str = val_str[:-1]
    elif val_str.endswith('M'):
        multiplier = 1000000
        val_str = val_str[:-1]
    elif val_str.endswith('B'):
        multiplier = 1000000000
        val_str = val_str[:-1]
    try:
        return int(float(val_str) * multiplier)
    except:
        return 0

def main():
    if len(sys.argv) < 3:
        print(json.dumps({"error": "Missing args"}))
        sys.exit(1)
        
    url = sys.argv[1]
    platform = sys.argv[2]
    
    options = Options()
    options.add_argument("--headless=new")
    options.add_argument("--disable-gpu")
    options.add_argument("--no-sandbox")
    options.add_argument("--disable-extensions")
    options.add_argument("--window-size=1366,768")
    options.add_argument("--disable-blink-features=AutomationControlled")
    options.add_argument("--user-agent=Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/121.0.0.0 Safari/537.36")
    options.page_load_strategy = 'eager'
    
    try:
        driver = webdriver.Chrome(options=options)
        driver.set_page_load_timeout(20)
    except Exception as e:
        print(json.dumps({"error": f"Failed to start browser: {str(e)}"}))
        sys.exit(1)
        
    try:
        driver.get(url)
        time.sleep(5) # Wait for render
        
        html = driver.page_source or ""
        page_text = driver.execute_script("return document.body.innerText;") or ""
        
        screenshot = ""
        try:
            screenshot = driver.get_screenshot_as_base64()
        except: pass
        
        views, likes, comments, author = 0, 0, 0, ""
        
        if platform == "TikTok":
            for pattern, field in [
                (r'"playCount"\s*:\s*(\d+)', 'views'),
                (r'"diggCount"\s*:\s*(\d+)', 'likes'),
                (r'"commentCount"\s*:\s*(\d+)', 'comments')
            ]:
                m = re.search(pattern, html)
                if m:
                    if field == 'views': views = int(m.group(1))
                    elif field == 'likes': likes = int(m.group(1))
                    elif field == 'comments': comments = int(m.group(1))
                    
            if not likes:
                like_m = re.search(r'data-e2e="like-count"[^>]*>\s*([\d\.]+[KMB]?)\s*<', html, re.IGNORECASE)
                if like_m: likes = parse_kmb(like_m.group(1))
            if not comments:
                comment_m = re.search(r'data-e2e="comment-count"[^>]*>\s*([\d\.]+[KMB]?)\s*<', html, re.IGNORECASE)
                if comment_m: comments = parse_kmb(comment_m.group(1))
            if not views:
                view_m = re.search(r'data-e2e="video-views"[^>]*>\s*([\d\.]+[KMB]?)\s*<', html, re.IGNORECASE)
                if view_m: views = parse_kmb(view_m.group(1))
                    
        elif platform == "YouTube":
            vm = re.search(r'"viewCount"\s*:\s*"(\d+)"', html)
            if vm: views = int(vm.group(1))
            lm = re.search(r'"likeCount"\s*:\s*"(\d+)"', html)
            if lm: likes = int(lm.group(1))
            
        elif platform == "Twitter/X":
            vm = re.search(r'"views"\s*:\s*\{\s*"count"\s*:\s*"(\d+)"', html)
            if vm: views = int(vm.group(1))
            lm = re.search(r'"favorite_count"\s*:\s*(\d+)', html)
            if lm: likes = int(lm.group(1))
            
        print(json.dumps({
            "views": views,
            "likes": likes,
            "comments": comments,
            "author": author,
            "screenshot": screenshot
        }))
        
    except Exception as e:
        print(json.dumps({"error": str(e)}))
    finally:
        driver.quit()

if __name__ == "__main__":
    main()