# -*- coding: utf-8 -*-
"""
通用爬虫引擎 (Playwright + stealth + 系统 Chrome)
- 批量模式: 逐条爬取网址列表
- 自动模式: 从起始网址按匹配规则自动发现链接并 BFS 爬取
- 支持: 随机延迟 / 重试 / 图片下载 / 暂停恢复 / 终止 / 配置热更新 / cookie 复用
"""
import json
import os
import random
import re
import threading
import time
import urllib.parse
from datetime import datetime

from playwright.sync_api import sync_playwright
from playwright_stealth import Stealth

import notify
import store

HERE = os.path.dirname(os.path.abspath(__file__))
DATA_DIR = os.path.join(HERE, "data")
CHROME = "/usr/bin/google-chrome"
IMG_EXTS = (".jpg", ".jpeg", ".png", ".gif", ".webp", ".bmp")
DEFAULT_UA = ("Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 "
              "(KHTML, like Gecko) Chrome/125.0.0.0 Safari/537.36")

_CHALLENGE_MARKS = [
    "access denied", "403 forbidden", "just a moment", "attention required",
    "captcha", "bot check", "cf-challenge", "verify you are human",
    "checking your browser", "enable javascript and cookies",
]


def is_challenge_page(title, html):
    low = html.lower()
    t = (title or "").lower()
    for mark in _CHALLENGE_MARKS:
        if mark in t or mark in low:
            return True
    # 极小页面 + 无正文结构 -> 疑似验证壳
    if len(html) < 5000 and ("<article" not in low and "<main" not in low):
        return True
    return False


def safe_name(url, idx):
    host = urllib.parse.urlparse(url).netloc.replace("www.", "").replace(".", "_")
    ts = datetime.now().strftime("%Y%m%d_%H%M%S")
    return f"{idx:04d}_{host}_{ts}"


