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

from bs4 import BeautifulSoup
import httpx
import trafilatura
from lxml_html_clean import Cleaner
import markdownify
from duckduckgo_search import DDGS
from patchright.async_api import async_playwright, Browser, BrowserContext, Page

from mcp.server.fastmcp import FastMCP
try:
    from mcp.server.transport_security import TransportSecuritySettings
    security_settings = TransportSecuritySettings(
        enable_dns_rebinding_protection=False,
        allowed_hosts=["*"]
    )
except ImportError:
    security_settings = None

from starlette.applications import Starlette
from starlette.routing import Route, Mount
from starlette.responses import JSONResponse
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
# ==========================================
# Initialize FastMCP with transport_security disabled for public cloud proxying
if security_settings:
    mcp = FastMCP("WebScraper", host=HOST, transport_security=security_settings)
else:
    mcp = FastMCP("WebScraper", host=HOST)

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

@mcp.tool(
    name="search_and_scrape",
    description="Searches DuckDuckGo and concurrently extracts content from top results."
)
async def search_and_scrape(query: str, max_results: int = 3) -> dict:
    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(
    name="take_screenshot",
    description="Takes a full-page or viewport screenshot of a webpage as Base64 PNG."
)
async def take_screenshot(url: str, full_page: bool = True, wait_seconds: float = 1.0) -> dict:
    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(
    name="interact_page",
    description="Executes a list of browser actions (click, type, press, wait, evaluate)."
)
async def interact_page(url: str, actions: list[dict[str, Any]], extract_markdown: bool = True) -> dict:
    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(
    name="extract_structured_data",
    description="Parses structured JSON-LD schemas and arbitrary CSS selector targets."
)
async def extract_structured_data(url: str, css_selectors: dict[str, str] | None = None) -> dict:
    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 & Tool Inspection
# ==========================================
async def health_check(request):
    return JSONResponse({
        "status": "healthy",
        "mcp_sse_endpoint": "/sse",
        "tools_endpoint": "/tools"
    })

async def list_tools_endpoint(request):
    """Direct HTTP endpoint exposing all registered MCP tools and schemas."""
    tools_list = []
    for tool_name, tool_obj in mcp._tool_manager._tools.items():
        tools_list.append({
            "name": tool_name,
            "description": getattr(tool_obj, "description", "") or "",
            "parameters": getattr(tool_obj, "parameters", {}) or {}
        })
    return JSONResponse({"tools": tools_list})

starlette_app = Starlette(
    routes=[
        Route("/", endpoint=health_check),
        Route("/healthz", endpoint=health_check),
        Route("/tools", endpoint=list_tools_endpoint),
        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")