import asyncio
import datetime
import hashlib
import json
import logging
import os
import re
import urllib.parse
from pyquery import PyQuery as pq
from playwright.async_api import async_playwright

logger = logging.getLogger(__name__)

UNLIMITED = 2**31


class PageScraper:
    def __init__(
        self,
        url,
        output_dir="data",
        max_depth=5,
        max_pages=None,
        max_assets=None,
        max_parallel=1,
        timeout=60000,
        page_timeout=120,
    ):
        self.url = url
        self.output_dir = output_dir
        self.max_depth = max_depth
        self.max_pages = max_pages
        self.max_assets = max_assets
        self.max_parallel = max_parallel
        self.timeout = timeout
        self.page_timeout = page_timeout
        self.page_count = 0
        self.asset_count = 0
        self.enqueued = set()
        self._pages = {}
        self._queue = asyncio.Queue()
        self._busy = 0
        self._origin_base = ".".join(
            urllib.parse.urlparse(url).netloc.rstrip(".").split(".")[-2:]
        )

    def _url_dir(self, url):
        parsed = urllib.parse.urlparse(url)
        path = parsed.path.strip("/")
        return (
            os.path.join(self.output_dir, parsed.netloc, path)
            if path
            else os.path.join(self.output_dir, parsed.netloc)
        )

    def _read_meta(self, url):
        meta_path = os.path.join(self._url_dir(url), "initial.json")
        if not os.path.isfile(meta_path):
            return None
        try:
            with open(meta_path) as f:
                return json.load(f)
        except (json.JSONDecodeError, FileNotFoundError):
            return None

    def _content_changed(self, url, content):
        new_hash = hashlib.sha256(content.encode()).hexdigest()
        previous = self._read_meta(url)
        if previous:
            return previous.get("content", {}).get("initial_sha256") != new_hash
        return True

    def _is_asset(self, resource_type):
        return resource_type in (
            "stylesheet",
            "image",
            "media",
            "font",
            "script",
            "fetch",
            "xhr",
            "other",
        )

    def _asset_path(self, url, content_type):
        parsed = urllib.parse.urlparse(url)
        netloc = parsed.netloc or parsed.path.split("/")[0]
        path = parsed.path.strip("/") if parsed.path else netloc
        if path.startswith(netloc):
            path = "/" + path[len(netloc) :]
        ext_map = {
            "image/png": "png",
            "image/jpeg": "jpg",
            "image/webp": "webp",
            "image/gif": "gif",
            "image/svg+xml": "svg",
            "image/x-icon": "ico",
            "application/javascript": "js",
            "text/javascript": "js",
            "text/css": "css",
            "font/woff2": "woff2",
            "font/woff": "woff",
            "application/json": "json",
            "application/pdf": "pdf",
        }
        _, ext = os.path.splitext(path)
        if ext:
            return os.path.join(self.output_dir, netloc, path)
        mime_ext = ext_map.get(content_type.split(";")[0].strip().lower())
        if mime_ext:
            return os.path.join(self.output_dir, netloc, path, f"index.{mime_ext}")
        return os.path.join(self.output_dir, netloc, path, "index.dat")

    async def _on_response(self, response, current_url_ref):
        try:
            url = response.url
            if not self._is_asset(response.request.resource_type):
                return
            if self.max_assets is not None and self.asset_count >= self.max_assets:
                return
            content_type = response.headers.get("content-type", "")
            asset_path = self._asset_path(url, content_type)
            if os.path.isfile(asset_path):
                return
            self.asset_count += 1
            try:
                body = await response.body()
            except Exception:
                return
            try:
                os.makedirs(os.path.dirname(asset_path), exist_ok=True)
            except OSError:
                return
            with open(asset_path, "wb") as f:
                f.write(body)
            asset_hash = hashlib.sha256(body).hexdigest()
            meta = {
                "url": url,
                "request": {
                    "method": response.request.method,
                    "headers": dict(response.request.headers),
                },
                "response": {
                    "status": response.status,
                    "headers": dict(response.headers),
                },
                "content": {"sha256": asset_hash},
                "scrape": {
                    "timestamp": datetime.datetime.now(
                        datetime.timezone.utc
                    ).isoformat(),
                },
            }
            with open(asset_path + ".json", "w", encoding="utf-8") as f:
                json.dump(meta, f, indent=2)
            page_url = current_url_ref.get()
            if page_url and page_url in self._pages:
                self._pages[page_url]["assets"].append(
                    {
                        "url": url,
                    }
                )
        except Exception:
            pass

    def _extract_article_links(self, doc, page_url):
        links = set()
        skip_patterns = [
            re.compile(
                r"^/(about|contact|privacy|terms|ad|subscribe|newsletter|"
                r"login|signup|careers|feedback|site|help|support|faq|"
                r"sitemap|robots|feed|rss|api|static|assets|img|images|"
                r"css|js|scripts|fonts|cdn|media|video|audio)$"
            ),
            re.compile(
                r"\.(png|jpg|jpeg|gif|svg|ico|css|js|woff|woff2|ttf|pdf|xml|json|html)$"
            ),
            re.compile(r"^#"),
            re.compile(r"^mailto:"),
            re.compile(r"^javascript:"),
            re.compile(r"^/$"),
        ]
        for a in doc("a").items():
            href = a.attr("href") or ""
            resolved = urllib.parse.urljoin(page_url, href)
            parsed = urllib.parse.urlparse(resolved)
            parsed_base = ".".join(parsed.netloc.rstrip(".").split(".")[-2:])
            if parsed_base != self._origin_base:
                continue
            path = parsed.path.strip("/")
            if not path:
                continue
            if any(s.search(parsed.path) for s in skip_patterns):
                continue
            segments = [s for s in path.split("/") if s]
            if len(segments) >= 2:
                links.add(resolved)
        links.discard(self.url)
        for seen in self.enqueued:
            links.discard(seen)
        return links

    async def _process_page(self, page, page_url, depth):
        headers = {}
        previous = self._read_meta(page_url)
        if previous:
            etag = previous.get("response", {}).get("etag")
            if etag:
                headers["if-none-match"] = etag
            timestamp = previous.get("scrape", {}).get("timestamp")
            if timestamp:
                headers["if-modified-since"] = timestamp
        try:
            await page.set_extra_http_headers(headers)
            resp = await page.goto(page_url, wait_until="domcontentloaded")
            await page.set_extra_http_headers({})
        except Exception as e:
            logger.error(f"Failed to navigate to {page_url}: {e}")
            return
        page_url = resp.url
        if resp.request.redirected_from:
            redirect_url = resp.request.redirected_from.url
            page_dir = self._url_dir(redirect_url)
            meta_path = os.path.join(page_dir, "initial.json")
            if not os.path.exists(meta_path):
                os.makedirs(page_dir, exist_ok=True)
                prev_resp = resp.request.redirected_from.existing_response
                meta = {
                    "request": {
                        "method": resp.request.method,
                        "url": redirect_url,
                        "final_url": page_url,
                        "redirect_chain": [],
                        "headers": dict(resp.request.headers),
                    },
                    "response": {
                        "status": prev_resp.status if prev_resp else None,
                        "headers": dict(resp.headers),
                    },
                    "content": {"sha256": ""},
                    "scrape": {
                        "timestamp": datetime.datetime.now(
                            datetime.timezone.utc
                        ).isoformat(),
                    },
                }
                with open(meta_path, "w", encoding="utf-8") as f:
                    json.dump(meta, f, indent=2)
        if resp.status == 304:
            page_dir = self._url_dir(page_url)
            initial_path = os.path.join(page_dir, "initial.html")
            if os.path.isfile(initial_path):
                with open(initial_path, encoding="utf-8") as f:
                    local_html = f.read()
                await self._enqueue_links(local_html, page_url, depth)
            return
        initial_html = await page.content()
        content_hash = hashlib.sha256(initial_html.encode()).hexdigest()
        changed = self._content_changed(page_url, initial_html)

        if changed:
            at_limit = self.max_pages is not None and self.page_count >= self.max_pages
            if at_limit:
                return
            self.page_count += 1
            logger.info(f"Scraping page {self.page_count}: {page_url}")
            page_dir = self._url_dir(page_url)
            os.makedirs(page_dir, exist_ok=True)
            with open(
                os.path.join(page_dir, "initial.html"), "w", encoding="utf-8"
            ) as f:
                f.write(initial_html)

            await self._wait_for_hydrated(page)
            await self._scroll_page(page)
            final_html = await page.content()
            final_hash = hashlib.sha256(final_html.encode()).hexdigest()
            with open(os.path.join(page_dir, "index.html"), "w", encoding="utf-8") as f:
                f.write(final_html)
            page_data = {"assets": []}
            self._pages[page_url] = page_data
            meta = {
                "request": {
                    "method": resp.request.method,
                    "url": page_url,
                    "final_url": page_url,
                    "redirect_chain": [],
                    "headers": dict(resp.request.headers),
                },
                "response": {
                    "status": resp.status,
                    "headers": dict(resp.headers),
                    "etag": resp.headers.get("etag"),
                },
                "content": {
                    "sha256": final_hash,
                    "initial_sha256": content_hash,
                },
                "scrape": {
                    "timestamp": datetime.datetime.now(
                        datetime.timezone.utc
                    ).isoformat(),
                },
            }
            with open(
                os.path.join(page_dir, "initial.json"), "w", encoding="utf-8"
            ) as f:
                json.dump(meta, f, indent=2)
            await self._enqueue_links(initial_html, page_url, depth)
            await self._enqueue_links(final_html, page_url, depth)
        else:
            logger.info(f"Page unchanged, skipping: {page_url}")
            page_dir = self._url_dir(page_url)
            initial_path = os.path.join(page_dir, "initial.html")
            if os.path.isfile(initial_path):
                with open(initial_path, encoding="utf-8") as f:
                    local_html = f.read()
                await self._enqueue_links(local_html, page_url, depth)

    async def _wait_for_hydrated(self, page):
        prev_size = -1
        stable = 0
        for _ in range(10):
            html = await page.content()
            curr_size = len(html)
            if abs(curr_size - prev_size) < 500:
                stable += 1
            else:
                stable = 0
            prev_size = curr_size
            if stable >= 2:
                break
            await asyncio.sleep(1)

    async def _scroll_page(self, page):
        """Scroll to bottom and back to trigger lazy loading."""
        await page.evaluate("window.scrollTo(0, document.body.scrollHeight)")
        await asyncio.sleep(1)
        await page.evaluate("window.scrollTo(0, 0)")

    async def _enqueue_links(self, html, page_url, depth):
        if self.max_pages is not None and self.page_count >= self.max_pages:
            return
        if depth >= self.max_depth:
            return
        doc = pq(html)
        remaining = (
            (self.max_pages - self.page_count)
            if self.max_pages is not None
            else UNLIMITED
        )
        for link_url in self._extract_article_links(doc, page_url):
            if remaining <= 0:
                break
            if link_url not in self.enqueued:
                self.enqueued.add(link_url)
                remaining -= 1
                await self._queue.put((link_url, depth + 1))

    async def _worker(self, worker_id, pages):
        page = pages[worker_id]
        current_url_ref = {"url": None}
        page.on("response", lambda resp, r=current_url_ref: self._on_response(resp, r))
        while True:
            if self.max_pages is not None and self.page_count >= self.max_pages:
                break
            item = await self._queue.get()
            if item is None:
                break
            url, depth = item
            if depth >= self.max_depth:
                self._queue.task_done()
                continue
            self._busy += 1
            try:
                if current_url_ref["url"] and current_url_ref["url"] in self._pages:
                    self._pages[current_url_ref["url"]]["assets"] = self._pages[
                        current_url_ref["url"]
                    ].get("assets", [])
                current_url_ref["url"] = url
                await asyncio.wait_for(
                    self._process_page(page, url, depth),
                    timeout=self.page_timeout,
                )
                await asyncio.sleep(0.3)
            except asyncio.TimeoutError:
                logger.warning(f"Timeout processing {url}")
            except Exception as e:
                logger.error(f"Error processing {url}: {e}")
            finally:
                self._busy -= 1
                self._queue.task_done()

    async def run(self):
        async with async_playwright() as p:
            browser = await p.chromium.launch(
                headless=True,
                args=["--disable-dev-shm-usage"],
                executable_path="/usr/bin/chromium-browser",
            )
            context = await browser.new_context(
                user_agent=(
                    "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 "
                    "(KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36"
                ),
                viewport={"width": 1920, "height": 1080},
            )
            await context.add_init_script(
                "Object.defineProperty(navigator, 'webdriver', {get: () => undefined})"
            )
            pages = []
            for i in range(self.max_parallel):
                pg = await context.new_page()
                pg.set_default_timeout(self.timeout)
                pages.append(pg)
            logger.info(
                f"Scraping {self.url} (depth={self.max_depth}, parallel={self.max_parallel})"
            )
            await self._queue.put((self.url, 0))
            workers = [
                asyncio.create_task(self._worker(i, pages))
                for i in range(self.max_parallel)
            ]
            watchdog = 0
            try:
                while True:
                    await asyncio.sleep(0.5)
                    watchdog += 1
                    if watchdog >= 300:
                        logger.warning("Watchdog: queue not draining")
                        watchdog = 0
                    if self._queue.empty() and self._busy == 0:
                        for _ in range(self.max_parallel):
                            await self._queue.put(None)
                        break
                    if self.max_pages is not None and self.page_count >= self.max_pages:
                        for _ in range(self.max_parallel):
                            await self._queue.put(None)
                        break
                await asyncio.gather(*workers)
            except Exception:
                logger.warning("Browser crashed, stopping scrape")
                for _ in range(self.max_parallel):
                    await self._queue.put(None)
                await asyncio.gather(*workers, return_exceptions=True)
            logger.info(f"Done: {self.page_count} pages, {self.asset_count} assets")
            try:
                await browser.close()
            except Exception:
                pass