class CrawlJob:
    """一次爬取执行 (独立线程运行)"""

    def __init__(self, task, run, persist):
        self.task = task
        self.run = run
        self.persist = persist          # callable(task_id, run)
        self._stop = threading.Event()
        self._pause = threading.Event()
        self._cfg_lock = threading.RLock()
        self.thread = None

    # ---------------- 控制接口 ----------------
    def start(self):
        self.thread = threading.Thread(target=self._run_loop, daemon=True)
        self.thread.start()

    def stop(self):
        self._stop.set()

    def pause(self):
        self._pause.set()

    def resume(self):
        self._pause.clear()

    def is_running(self):
        return self.thread is not None and self.thread.is_alive()

    def update_config(self, patch):
        with self._cfg_lock:
            self.task["config"].update(patch)
        self._log("info", f"配置已热更新: {', '.join(patch.keys())}")

    # ---------------- 内部工具 ----------------
    def _cfg(self, key, default=None):
        with self._cfg_lock:
            return self.task["config"].get(key, default)

    def _wait_if_paused(self):
        was_paused = False
        while self._pause.is_set() and not self._stop.is_set():
            if not was_paused:
                self.run["status"] = "paused"
                self._persist()
                was_paused = True
            time.sleep(0.5)
        if was_paused:
            self.run["status"] = "running"
            self._persist()

    def _log(self, level, msg):
        logs = self.run.setdefault("logs", [])
        logs.append({"ts": datetime.now().strftime("%H:%M:%S"), "level": level, "msg": msg})
        if len(logs) > 500:
            del logs[:len(logs) - 500]
        self.persist(self.task["id"], self.run)

    def _persist(self):
        self.persist(self.task["id"], self.run)

    def _resolve_out_dir(self):
        cfg = self._cfg("out_dir") or ""
        if cfg.strip():
            return cfg.strip()
        return os.path.join(HERE, "out", self.task["id"])

    def _open_browser(self):
        p = sync_playwright().start()
        browser = p.chromium.launch(
            headless=True,
            executable_path=CHROME,
            args=["--disable-blink-features=AutomationControlled",
                  "--no-sandbox", "--disable-gpu", "--disable-dev-shm-usage"],
        )
        ctx = browser.new_context(
            user_agent=self._cfg("user_agent", DEFAULT_UA),
            viewport={"width": 1920, "height": 1080},
            locale="en-US",
        )
        cookie_file = os.path.join(DATA_DIR, f"cookies_{self.task['id']}.json")
        if os.path.exists(cookie_file):
            try:
                ctx.add_cookies(json.load(open(cookie_file)))
                self._log("info", "已复用上次会话 cookie")
            except Exception:
                pass
        Stealth().apply_stealth_sync(ctx)
        page = ctx.new_page()
        return p, browser, ctx, page, cookie_file

    def _close_browser(self, p, browser, ctx, cookie_file):
        try:
            json.dump(ctx.cookies(), open(cookie_file, "w"))
        except Exception:
            pass
        try:
            browser.close()
        except Exception:
            pass
        try:
            p.stop()
        except Exception:
            pass

    def _wait_page_settle(self, page, timeout_s):
        last_title, stable = "", 0
        start = time.time()
        while time.time() - start < timeout_s:
            if self._stop.is_set():
                raise RuntimeError("任务已终止")
            self._wait_if_paused()
            time.sleep(1)
            try:
                title = page.title()
                html = page.content()
            except Exception:
                continue                      # 正在跳转
            if is_challenge_page(title, html):
                stable = 0
                continue
            if title == last_title:
                stable += 1
                if stable >= 2 and len(html) > 1000:
                    return True, title, html
            else:
                stable = 0
            last_title = title
        return True, page.title(), page.content()

    def _crawl_one(self, page, url, timeout_s):
        page.goto(url, wait_until="domcontentloaded", timeout=timeout_s * 1000)
        ok, title, html = self._wait_page_settle(page, timeout_s)
        if not ok:
            raise RuntimeError(f"页面加载超时({timeout_s}s)")
        if is_challenge_page(title, html):
            raise RuntimeError(f"仍被反爬拦截: title={title!r} size={len(html)}")
        try:
            text = page.inner_text("body")
        except Exception:
            text = ""
        return title, html, text

    def _crawl_images(self, ctx, page, out_dir, base):
        """下载页面图片, 返回 [{file, url, size}]"""
        try:
            urls = page.evaluate(
                "() => Array.from(document.querySelectorAll('img'))"
                ".map(i => i.currentSrc || i.src).filter(Boolean)"
            )
        except Exception:
            return []
        saved = []
        img_dir = os.path.join(out_dir, base + "_img")
        for n, u in enumerate(urls, 1):
            if self._stop.is_set():
                break
            if not str(u).startswith("http"):
                continue
            path = urllib.parse.urlparse(u).path.lower()
            if not path.endswith(IMG_EXTS):
                continue
            try:
                resp = ctx.request.get(u, timeout=20000)
                if resp.ok and resp.body():
                    ext = os.path.splitext(path)[1] or ".jpg"
                    fname = f"img_{n:04d}{ext}"
                    os.makedirs(img_dir, exist_ok=True)
                    with open(os.path.join(img_dir, fname), "wb") as f:
                        f.write(resp.body())
                    saved.append({"file": f"{base}_img/{fname}", "url": u, "size": len(resp.body())})
            except Exception:
                continue
        return saved

    def _retry_crawl(self, page, ctx, url, idx, out_dir):
        """带重试的单页爬取, 返回结果 entry"""
        timeout = int(self._cfg("timeout", 60))
        retries = int(self._cfg("retry_count", 2))
        retry_wait = float(self._cfg("retry_interval", 3))
        crawl_images = bool(self._cfg("crawl_images", False))

        entry = {"url": url, "title": "", "status": "FAIL", "error": "",
                 "html_file": "", "txt_file": "", "images": []}
        base = safe_name(url, idx)
        for attempt in range(retries + 1):
            if self._stop.is_set():
                entry["error"] = "任务已终止"
                break
            self._wait_if_paused()
            try:
                title, html, text = self._crawl_one(page, url, timeout)
                html_path = os.path.join(out_dir, base + ".html")
                txt_path = os.path.join(out_dir, base + ".txt")
                with open(html_path, "w", encoding="utf-8") as f:
                    f.write(html)
                with open(txt_path, "w", encoding="utf-8") as f:
                    f.write(text)
                entry.update(title=title, status="OK",
                             html_file=base + ".html", txt_file=base + ".txt")
                if crawl_images:
                    entry["images"] = self._crawl_images(ctx, page, out_dir, base)
                self._log("info", f"OK {title[:50]!r} html={len(html)//1024}KB 图片={len(entry['images'])}")
                break
            except Exception as e:
                entry["error"] = str(e)
                self._log("warn", f"第{attempt + 1}次失败 {url}: {e}")
                if attempt < retries:
                    self._wait_if_paused()
                    t0 = time.time()
                    while time.time() - t0 < retry_wait:
                        if self._stop.is_set():
                            break
                        self._wait_if_paused()
                        time.sleep(0.3)
        return entry

    def _delay(self):
        dmin = float(self._cfg("delay_min", 2))
        dmax = float(self._cfg("delay_max", 5))
        total = random.uniform(max(0.1, dmin), max(dmin + 0.1, dmax))
        end = time.time() + total
        while time.time() < end:
            if self._stop.is_set():
                return
            self._wait_if_paused()
            time.sleep(min(0.5, end - time.time()))

    def _bump_stats(self, entry):
        st = self.run.setdefault("stats", {"ok": 0, "fail": 0, "images": 0})
        if entry["status"] == "OK":
            st["ok"] += 1
        else:
            st["fail"] += 1
        st["images"] = st.get("images", 0) + len(entry.get("images", []))

    # ---------------- 主流程 ----------------
    def _run_loop(self):
        run, task = self.run, self.task
        run["status"] = "running"
        run["started_at"] = store.now_str()
        self._persist()
        try:
            if task.get("mode") == "auto":
                self._crawl_auto()
            else:
                self._crawl_list(task.get("urls", []))
            if self._stop.is_set():
                run["status"] = "stopped"
            else:
                run["status"] = "completed"
        except Exception as e:
            run["status"] = "failed"
            self._log("error", f"任务异常终止: {e}")
        run["finished_at"] = store.now_str()
        self._persist()
        if self._cfg("notify", False):
            try:
                ok, msg = notify.notify_email(task, run)
                if ok:
                    self._log("info", "完成通知邮件已发送")
                else:
                    self._log("error", f"邮件通知失败: {msg}")
            except Exception as e:
                self._log("error", f"邮件通知异常: {e}")
                self._persist()
        self._log("info", f"任务结束: {run['status']} 成功{run['stats'].get('ok', 0)} 失败{run['stats'].get('fail', 0)}")
        self._persist()

    def _crawl_list(self, urls):
        run = self.run
        total = len(urls)
        run["progress"]["total"] = total
        out_dir = self._resolve_out_dir()
        run["out_dir"] = out_dir
        os.makedirs(out_dir, exist_ok=True)
        self._persist()

        p, browser, ctx, page, cookie_file = self._open_browser()
        try:
            for i, url in enumerate(urls, 1):
                if self._stop.is_set():
                    self._log("info", "收到终止信号, 停止爬取")
                    break
                self._wait_if_paused()
                run["progress"]["current_url"] = url
                run["progress"]["done"] = i - 1
                self._persist()
                entry = self._retry_crawl(page, ctx, url, i, out_dir)
                run["results"].append(entry)
                self._bump_stats(entry)
                run["progress"]["done"] = i
                self._persist()
                if entry["status"] == "OK":
                    self._delay()
        finally:
            self._close_browser(p, browser, ctx, cookie_file)

    def _discover_links(self, page):
        """从当前页面提取符合规则的链接"""
        auto = self.task.get("auto", {})
        include = [x.strip() for x in (auto.get("include") or []) if x.strip()]
        exclude = [x.strip() for x in (auto.get("exclude") or []) if x.strip()]
        use_regex = bool(auto.get("use_regex", False))
        same_domain = auto.get("same_domain", True)
        seed_host = urllib.parse.urlparse(auto.get("seed_url", "")).hostname or ""
        try:
            hrefs = page.evaluate(
                "() => Array.from(document.querySelectorAll('a[href]')).map(a => a.href)"
            )
        except Exception:
            return []
        out = []
        for h in hrefs:
            if not str(h).startswith("http"):
                continue
            if same_domain and seed_host:
                host = urllib.parse.urlparse(h).hostname or ""
                if host != seed_host and not host.endswith("." + seed_host):
                    continue
            if use_regex:
                if include and not any(re.search(p, h) for p in include):
                    continue
                if any(re.search(p, h) for p in exclude):
                    continue
            else:
                if include and not any(p.lower() in h.lower() for p in include):
                    continue
                if any(p.lower() in h.lower() for p in exclude):
                    continue
            out.append(h)
        return out

    def _crawl_auto(self):
        run = self.run
        auto = self.task.get("auto", {})
        seed = auto.get("seed_url", "")
        max_pages = int(auto.get("max_pages", 50) or 50)
        max_depth = int(auto.get("max_depth", 2) or 2)
        run["progress"]["total"] = max_pages
        out_dir = self._resolve_out_dir()
        run["out_dir"] = out_dir
        os.makedirs(out_dir, exist_ok=True)
        self._persist()

        p, browser, ctx, page, cookie_file = self._open_browser()
        queue = [(seed, 0)]
        visited = set()
        queued = set([seed])
        idx = 0
        try:
            while queue and not self._stop.is_set():
                self._wait_if_paused()
                url, depth = queue.pop(0)
                if url in visited:
                    continue
                if len(visited) >= max_pages:
                    break
                visited.add(url)
                idx += 1
                run["progress"]["current_url"] = url
                run["progress"]["done"] = len(visited)
                self._persist()
                entry = self._retry_crawl(page, ctx, url, idx, out_dir)
                run["results"].append(entry)
                self._bump_stats(entry)
                self._persist()
                if entry["status"] == "OK" and depth < max_depth:
                    for link in self._discover_links(page):
                        if link not in visited and link not in queued:
                            queued.add(link)
                            queue.append((link, depth + 1))
                if entry["status"] == "OK":
                    self._delay()
        finally:
            self._close_browser(p, browser, ctx, cookie_file)
        run["progress"]["total"] = len(visited)
        self._persist()