# /home/mike/repos/pipulate/tools/scraper_tools.py
import asyncio
import json
import os
import sys
import shutil
import tempfile
from datetime import datetime
from pathlib import Path
from urllib.parse import quote, urlparse
import random
import time

from loguru import logger
import undetected_chromedriver as uc
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from selenium.webdriver.common.by import By

from tools import auto_tool
from . import dom_tools

async def generate_optics_subprocess(target_dir_path: str):
    """Isolated wrapper to call llm_optics.py as a subprocess, protecting the event loop."""
    script_path = Path(__file__).resolve().parent / "llm_optics.py"
    
    proc = await asyncio.create_subprocess_exec(
        sys.executable, str(script_path), str(target_dir_path),
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE
    )
    
    stdout, stderr = await proc.communicate()
    
    if proc.returncode == 0:
        return {"success": True, "output": stdout.decode()}
    else:
        return {"success": False, "error": stderr.decode()}

def get_safe_path_component(url: str) -> tuple[str, str]:
    """Converts a URL into filesystem-safe components for directory paths."""
    parsed = urlparse(url)
    domain = parsed.netloc
    path = parsed.path
    if not path or path == '/':
        path_slug = "%2F"
    else:
        path_slug = quote(path, safe='').replace('/', '_')[:100]
    return domain, path_slug


def _simplify_html_for_llm(html_content, default_title=""):
    """Applies a symmetrical, opinionated filter to HTML for LLM consumption."""
    try:
        from bs4 import BeautifulSoup
        soup = BeautifulSoup(html_content, 'html.parser')
        
        # Remove all noise elements that confuse LLMs (Added 'svg' to the hit list!)
        for tag in soup(['script', 'style', 'noscript', 'meta', 'link', 'head', 'svg']):
            tag.decompose()

        # Clean up attributes - keep only automation-relevant ones
        for element in soup.find_all():
            attrs_to_keep = {}
            for attr, value in element.attrs.items():
                # Added 'rel' and 'target' to preserve SEO link data!
                if attr in ['id', 'role', 'data-testid', 'name', 'type', 'href', 'src', 'class', 'for', 'value', 'placeholder', 'title', 'rel', 'target'] or attr.startswith('aria-'):
                    attrs_to_keep[attr] = value
            element.attrs = attrs_to_keep
        
        simple_html = soup.prettify()
    except Exception as e:
        logger.warning(f"⚠️ DOM simplification failed, using fallback: {e}")
        simple_html = html_content

    # Add minimal metadata wrapper
    title = soup.title.string if soup and soup.title else default_title
    final_html = f"<html>\n<head><title>{title}</title></head>\n<body>\n{simple_html}\n</body>\n</html>"
    return final_html


