import asyncio
import os
import sys
import json
import base64
from datetime import datetime
from playwright.async_api import async_playwright
from urllib.parse import urlparse

# Import browser finding utilities
try:
    from tools.find_browsers import get_all_browsers
except ImportError:
    sys.path.append(os.path.join(os.path.dirname(__file__), "tools"))
    from find_browsers import get_all_browsers

# ==========================================
# SETTINGS
# ==========================================
SHOW_AND_ASK_BROWSER = False
DEFAULT_BROWSER_INDEX = 4    # Camoufox Cache
DEFAULT_URL = "https://www.vap.org.tr/fon-turleri-bazinda-nakit-akisi"
ASK_OPENING_URL = False

class SessionManager:
    """Manages multi-page, multi-tab forensic sessions with deep sniffing."""
    def __init__(self, url):
        domain = urlparse(url).netloc.replace('.', '_') or "local"
        timestamp = datetime.now().strftime("%Y%m%d_%H%M")
        self.session_dir = os.path.join("sessions", f"capture_{timestamp}_{domain}")
        os.makedirs(self.session_dir, exist_ok=True)
        
        self.raw_log_path = os.path.join(self.session_dir, "raw_ui.jsonl")
        self.network_log_path = os.path.join(self.session_dir, "network_traffic.jsonl")
        self.websocket_log_path = os.path.join(self.session_dir, "websocket_traffic.jsonl")
        self.baseline_path = os.path.join(self.session_dir, "baseline.json")
        self.screenshot_dir = os.path.join(self.session_dir, "screenshots")
        self.js_dir = os.path.join(self.session_dir, "js_resources")
        self.auth_path = os.path.join(self.session_dir, "storage_state.json")
        
        os.makedirs(self.screenshot_dir, exist_ok=True)
        os.makedirs(self.js_dir, exist_ok=True)
        self.baseline_counter = 0
        self.action_counter = {} # Per page action counter
        self.captured_js_urls = set()

    def log_event(self, data, page_id="N/A", stream="ui"):
        """Logs data to UI, Network, or WebSocket streams."""
        try:
            if stream == "ui": target_path = self.raw_log_path
            elif stream == "network": target_path = self.network_log_path
            elif stream == "websocket": target_path = self.websocket_log_path
            else: target_path = self.raw_log_path

            payload = {
                "system_ts": datetime.now().isoformat(),
                "page_id": page_id,
                **data
            }
            with open(target_path, 'a', encoding='utf-8') as f:
                f.write(json.dumps(payload, ensure_ascii=False) + "\n")
                f.flush()
        except: pass

    async def capture_page_baseline(self, page, page_id, trigger_reason="INITIAL"):
        self.baseline_counter += 1
        safe_url = urlparse(page.url).path.replace('/', '_') or "root"
        if len(safe_url) > 20: safe_url = safe_url[:20]
        b_name = f"baseline_{self.baseline_counter}_{page_id}_{safe_url}"
        
        print(f"💎 [{page_id}] Snapshot ({self.baseline_counter}): {trigger_reason}")
        
        try:
            img_path = os.path.join(self.screenshot_dir, f"{b_name}.png")
            await page.screenshot(path=img_path, full_page=True)
            
            baseline = await page.evaluate("""async () => {
                return {
                    inventory: window.__WCS_DISCOVERY__ ? window.__WCS_DISCOVERY__.scan() : [],
                    state: { localStorage: {...localStorage}, sessionStorage: {...sessionStorage} },
                    meta: window.__FORENSIC__ ? window.__FORENSIC__.collectMetadata() : {}
                };
            }""")
            
            baseline.update({
                "url": page.url, "trigger": trigger_reason, "ts": datetime.now().isoformat(),
                "cookies": await page.context.cookies()
            })

            with open(os.path.join(self.session_dir, f"{b_name}.json"), 'w', encoding='utf-8') as f:
                json.dump(baseline, f, indent=2, ensure_ascii=False)
                
            self.log_event({"type": "BASELINE", "file": f"{b_name}.json", "reason": trigger_reason}, page_id)
        except Exception as e:
            print(f"  ⚠️ Snapshot Failed: {e}")

