"""High-Performance Stealth Web Scraper & Browser Automation MCP Server

Directly serves SSE & HTTP endpoints via Uvicorn on 0.0.0.0:7860.
"""

import asyncio
import base64
import json
import logging
import os
import sys
from contextlib import asynccontextmanager
from typing import Any, AsyncGenerator

from bs4 import BeautifulSoup
from duckduckgo_search import DDGS
import httpx
from lxml_html_clean import Cleaner
import markdownify
from mcp.server.fastmcp import FastMCP
from patchright.async_api import Browser, BrowserContext, Page, async_playwright
from starlette.applications import Starlette
from starlette.responses import JSONResponse
from starlette.routing import Mount, Route
import trafilatura
import uvicorn

# ==========================================
# 1. Logging & Configuration
# ==========================================
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
    handlers=[logging.StreamHandler(sys.stdout)],
)
logger = logging.getLogger("mcp_server")

PORT = int(os.getenv("PORT", 7860))
HOST = "0.0.0.0"
MAX_CONCURRENT_BROWSERS = 4
MAX_RECYCLE_REQUESTS = 100
MAX_CONTENT_CHARS = 100_000
HTTP_TIMEOUT_SEC = 12.0
BROWSER_TIMEOUT_MS = 30000

# ==========================================
# 2. Browser Pool
# ==========================================


class StealthBrowserPool:

    def __init__(self):
        self._playwright = None
        self._browser: Browser | None = None
        self._semaphore = asyncio.Semaphore(MAX_CONCURRENT_BROWSERS)
        self._request_counter = 0
        self._lock = asyncio.Lock()

    async def get_browser(self) -> Browser:
        async with self._lock:
            if not self._playwright:
                self._playwright = await async_playwright().start()
            if not self._browser or not self._browser.is_connected():
                logger.info("Initializing Chromium Stealth instance...")
                self._browser = await self._playwright.chromium.launch(
                    headless=True,
                    args=[
                        "--no-sandbox",
                        "--disable-setuid-sandbox",
                        "--disable-dev-shm-usage",
                        "--disable-gpu",
                        "--disable-blink-features=AutomationControlled",
                        "--no-first-run",
                        "--window-size=1920,1080",
                    ],
                )
                self._request_counter = 0
            return self._browser

    async def _check_recycle(self):
        async with self._lock:
            self._request_counter += 1
            if self._request_counter >= MAX_RECYCLE_REQUESTS:
                if self._browser:
                    try:
                        await self._browser.close()
                    except Exception:
                        pass
                self._browser = None

    @asynccontextmanager
    async def get_page(self) -> AsyncGenerator[Page, None]:
        await self._semaphore.acquire()
        context: BrowserContext | None = None
        page: Page | None = None
        try:
            browser = await self.get_browser()
            context = await browser.new_context(
                viewport={"width": 1920, "height": 1080},
                user_agent=(
                    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
                    "AppleWebKit/537.36 (KHTML, like Gecko) "
                    "Chrome/131.0.0.0 Safari/537.36"
                ),
                locale="en-US",
                timezone_id="America/New_York",
            )
            await context.add_init_script("""
                Object.defineProperty(navigator, 'webdriver', { get: () => undefined });
                Object.defineProperty(navigator, 'plugins', { get: () => [1, 2, 3, 4, 5] });
                Object.defineProperty(navigator, 'languages', { get: () => ['en-US', 'en'] });
                window.chrome = { runtime: {} };
            """)
            page = await context.new_page()
            page.set_default_navigation_timeout(BROWSER_TIMEOUT_MS)
            page.set_default_timeout(BROWSER_TIMEOUT_MS)
            yield page
        finally:
            if page:
                try:
                    await page.close()
                except Exception:
                    pass
            if context:
                try:
                    await context.close()
                except Exception:
                    pass
            self._semaphore.release()
            await self._check_recycle()


browser_pool = StealthBrowserPool()

# ==========================================
# 3. Content Extraction Pipeline
# ==========================================


def clean_html(raw_html: str) -> str:
    cleaner = Cleaner(
        scripts=True,
        javascript=True,
        comments=True,
        style=True,
        links=False,
        meta=False,
        page_structure=False,
        safe_attrs_only=False,
    )
    return cleaner.clean_html(raw_html)


