# -*- coding: utf-8 -*-
"""
通用爬虫引擎 (Playwright + stealth + 系统 Chrome)
- 批量模式: 逐条爬取网址列表
- 自动模式: 从起始网址按匹配规则自动发现链接并 BFS 爬取
- 试爬取: 仅抓取起始页, 列出按规则将爬取的链接(不保存文件)
- 每个页面/图片生成 .meta.json 操作信息(模式/时间/网址/来源链接/深度等)
- 支持: 随机延迟 / 重试 / 图片下载 / 暂停恢复 / 终止 / 配置热更新 / 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
import db

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
    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}"


def _settle_wait(page, timeout_s):
    """等待页面稳定(无 stop/pause 检查, 供试爬取等独立流程使用)"""
    last_title, stable = "", 0
    start = time.time()
    while time.time() - start < timeout_s:
        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()


_TRACKING_PARAMS = {
    "utm_source", "utm_medium", "utm_campaign", "utm_term", "utm_content",
    "fbclid", "gclid", "yclid", "mc_cid", "mc_eid", "ref", "ref_src",
}


def normalize_url(url):
    """URL 规范化 (用于去重): 去锚点/跟踪参数/尾部斜杠/默认端口, host 小写"""
    try:
        p = urllib.parse.urlparse(str(url))
        host = (p.hostname or "").lower()
        if not host:
            return str(url)
        port = ""
        if p.port and p.port not in (80, 443):
            port = f":{p.port}"
        path = p.path or "/"
        if len(path) > 1 and path.endswith("/"):
            path = path.rstrip("/")
        query = ""
        if p.query:
            kept = [kv for kv in p.query.split("&")
                    if kv.split("=", 1)[0].lower() not in _TRACKING_PARAMS]
            if kept:
                query = "?" + "&".join(kept)
        return f"{p.scheme.lower()}://{host}{port}{path}{query}"
    except Exception:
        return str(url)


def filter_links(hrefs, seed_url, include=None, exclude=None,
                 same_domain=True, use_regex=False):
    """按规则过滤链接, 返回 (included, excluded); excluded 含排除原因"""
    include = include or []
    exclude = exclude or []
    included, excluded = [], []
    seed_host = urllib.parse.urlparse(seed_url).hostname or ""
    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):
                excluded.append({"url": h, "reason": "不在同域名内"})
                continue
        if use_regex:
            if include and not any(re.search(p, h) for p in include):
                excluded.append({"url": h, "reason": "未匹配包含规则"})
                continue
            hit = next((p for p in exclude if re.search(p, h)), None)
            if hit:
                excluded.append({"url": h, "reason": f"命中排除规则: {hit}"})
                continue
        else:
            if include and not any(p.lower() in h.lower() for p in include):
                excluded.append({"url": h, "reason": "未匹配包含规则"})
                continue
            hit = next((p for p in exclude if p.lower() in h.lower()), None)
            if hit:
                excluded.append({"url": h, "reason": f"命中排除规则: {hit}"})
                continue
        included.append(h)
    return included, excluded


def probe_links(seed_url, include=None, exclude=None,
                same_domain=True, use_regex=False, timeout=45):
    """试爬取: 抓取起始页并列出按规则将爬取的链接 (不保存任何文件)"""
    result = {"ok": False, "error": "", "seed_url": seed_url, "title": "",
              "crawled_at": "", "total_links": 0,
              "total_included": 0, "total_excluded": 0,
              "included": [], "excluded": []}
    try:
        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=DEFAULT_UA,
            viewport={"width": 1920, "height": 1080},
            locale="en-US",
        )
        Stealth().apply_stealth_sync(ctx)
        page = ctx.new_page()
        try:
            page.goto(seed_url, wait_until="domcontentloaded", timeout=timeout * 1000)
            _ok, title, html = _settle_wait(page, timeout)
            if is_challenge_page(title, html):
                result["error"] = f"起始页被反爬拦截: title={title!r}"
            else:
                try:
                    hrefs = page.evaluate(
                        "() => Array.from(document.querySelectorAll('a[href]')).map(a => a.href)"
                    )
                except Exception:
                    hrefs = []
                included, excluded = filter_links(
                    hrefs, seed_url, include, exclude, same_domain, use_regex)
                result.update(
                    ok=True, title=title, crawled_at=store.now_str(),
                    total_links=len(hrefs),
                    total_included=len(included), total_excluded=len(excluded),
                    included=included[:200], excluded=excluded[:200],
                )
        except Exception as e:
            result["error"] = f"试爬取失败: {e}"
        finally:
            try:
                browser.close()
            except Exception:
                pass
            try:
                p.stop()
            except Exception:
                pass
    except Exception as e:
        result["error"] = f"浏览器启动失败: {e}"
    return result


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, page_url, source_url):
        """下载页面图片并生成图片集 meta.json, 返回 [{file,url,size,download_time}]"""
        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()), "download_time": store.now_str(),
                    })
            except Exception:
                continue
        if saved:
            try:
                meta = {
                    "type": "images",
                    "mode": self.run.get("mode"),
                    "task_id": self.task["id"],
                    "task_name": self.task.get("name", ""),
                    "run_id": self.run.get("id"),
                    "crawl_time": store.now_str(),
                    "page_url": page_url,
                    "source_url": source_url,
                    "images": saved,
                }
                with open(os.path.join(img_dir, "meta.json"), "w", encoding="utf-8") as f:
                    json.dump(meta, f, ensure_ascii=False, indent=2)
            except Exception:
                pass
        return saved

    def _write_page_meta(self, out_dir, base, entry):
        """为每个爬取页面生成操作信息 meta.json"""
        meta = {
            "type": "page",
            "mode": self.run.get("mode"),
            "task_id": self.task["id"],
            "task_name": self.task.get("name", ""),
            "run_id": self.run.get("id"),
            "crawl_time": entry.get("crawl_time", ""),
            "url": entry.get("url", ""),
            "source_url": entry.get("source_url", ""),
            "depth": entry.get("depth"),
            "title": entry.get("title", ""),
            "status": entry.get("status", ""),
            "error": entry.get("error", ""),
            "attempts": entry.get("attempts", 1),
            "html_file": entry.get("html_file", ""),
            "txt_file": entry.get("txt_file", ""),
            "images": entry.get("images", []),
        }
        try:
            with open(os.path.join(out_dir, base + ".meta.json"), "w", encoding="utf-8") as f:
                json.dump(meta, f, ensure_ascii=False, indent=2)
        except Exception:
            pass

    def _retry_crawl(self, page, ctx, url, idx, out_dir, source_url="", depth=None):
        """带重试的单页爬取, 返回结果 entry (含 meta 信息)"""
        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))

        base = safe_name(url, idx)
        entry = {
            "url": url, "title": "", "status": "FAIL", "error": "",
            "html_file": "", "txt_file": "", "meta_file": base + ".meta.json",
            "crawl_time": store.now_str(),
            "source_url": source_url, "depth": depth,
            "images": [], "attempts": 0,
        }
        for attempt in range(retries + 1):
            if self._stop.is_set():
                entry["error"] = "任务已终止"
                break
            self._wait_if_paused()
            entry["attempts"] += 1
            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",
                             error="", crawl_time=store.now_str())
                if crawl_images:
                    entry["images"] = self._crawl_images(ctx, page, out_dir, base, url, source_url)
                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)
                entry["crawl_time"] = store.now_str()
                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)
        self._write_page_meta(out_dir, base, entry)
        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._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", {})
        try:
            hrefs = page.evaluate(
                "() => Array.from(document.querySelectorAll('a[href]')).map(a => a.href)"
            )
        except Exception:
            return []
        included, _excluded = filter_links(
            hrefs, auto.get("seed_url", ""),
            auto.get("include", []), auto.get("exclude", []),
            auto.get("same_domain", True), bool(auto.get("use_regex", False)),
        )
        return included

    def _crawl_auto(self):
        """自动爬取: 持续递归 爬取->发现链接->爬取... 直到无新链接可爬或达到上限
        - max_depth: 0=无限制, N=只爬 N 层
        - max_pages: 0=无限制, N=安全上限
        - 提取但未爬取的链接持久化到 auto.pending, 已爬集合持久化到 auto.visited
        - 停止后再次运行从缓存队列继续爬; 起始网址每次运行都重新爬(不去重)
        """
        run = self.run
        auto = self.task.get("auto", {})
        seed = auto.get("seed_url", "")
        max_pages = int(auto.get("max_pages", 1000) or 0)     # 0 = 无限制
        max_depth = int(auto.get("max_depth", 0) or 0)        # 0 = 无限制
        out_dir = self._resolve_out_dir()
        run["out_dir"] = out_dir
        os.makedirs(out_dir, exist_ok=True)

        # 恢复持久化状态: 已爬集合 + 上次未爬完的缓存队列
        visited = set(auto.get("visited", []) or [])
        pending = auto.get("pending", []) or []
        # 起始网址每次运行都爬(不做去重), 缓存队列继续消费
        queue = [(seed, 0, "")]
        if pending:
            queue.extend((p["url"], p.get("depth", 0), p.get("source", "")) for p in pending)
        queued = set(visited)
        for u, _d, _s in queue:
            queued.add(normalize_url(u))

        run["progress"]["total"] = len(queue)
        self._persist()

        p, browser, ctx, page, cookie_file = self._open_browser()
        try:
            while queue and not self._stop.is_set():
                self._wait_if_paused()
                if max_pages > 0 and len(visited) >= max_pages:
                    self._log("info", f"达到最大页数上限 {max_pages}, 停止")
                    break
                url, depth, src = queue.pop(0)
                key = normalize_url(url)
                if key in visited and url != seed:     # 起始网址不去重, 其余已爬跳过
                    continue
                visited.add(key)
                idx = len(visited)
                run["progress"]["current_url"] = url
                run["progress"]["done"] = len(visited)
                self._persist()
                entry = self._retry_crawl(page, ctx, url, idx, out_dir,
                                          source_url=src, depth=depth)
                run["results"].append(entry)
                self._bump_stats(entry)
                self._persist()
                # 无深度限制或未达深度限制时持续发现链接
                if entry["status"] == "OK" and (max_depth == 0 or depth < max_depth):
                    for link in self._discover_links(page):
                        lk = normalize_url(link)
                        if lk not in visited and lk not in queued:
                            queued.add(lk)
                            queue.append((link, depth + 1, url))
                if entry["status"] == "OK":
                    self._delay()
            run["progress"]["total"] = len(visited)
            if not self._stop.is_set() and len(queue) == 0:
                self._log("info", f"无新链接可爬, 任务结束 (共 {len(visited)} 页)")
        finally:
            self._close_browser(p, browser, ctx, cookie_file)
            # 无论完成/停止/异常, 都保存缓存队列与已爬集合, 便于下次继续
            try:
                auto["pending"] = [
                    {"url": u, "depth": d, "source": s} for u, d, s in queue]
                auto["visited"] = list(visited)
                store.upsert_task(self.task)
                db.upsert_task(self.task)
            except Exception:
                pass
        self._persist()