async def setup_network_sniffer(page, sm, page_id):
    """Deep inspection with Trace ID correlation."""
    
    async def request_handler(request):
        try:
            # Try to get the Trace ID from the page context
            trace_id = await page.evaluate("window.__WCS_LAST_TRACE_ID__ || null")
            
            req_data = {
                "type": "NET_REQUEST",
                "trace_id": trace_id,
                "method": request.method,
                "url": request.url,
                "resource_type": request.resource_type,
                "headers": request.headers,
                "post_data": request.post_data
            }
            sm.log_event(req_data, page_id, stream="network")
            
            # Live print for XHR/Fetch requests
            if request.resource_type in ["fetch", "xhr"]:
                print(f"🌐 [{page_id}] {request.method} {request.url[:60]}...")
        except: pass

    async def response_handler(response):
        try:
            # We don't re-query trace_id for response to maintain consistency with trigger
            res_data = {
                "type": "NET_RESPONSE",
                "url": response.url,
                "status": response.status,
                "headers": response.headers
            }
            if response.request.resource_type in ["fetch", "xhr", "script"]:
                try:
                    body = await response.body()
                    
                    # Special handling for JS capture
                    if response.request.resource_type == "script" and response.url not in sm.captured_js_urls:
                        sm.captured_js_urls.add(response.url)
                        try:
                            # Generate a safe filename
                            parsed = urlparse(response.url)
                            fname = os.path.basename(parsed.path) or "script.js"
                            if not fname.endswith(".js"): fname += ".js"
                            
                            # Add hash subset or timestamp to ensure uniqueness if needed, but for now simple unique name
                            safe_name = f"js_{len(sm.captured_js_urls)}_{fname}"
                            js_path = os.path.join(sm.js_dir, safe_name)
                            
                            with open(js_path, 'wb') as f:
                                f.write(body)
                            
                            sm.log_event({
                                "type": "JS_CAPTURE",
                                "url": response.url,
                                "file": safe_name,
                                "status": "SAVED"
                            }, page_id)
                            print(f"📦 [{page_id}] JS Captured: {safe_name}")
                        except Exception as js_err:
                            print(f"  ⚠️ JS Save Failed ({response.url}): {js_err}")

                    try:
                        res_data["body"] = body.decode('utf-8')
                        res_data["body_encoding"] = "utf-8"
                    except:
                        res_data["body_base64"] = base64.b64encode(body).decode('utf-8')
                        res_data["body_encoding"] = "base64"
                except: pass

            sm.log_event(res_data, page_id, stream="network")
        except: pass

    page.on("request", request_handler)
    page.on("response", response_handler)

async def setup_websocket_sniffer(page, sm, page_id):
    """Sniffs WebSocket with Trace ID support."""
    def handle_ws(ws):
        ws_id = f"ws_{id(ws)}"
        print(f"🔌 [{page_id}] WebSocket Opened.")

        async def frame_logged(direction, payload):
            tid = await page.evaluate("window.__WCS_LAST_TRACE_ID__ || null")
            sm.log_event({
                "type": f"WS_FRAME_{direction}", 
                "ws_id": ws_id, 
                "trace_id": tid,
                "payload": payload if isinstance(payload, str) else base64.b64encode(payload).decode('utf-8'),
                "is_binary": not isinstance(payload, str)
            }, page_id, stream="websocket")

        ws.on("framesent", lambda p: asyncio.create_task(frame_logged("SENT", p)))
        ws.on("framereceived", lambda p: asyncio.create_task(frame_logged("RECEIVED", p)))
        ws.on("close", lambda: sm.log_event({"type": "WS_CLOSED", "ws_id": ws_id}, page_id, stream="websocket"))

    page.on("websocket", handle_ws)