@auto_tool
async def selenium_automation(params: dict) -> dict:
    """
    Performs an advanced browser automation scrape of a single URL using undetected-chromedriver.
    Checks for cached data before initiating a new scrape.
    ...
    """
    url = params.get("url")
    domain = params.get("domain")
    url_path_slug = params.get("url_path_slug")
    take_screenshot = params.get("take_screenshot", False)
    headless = params.get("headless", True)
    is_notebook_context = params.get("is_notebook_context", False)
    persistent = params.get("persistent", False)
    profile_name = params.get("profile_name", "default")
    verbose = params.get("verbose", True)
    override_cache = params.get("override_cache", False)
    delay_range = params.get("delay_range")

    if not all([url, domain, url_path_slug is not None]):
        return {"success": False, "error": "URL, domain, and url_path_slug parameters are required."}

    base_dir = Path("browser_cache/")
    if not is_notebook_context:
        base_dir = base_dir / "looking_at"
    
    output_dir = base_dir / domain / url_path_slug
    artifacts = {}

    # --- CACHE OVERRIDE LOGIC ---
    if override_cache and output_dir.exists():
        if verbose:
            logger.info(f"🧹 override_cache is True. Clearing existing directory: {output_dir}")
        try:
            shutil.rmtree(output_dir)
        except Exception as e:
            logger.error(f"Failed to clear cache directory: {e}")

    # --- IDEMPOTENCY CHECK ---
    # Check if the primary artifact (hydrated_dom.html) already exists.
    dom_path = output_dir / "hydrated_dom.html"
    if dom_path.exists():
        if verbose:
            logger.info(f"✅ Using cached data from: {output_dir}")

        # Gather paths of existing artifacts
        for artifact_name in [
            "hydrated_dom.html", 
            "source.html", 
            "simple_source_html.html",
            "simple_hydrated_dom.html",
            "diff_boxes.txt",
            "diff_boxes.html",
            "diff_hierarchy.txt",
            "diff_hierarchy.html",
            "diff_simple_dom.txt",
            "diff_simple_dom.html",
            "screenshot.png", 
            "seo.md",
            "source_dom_layout_boxes.txt", 
            "source_dom_layout_boxes.html", 
            "source_dom_hierarchy.txt", 
            "source_dom_hierarchy.html", 
            "hydrated_dom_layout_boxes.txt", 
            "hydrated_dom_layout_boxes.html", 
            "hydrated_dom_hierarchy.txt", 
            "hydrated_dom_hierarchy.html", 
            "accessibility_tree.json", 
            "accessibility_tree_summary.txt"
        ]:
            artifact_path = output_dir / artifact_name
            if artifact_path.exists():
                 artifacts[Path(artifact_name).stem] = str(artifact_path)

        return {"success": True, "looking_at_files": artifacts, "cached": True}

    # --- Fuzzed Delay Logic (only runs if not cached) ---
    if delay_range and isinstance(delay_range, (tuple, list)) and len(delay_range) == 2:
        min_delay, max_delay = delay_range
        if isinstance(min_delay, (int, float)) and isinstance(max_delay, (int, float)) and min_delay <= max_delay:
            delay = random.uniform(min_delay, max_delay)
            if verbose:
                logger.info(f"⏳ Waiting for {delay:.3f} seconds before next request...")
            await asyncio.sleep(delay)
        else:
            logger.warning(f"⚠️ Invalid delay_range provided: {delay_range}. Must be a tuple of two numbers (min, max).")

    driver = None
    profile_path = None
    temp_profile = False

    # --- Find the browser executable path (Platform-Specific) ---
    effective_os = os.environ.get("EFFECTIVE_OS") # This is set by your flake.nix
    
    # Fallback in case script is run outside of the Nix shell
    if not effective_os:
        import platform
        effective_os = platform.system().lower()

    browser_path = None
    driver_path = None

    if effective_os == "linux":
        if verbose: logger.info("🐧 Linux platform detected. Looking for Nix-provided Chromium...")
        browser_path = shutil.which("chromium")
        driver_path = shutil.which("undetected-chromedriver")
        if not browser_path:
            browser_path = shutil.which("chromium-browser")
        
        if not browser_path:
            logger.error("❌ Could not find Nix-provided chromium or chromium-browser.")
            return {"success": False, "error": "Chromium executable not found in Nix environment."}
        if not driver_path:
            logger.error("❌ Could not find Nix-provided 'undetected-chromedriver'.")
            return {"success": False, "error": "undetected-chromedriver not found in Nix environment."}

    elif effective_os == "darwin":
        if verbose: logger.info("🍏 macOS platform detected. Looking for host-installed Google Chrome...")
        # On macOS, we rely on the user's host-installed Google Chrome.
        # undetected-chromedriver will use webdriver-manager to find/download the driver.
        browser_path = "/Applications/Google Chrome.app/Contents/MacOS/Google Chrome"
        driver_path = None # This tells uc to find/download the driver automatically

        if not Path(browser_path).exists():
            # Fallback for Chrome Canary
            browser_path_canary = "/Applications/Google Chrome Canary.app/Contents/MacOS/Google Chrome Canary"
            if Path(browser_path_canary).exists():
                browser_path = browser_path_canary
                if verbose: logger.info("  -> Google Chrome not found, using Google Chrome Canary.")
            else:
                logger.error(f"❌ Google Chrome not found at default path: {browser_path}")
                logger.error("   Please install Google Chrome on your Mac to continue.")
                return {"success": False, "error": "Google Chrome not found on macOS."}
        
        # Check if webdriver-manager is installed (it's a dependency of undetected-chromedriver)
        try:
            import webdriver_manager
        except ImportError:
            logger.error("❌ 'webdriver-manager' package not found.")
            logger.error("   Please add 'webdriver-manager' to requirements.txt and re-run 'nix develop'.")
            return {"success": False, "error": "webdriver-manager Python package missing."}
    
    else:
        logger.error(f"❌ Unsupported EFFECTIVE_OS: '{effective_os}'. Check flake.nix.")
        return {"success": False, "error": "Unsupported operating system."}

    if verbose: 
        logger.info(f"🔍 Using browser executable at: {browser_path}")
        if driver_path:
            logger.info(f"🔍 Using driver executable at: {driver_path}")
        else:
            logger.info(f"🔍 Using driver executable from webdriver-manager (uc default).")

    try:
        # Create directory only if we are actually scraping
        output_dir.mkdir(parents=True, exist_ok=True)
        if verbose: logger.info(f"💾 Saving new artifacts to: {output_dir}")

        options = uc.ChromeOptions()
        if headless:
            options.add_argument("--headless")
        options.add_argument("--no-sandbox")
        options.add_argument("--disable-dev-shm-usage")
        options.add_argument("--start-maximized")
        options.add_argument("--window-size=1920,1080")

        if persistent:
            profile_path = Path(f"data/uc_profiles/{profile_name}")
            profile_path.mkdir(parents=True, exist_ok=True)
            logger.info(f"🔒 Using persistent profile: {profile_path}")
        else:
            profile_path = tempfile.mkdtemp(prefix='pipulate_automation_')
            temp_profile = True
            logger.info(f"👻 Using temporary profile: {profile_path}")
        
        logger.info(f"🚀 Initializing undetected-chromedriver (Headless: {headless})...")
        try:
            driver = uc.Chrome(options=options, 
                               user_data_dir=str(profile_path), 
                               browser_executable_path=browser_path,
                               driver_executable_path=driver_path)
        except Exception as e:
            error_msg = str(e)
            if "Current browser version is" in error_msg:
                import re
                match = re.search(r'Current browser version is (\d+)', error_msg)
                if match:
                    fallback_version = int(match.group(1))
                    logger.warning(f"⚠️ Chrome version mismatch detected. Auto-healing with version_main={fallback_version}")
                    # UC consumes the options object. We must forge a fresh one.
                    fresh_options = uc.ChromeOptions()
                    if headless:
                        fresh_options.add_argument("--headless")
                    fresh_options.add_argument("--no-sandbox")
                    fresh_options.add_argument("--disable-dev-shm-usage")
                    fresh_options.add_argument("--start-maximized")
                    fresh_options.add_argument("--window-size=1920,1080")
                    
                    driver = uc.Chrome(options=fresh_options, 
                                       user_data_dir=str(profile_path), 
                                       browser_executable_path=browser_path,
                                       driver_executable_path=driver_path,
                                       version_main=fallback_version)
                else:
                    raise
            else:
                raise

        logger.info(f"Navigating to: {url}")
        driver.get(url)

        try:
            if verbose: logger.info("Waiting for security challenge to trigger a reload (Stage 1)...")
            initial_body = driver.find_element(By.TAG_NAME, 'body')
            WebDriverWait(driver, 20).until(EC.staleness_of(initial_body))
            if verbose: logger.success("✅ Page reload detected!")
            
            if verbose: logger.info("Waiting for main content to appear after reload (Stage 2)...")
            WebDriverWait(driver, 10).until(EC.presence_of_element_located((By.CSS_SELECTOR, "body")))
            if verbose: logger.success("✅ Main content located!")
        except Exception as e:
            if verbose: logger.info(f"Did not detect a page reload for security challenge. Proceeding anyway.")

        # --- Capture Core Artifacts ---
        dom_content = driver.execute_script("return document.documentElement.outerHTML;")
        dom_path = output_dir / "hydrated_dom.html"
        dom_path.write_text(dom_content, encoding='utf-8')
        artifacts['hydrated_dom'] = str(dom_path)
        
        # 1. Native Header & TRUE Raw Source Capture (The XHR Hack Lens)
        if verbose: logger.info("🌐 Extracting native headers and true raw source via XHR injection...")
        try:
            # We use a single XHR call to grab both the raw headers and the untouched responseText
            network_data_json = driver.execute_script("""
                var req = new XMLHttpRequest();
                req.open('GET', document.location.href, false);
                req.send(null);
                
                var headers = req.getAllResponseHeaders().toLowerCase();
                var arr = headers.trim().split(/[\\r\\n]+/);
                var headerMap = {};
                arr.forEach(function (line) {
                    var parts = line.split(': ');
                    var header = parts.shift();
                    if (header) headerMap[header] = parts.join(': ');
                });
                
                return JSON.stringify({
                    headers: headerMap,
                    raw_source: req.responseText
                });
            """)
            network_data = json.loads(network_data_json)
            actual_headers = network_data.get("headers", {})
            true_raw_source = network_data.get("raw_source", "")
            
            # Fallback to page_source if responseText is somehow empty
            if not true_raw_source.strip():
                true_raw_source = driver.page_source
                
        except Exception as e:
            if verbose: logger.warning(f"⚠️ Failed to extract native network data: {e}")
            actual_headers = {"error": "Could not extract headers without proxy"}
            true_raw_source = driver.page_source  # Fallback to the live DOM
            
        # Save True Raw Source
        source_html_path = output_dir / "source.html"
        source_html_path.write_text(true_raw_source, encoding='utf-8')
        artifacts['source_html'] = str(source_html_path)
        
        # Save Headers
        headers_data = {
            "url": url,
            "title": driver.title,
            "timestamp": datetime.now().isoformat(),
            "status": "success",
            "headers": actual_headers
        }
        headers_path = output_dir / "headers.json"
        headers_path.write_text(json.dumps(headers_data, indent=2), encoding='utf-8')
        artifacts['headers'] = str(headers_path)

        if take_screenshot:
            screenshot_path = output_dir / "screenshot.png"
            driver.save_screenshot(str(screenshot_path))
            artifacts['screenshot'] = str(screenshot_path)

        # 1. Native Header Capture via Performance API (The XHR Hack Lens)
        if verbose: logger.info("🌐 Extracting headers via XHR injection...")
        try:
            headers_json = driver.execute_script("""
                var req = new XMLHttpRequest();
                req.open('GET', document.location, false);
                req.send(null);
                var headers = req.getAllResponseHeaders().toLowerCase();
                var arr = headers.trim().split(/[\\r\\n]+/);
                var headerMap = {};
                arr.forEach(function (line) {
                    var parts = line.split(': ');
                    var header = parts.shift();
                    if (header) headerMap[header] = parts.join(': ');
                });
                return JSON.stringify(headerMap);
            """)
            actual_headers = json.loads(headers_json)
        except Exception as e:
            if verbose: logger.warning(f"⚠️ Failed to extract headers: {e}")
            actual_headers = {"error": "Could not extract headers without proxy"}
        
        headers_data = {
            "url": url,
            "title": driver.title,
            "timestamp": datetime.now().isoformat(),
            "status": "success",
            "headers": actual_headers
        }
        headers_path = output_dir / "headers.json"
        headers_path.write_text(json.dumps(headers_data, indent=2), encoding='utf-8')
        artifacts['headers'] = str(headers_path)

        # 2. Create LLM-Optimized Simplified DOMs (The Symmetrical Lens)
        if verbose: logger.info("🧠 Creating LLM-optimized simplified DOMs (Symmetrical Lens)...")
        
        simple_source_content = _simplify_html_for_llm(true_raw_source, driver.title)
        simple_source_path = output_dir / "simple_source_html.html"
        simple_source_path.write_text(simple_source_content, encoding='utf-8')
        artifacts['simple_source'] = str(simple_source_path)

        simple_hydrated_content = _simplify_html_for_llm(dom_content, driver.title)
        simple_hydrated_path = output_dir / "simple_hydrated_dom.html"
        simple_hydrated_path.write_text(simple_hydrated_content, encoding='utf-8')
        artifacts['simple_hydrated'] = str(simple_hydrated_path)

        # --- Generate Accessibility Tree Artifact ---
        if verbose: logger.info("🌲 Extracting accessibility tree...")
        try:
            driver.execute_cdp_cmd("Accessibility.enable", {})
            ax_tree_result = driver.execute_cdp_cmd("Accessibility.getFullAXTree", {})
            ax_tree = ax_tree_result.get("nodes", [])
            ax_tree_path = output_dir / "accessibility_tree.json"
            ax_tree_path.write_text(json.dumps({"success": True, "node_count": len(ax_tree), "accessibility_tree": ax_tree}, indent=2), encoding='utf-8')
            artifacts['accessibility_tree'] = str(ax_tree_path)

            summary_result = await dom_tools.summarize_accessibility_tree({"file_path": str(ax_tree_path)})
            if summary_result.get("success"):
                summary_path = output_dir / "accessibility_tree_summary.txt"
                summary_path.write_text(summary_result["output"], encoding='utf-8')
                artifacts['accessibility_tree_summary'] = str(summary_path)
        except Exception as ax_error:
            logger.warning(f"⚠️ Could not extract accessibility tree: {ax_error}")

        # --- Generate LLM Optics (Subprocess Bulkhead) ---
        if verbose: logger.info("👁️‍🗨️ Running LLM Optics Engine (Subprocess Bulkhead)...")
        # We pass the output_dir, not the dom_path
        optics_result = await generate_optics_subprocess(str(output_dir))
        
        if optics_result.get('success'):
            if verbose: logger.success("✅ LLM Optics Engine completed successfully.")
            # Append new optical artifacts to the result dictionary
            for optic_key, filename in [
                ('seo_md', 'seo.md'),
                ('source_hierarchy_txt', 'source_dom_hierarchy.txt'),
                ('source_hierarchy_html', 'source_dom_hierarchy.html'),
                ('source_boxes_txt', 'source_dom_layout_boxes.txt'),
                ('source_boxes_html', 'source_dom_layout_boxes.html'),
                ('hydrated_hierarchy_txt', 'hydrated_dom_hierarchy.txt'),
                ('hydrated_hierarchy_html', 'hydrated_dom_hierarchy.html'),
                ('hydrated_boxes_txt', 'hydrated_dom_layout_boxes.txt'),
                ('hydrated_boxes_html', 'hydrated_dom_layout_boxes.html'),
                ('diff_hierarchy_txt', 'diff_hierarchy.txt'),
                ('diff_hierarchy_html', 'diff_hierarchy.html'),
                ('diff_boxes_txt', 'diff_boxes.txt'),
                ('diff_boxes_html', 'diff_boxes.html'),
                ('diff_simple_txt', 'diff_simple_dom.txt'),
                ('diff_simple_html', 'diff_simple_dom.html')
            ]:
                optic_path = output_dir / filename
                if optic_path.exists():
                    artifacts[optic_key] = str(optic_path)
        else:
            if verbose: logger.warning(f"⚠️ LLM Optics Engine partially failed: {optics_result.get('error')}")

        logger.success(f"✅ Scrape successful for {url}")
        return {"success": True, "looking_at_files": artifacts, "cached": False}

    finally:
        if driver:
            try:
                driver.quit()
                if verbose: logger.info("Browser closed.")
            except Exception as e:
                logger.warning(f"Error while quitting browser: {e}")

        if temp_profile and profile_path and os.path.exists(profile_path):
            try:
                # Add ignore_errors=True to prevent ghost processes from crashing the cleanup
                shutil.rmtree(profile_path, ignore_errors=True)
                if verbose: logger.info(f"Cleaned up temporary profile: {profile_path}")
            except Exception as e:
                logger.warning(f"Could not completely remove temp profile (this is normal): {e}")