def extract_content(html: str, url: str = "") -> dict[str, Any]:
    extracted = trafilatura.extract(
        html,
        url=url,
        include_links=True,
        include_images=True,
        output_format="markdown",
    )
    metadata = trafilatura.extract_metadata(html) or {}
    meta_dict = {
        "title": getattr(metadata, "title", "") or "",
        "author": getattr(metadata, "author", "") or "",
        "date": getattr(metadata, "date", "") or "",
        "description": getattr(metadata, "description", "") or "",
    }

    if not extracted or len(extracted.strip()) < 100:
        cleaned = clean_html(html)
        soup = BeautifulSoup(cleaned, "html.parser")
        for tag in soup(
            ["nav", "footer", "aside", "header", "script", "style", "noscript"]
        ):
            tag.decompose()
        if not meta_dict["title"] and soup.title:
            meta_dict["title"] = soup.title.string or ""
        extracted = markdownify.markdownify(
            str(soup), heading_style="ATX", strip=["svg"]
        )

    if len(extracted) > MAX_CONTENT_CHARS:
        extracted = (
            extracted[:MAX_CONTENT_CHARS]
            + f"\n\n... [Truncated: reached {MAX_CONTENT_CHARS} limit]"
        )

    return {"content": extracted.strip(), "metadata": meta_dict}


def extract_json_ld(html: str) -> list[dict]:
    soup = BeautifulSoup(html, "html.parser")
    schemas = []
    for script in soup.find_all("script", type="application/ld+json"):
        try:
            if script.string:
                data = json.loads(script.string.strip())
                schemas.append(data)
        except Exception:
            continue
    return schemas


# ==========================================
# 4. Hybrid Scraping Core
# ==========================================

FAST_HEADERS = {
    "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36",
    "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8",
    "Accept-Language": "en-US,en;q=0.9",
}


async def dismiss_popups(page: Page):
    selectors = [
        'button[id*="cookie" i]',
        'button[class*="cookie" i]',
        'button[aria-label*="accept" i]',
        'button:has-text("Accept all")',
        'button:has-text("Accept")',
        'button:has-text("I agree")',
        'button:has-text("Got it")',
    ]
    for sel in selectors:
        try:
            elem = await page.query_selector(sel)
            if elem and await elem.is_visible():
                await elem.click(timeout=800)
                await asyncio.sleep(0.2)
                break
        except Exception:
            continue


async def scrape_pipeline(
    url: str,
    force_browser: bool = False,
    auto_scroll: bool = True,
    wait_for_selector: str | None = None,
) -> dict:
    if not force_browser:
        try:
            async with httpx.AsyncClient(
                http2=True,
                timeout=HTTP_TIMEOUT_SEC,
                follow_redirects=True,
                headers=FAST_HEADERS,
            ) as client:
                resp = await client.get(url)
                if resp.status_code == 200:
                    html = resp.text
                    if not any(
                        s in html.lower()
                        for s in [
                            "cf-challenge",
                            "ray-id",
                            "just a moment...",
                        ]
                    ):
                        parsed = extract_content(html, url=url)
                        if len(parsed["content"]) >= 150:
                            return {
                                "url": str(resp.url),
                                "status": 200,
                                "engine": "fast-http",
                                "content": parsed["content"],
                                "metadata": parsed["metadata"],
                            }
        except Exception:
            pass

    async with browser_pool.get_page() as page:
        await page.goto(url, wait_until="domcontentloaded")
        try:
            cf_frame = await page.query_selector(
                "iframe[src*='challenges.cloudflare.com']"
            )
            if cf_frame:
                box = await cf_frame.bounding_box()
                if box:
                    await page.mouse.click(
                        box["x"] + box["width"] / 2,
                        box["y"] + box["height"] / 2,
                    )
                    await asyncio.sleep(2.0)
        except Exception:
            pass

        await dismiss_popups(page)

        if wait_for_selector:
            try:
                await page.wait_for_selector(
                    wait_for_selector, state="visible", timeout=8000
                )
            except Exception:
                pass

        if auto_scroll:
            for _ in range(4):
                await page.mouse.wheel(0, 1000)
                await asyncio.sleep(0.3)

        await asyncio.sleep(0.5)
        html = await page.content()
        final_url = page.url
        title = await page.title()

    parsed = extract_content(html, url=final_url)
    if not parsed["metadata"].get("title"):
        parsed["metadata"]["title"] = title

    return {
        "url": final_url,
        "status": 200,
        "engine": "patchright-stealth",
        "content": parsed["content"],
        "metadata": parsed["metadata"],
    }


# ==========================================
# 5. MCP Server & Tool Definitions
# ==========================================
mcp = FastMCP("WebScraper")


@mcp.tool()
async def scrape_url(
    url: str,
    force_browser: bool = False,
    auto_scroll: bool = True,
    wait_for_selector: str | None = None,
) -> dict:
    """Scrapes cleaned Markdown content and metadata from any URL."""
    return await scrape_pipeline(
        url, force_browser, auto_scroll, wait_for_selector
    )