async def handle_page(page, sm, master_js):
    page_id = f"tab_{id(page)}"
    print(f"🌟 [New Tab] {page_id}")
    
    state = {"ready": False, "injected_frames": set()}
    sm.action_counter[page_id] = 0
    sm.log_event({"type": "TAB_OPENED", "url": page.url}, page_id)

    await setup_network_sniffer(page, sm, page_id)
    await setup_websocket_sniffer(page, sm, page_id)

    # Injection function for frames
    async def inject(frame):
        if not state["ready"] or frame.url == "about:blank": return
        try:
            await frame.evaluate(master_js)
            state["injected_frames"].add(frame)
        except: pass

    async def console_handler(msg):
        txt = msg.text
        if "WCS_DELTA>>" in txt:
            try:
                data = json.loads(txt.split("WCS_DELTA>>")[1])
                sm.log_event(data, page_id)
                
                dtype = data.get("type")
                action = data.get("action", dtype)
                
                # SPA State Shift Management
                if dtype == "UI_EVENT" and data.get("event") in ["CLICK", "INPUT", "KEYDOWN"]:
                    sm.action_counter[page_id] += 1
                    if sm.action_counter[page_id] % 10 == 0:
                        asyncio.create_task(sm.capture_page_baseline(page, page_id, f"ACTION_THRESHOLD_{sm.action_counter[page_id]}"))
                
                if dtype == "POPUP_SHOW":
                    asyncio.create_task(sm.capture_page_baseline(page, page_id, "POPUP_DETECTED"))

                # Enhanced Live Print
                sel = data.get('sel', '')
                val = data.get('val', '')
                text = data.get('text', '')
                
                if dtype in ["UI_EVENT", "HOVER", "HOVER_OUT"]:
                    display = text or val or sel
                    print(f"[{page_id}] ✨ {action}: {display[:60]}")
                
                elif dtype == "VALUE_DNA":
                    origin_icon = "🧬" if data.get('origin') == 'USER_TYPED' else "🤖"
                    print(f"[{page_id}] {origin_icon} {action}: {sel} = '{val[:40]}'")
                
                elif dtype == "STATE_CHANGE":
                    if action == "USER_IDLE":
                        print(f"[{page_id}] 😴 {action}")
                    elif action == "USER_ACTIVE":
                        duration_ms = data.get("was_idle_for_ms", 0)
                        print(f"[{page_id}] ⏰ {action} (idle: {duration_ms/1000:.1f}s)")
                
                elif dtype == "POPUP_SHOW":
                    print(f"[{page_id}] 🎯 {action}: {text[:50]}")
                
                elif dtype == "SCROLL":
                    y = data.get('y', 0)
                    print(f"[{page_id}] 📜 {action}: y={y}")
                
                elif dtype == "FOCUS":
                    print(f"[{page_id}] 🎯 {action}: {sel}")
                
                elif dtype == "CLIPBOARD":
                    print(f"[{page_id}] 📋 {action}")
                
                elif dtype == "KEYDOWN":
                    key = data.get('key', '')
                    print(f"[{page_id}] ⌨️ {action}: {key}")
                
                elif dtype == "INIT_METADATA":
                    print(f"[{page_id}] 🔧 Forensic metadata initialized")
                    
            except: pass
        elif msg.type == "error":
            sm.log_event({"type": "CONSOLE", "level": "error", "text": txt}, page_id)
            print(f"🚨 [{page_id}] CONSOLE ERROR: {txt[:80]}...")

    page.on("console", lambda m: asyncio.create_task(console_handler(m)))
    page.on("frameattached", lambda f: asyncio.create_task(inject(f)))
    page.on("close", lambda: sm.log_event({"type": "TAB_CLOSED"}, page_id))

    async def on_navigated(frame):
        if frame == page.main_frame:
            await asyncio.sleep(2)
            state["ready"] = True
            await sm.capture_page_baseline(page, page_id, "NAVIGATION_OR_RELOAD")
            for f in page.frames: await inject(f)

    page.on("framenavigated", lambda f: asyncio.create_task(on_navigated(f)))

    if page.url != "about:blank":
        state["ready"] = True
        await inject(page.main_frame)
        asyncio.create_task(sm.capture_page_baseline(page, page_id, "INITIAL_DETECT"))

    async def maintenance():
        while not page.is_closed():
            await asyncio.sleep(5)
            for f in page.frames:
                if f not in state["injected_frames"]: await inject(f)
    asyncio.create_task(maintenance())

async def launch_forensic_browser(browser_info, url):
    sm = SessionManager(url)
    
    # Load Unified JS Sentinel
    js_path = os.path.join(os.path.dirname(__file__), "js_files", "sentinel_agent.js")
    if os.path.exists(js_path):
        with open(js_path, 'r', encoding='utf-8') as f:
            master_js = f.read()
    else:
        print(f"⚠️ Critical Error: {js_path} not found!")
        sys.exit(1)

    async with async_playwright() as p:
        engine = browser_info.get("type", "chromium")
        browser_type = p.chromium if engine != "firefox" else p.firefox
        launch_args = ["--start-maximized"] if engine != "firefox" else []
        
        print(f"\n🚀 Launching UNIFIED SENTINEL Forensic Suite...")
        print(f"🎬 Browser: {browser_info['name']}")
        print(f"📂 Session: {sm.session_dir}\n")

        browser = await browser_type.launch(executable_path=browser_info['path'], headless=False, args=launch_args)
        context = await browser.new_context(no_viewport=True, bypass_csp=True)
        
        await context.expose_binding("__py_event", lambda s, e: sm.log_event({"raw_py_event": e}, f"tab_{id(s['page'])}"))
        context.on("page", lambda p: asyncio.create_task(handle_page(p, sm, master_js)))

        page = await context.new_page()
        await page.goto(url, wait_until="commit")

        try:
            while len(context.pages) > 0:
                await asyncio.sleep(1)
        except KeyboardInterrupt: pass
        finally:
            try:
                await context.storage_state(path=sm.auth_path)
                await browser.close()
            except: pass
            print(f"\n👋 Done. Data saved in: {sm.session_dir}")

def main():
    browsers = get_all_browsers()
    if not browsers: return
    asyncio.run(launch_forensic_browser(browsers[DEFAULT_BROWSER_INDEX], DEFAULT_URL))

if __name__ == "__main__":
    main()