@mcp.tool()
async def search_and_scrape(query: str, max_results: int = 3) -> dict:
    """Searches DuckDuckGo and concurrently extracts content from top results."""
    max_results = min(max(1, max_results), 5)
    loop = asyncio.get_running_loop()

    def _search():
        with DDGS() as ddgs:
            return list(ddgs.text(query, max_results=max_results))

    results = await loop.run_in_executor(None, _search)
    if not results:
        return {"query": query, "results": []}

    tasks = [
        scrape_pipeline(res["href"], force_browser=False, auto_scroll=False)
        for res in results
    ]
    scraped_payloads = await asyncio.gather(*tasks, return_exceptions=True)

    enriched = []
    for meta, payload in zip(results, scraped_payloads):
        if isinstance(payload, dict):
            enriched.append(
                {
                    "title": meta.get("title", ""),
                    "url": meta.get("href", ""),
                    "snippet": meta.get("body", ""),
                    "content": payload.get("content", ""),
                }
            )
        else:
            enriched.append(
                {
                    "title": meta.get("title", ""),
                    "url": meta.get("href", ""),
                    "snippet": meta.get("body", ""),
                    "content": f"[Error: {str(payload)}]",
                }
            )
    return {"query": query, "results": enriched}


@mcp.tool()
async def take_screenshot(
    url: str, full_page: bool = True, wait_seconds: float = 1.0
) -> dict:
    """Takes a full-page or viewport screenshot of a webpage."""
    async with browser_pool.get_page() as page:
        await page.goto(url, wait_until="networkidle")
        if wait_seconds > 0:
            await asyncio.sleep(min(wait_seconds, 10.0))
        screenshot_bytes = await page.screenshot(
            full_page=full_page, type="png"
        )
        b64 = base64.b64encode(screenshot_bytes).decode("utf-8")
        return {
            "url": page.url,
            "title": await page.title(),
            "format": "image/png;base64",
            "base64_image": b64,
        }


@mcp.tool()
async def interact_page(
    url: str, actions: list[dict[str, Any]], extract_markdown: bool = True
) -> dict:
    """Executes a list of browser actions (click, type, press, wait, evaluate)."""
    async with browser_pool.get_page() as page:
        await page.goto(url, wait_until="domcontentloaded")
        logs = []

        for idx, act in enumerate(actions):
            act_type = act.get("type", "").lower()
            try:
                if act_type == "click":
                    sel = act["selector"]
                    await page.click(sel, timeout=6000)
                    logs.append(f"[{idx}] Clicked '{sel}'")
                elif act_type == "type":
                    sel = act["selector"]
                    txt = act["text"]
                    await page.fill(sel, txt, timeout=6000)
                    logs.append(f"[{idx}] Typed into '{sel}'")
                elif act_type == "press":
                    k = act["key"]
                    await page.keyboard.press(k)
                    logs.append(f"[{idx}] Pressed '{k}'")
                elif act_type == "wait":
                    s = min(float(act.get("seconds", 1.0)), 15.0)
                    await asyncio.sleep(s)
                    logs.append(f"[{idx}] Waited {s}s")
                elif act_type == "evaluate":
                    res = await page.evaluate(act["script"])
                    logs.append(f"[{idx}] Evaluated script -> {res}")
            except Exception as e:
                logs.append(f"[{idx}] Failed: {str(e)}")

        html = await page.content()
        final_url = page.url

    res = {"final_url": final_url, "logs": logs}
    if extract_markdown:
        parsed = extract_content(html, url=final_url)
        res["content"] = parsed["content"]
        res["metadata"] = parsed["metadata"]
    return res


@mcp.tool()
async def extract_structured_data(
    url: str, css_selectors: dict[str, str] | None = None
) -> dict:
    """Parses structured JSON-LD schemas and arbitrary CSS selector targets."""
    async with browser_pool.get_page() as page:
        await page.goto(url, wait_until="domcontentloaded")
        html = await page.content()

    soup = BeautifulSoup(html, "html.parser")
    custom_data = {}
    if css_selectors:
        for k, sel in css_selectors.items():
            elements = soup.select(sel)
            custom_data[k] = [el.get_text(strip=True) for el in elements]

    return {
        "url": url,
        "custom_fields": custom_data,
        "json_ld_schemas": extract_json_ld(html),
    }


# ==========================================
# 6. Starlette SSE Router & Server Runner
# ==========================================


async def health_check(request):
    return JSONResponse(
        {
            "status": "healthy",
            "mcp_sse_endpoint": "/sse",
            "transport": "Server-Sent Events (SSE)",
        }
    )


# Mount the MCP SSE application with health endpoints
starlette_app = Starlette(
    routes=[
        Route("/", endpoint=health_check),
        Route("/healthz", endpoint=health_check),
        Mount("/", app=mcp.sse_app()),
    ]
)

if __name__ == "__main__":
    logger.info(f"Starting MCP Server on http://{HOST}:{PORT} ...")
    uvicorn.run(starlette_app, host=HOST, port=PORT, log_level="info")