# -*- coding: UTF-8 -*-
'''
@Project :domainScanDemo
@File    :wayback_detector.py
@IDE     :PyCharm
@Author  :梦伴
@Date    :2026/4/8 23:57
@explain : Wayback检测器
'''

import html
import json
import re
import threading
import time
import zlib
from base64 import b64decode, b64encode
from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait

import redis
import requests
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry

from app.config import config
from app.detectors.base import BaseDetector


class WaybackDetector(BaseDetector):
    """
    Wayback检测器
    """

    TITLE_PATTERN = re.compile(r'<title[^>]*>(.*?)</title>', re.IGNORECASE | re.DOTALL)

    def __init__(self):
        """
        初始化Wayback检测器
        """
        super().__init__()
        self.cdx_api_url = 'https://web.archive.org/cdx/search/cdx'
        # 使用 id_ 直接取快照内容，避免回放页面额外壳层干扰正文匹配。
        self.snapshot_url = 'https://web.archive.org/web/{timestamp}id_/{domain}'
        self._timestamp_cache = {}
        self._title_cache = {}
        self._cache_lock = threading.Lock()
        self.session = self._build_session()
        self.redis_client = self._build_redis_client()

    def _build_session(self):
        session = requests.Session()
        retry = Retry(
            total=max(0, config.WAYBACK_RETRY_COUNT),
            backoff_factor=0.5,
            status_forcelist=(429, 500, 502, 503, 504),
            allowed_methods=frozenset(["GET"]),
            raise_on_status=False,
        )
        adapter = HTTPAdapter(max_retries=retry, pool_connections=10, pool_maxsize=10)
        session.mount('http://', adapter)
        session.mount('https://', adapter)
        session.headers.update({
            'User-Agent': config.WAYBACK_USER_AGENT,
            'Accept-Language': 'zh-CN,zh;q=0.9,en;q=0.8',
        })
        return session

    def _build_redis_client(self):
        try:
            client = redis.Redis(
                host=config.REDIS_HOST,
                port=config.REDIS_PORT,
                password=config.REDIS_PASSWORD,
                db=config.REDIS_DB,
                decode_responses=True,
                socket_connect_timeout=5,
                socket_timeout=5,
            )
            client.ping()
            return client
        except Exception:
            return None

    def _title_cache_key(self, domain, timestamp):
        return f"domain_tool:wayback_title:{domain}:{timestamp}"

    def _timestamp_cache_key(self, domain):
        return f"domain_tool:wayback_timestamps:{domain}"

    def _record_cache_key(self, domain):
        return f"domain_tool:wayback_records:{domain}"

    def _normalize_title(self, title):
        normalized = html.unescape(title or '')
        normalized = re.sub(r'\s+', ' ', normalized, flags=re.DOTALL).strip().lower()
        return normalized

    def _extract_title(self, content):
        match = self.TITLE_PATTERN.search(content or '')
        if not match:
            return ''
        return html.unescape(match.group(1)).strip()

    def _load_cached_title(self, domain, timestamp):
        cache_key = self._title_cache_key(domain, timestamp)
        with self._cache_lock:
            if cache_key in self._title_cache:
                return self._title_cache[cache_key]
        if self.redis_client:
            try:
                raw_value = self.redis_client.get(cache_key)
                if raw_value:
                    data = json.loads(raw_value)
                    with self._cache_lock:
                        self._title_cache[cache_key] = data
                    return data
            except Exception:
                pass
        return None

    def _save_cached_title(self, domain, timestamp, data):
        cache_key = self._title_cache_key(domain, timestamp)
        with self._cache_lock:
            self._title_cache[cache_key] = data
        if self.redis_client:
            try:
                self.redis_client.set(
                    cache_key,
                    json.dumps(data, ensure_ascii=False),
                    ex=max(0, config.WAYBACK_TITLE_CACHE_TTL) or None,
                )
            except Exception:
                pass

    def _load_cached_timestamps(self, domain):
        if domain in self._timestamp_cache:
            return list(self._timestamp_cache[domain])
        if not self.redis_client:
            return None
        try:
            raw_value = self.redis_client.get(self._timestamp_cache_key(domain))
            if not raw_value:
                return None
            compressed = b64decode(raw_value.encode('ascii'))
            timestamps = json.loads(zlib.decompress(compressed).decode('utf-8'))
            if isinstance(timestamps, list):
                self._timestamp_cache[domain] = tuple(timestamps)
                return list(timestamps)
        except Exception:
            return None
        return None

    def _save_cached_timestamps(self, domain, timestamps):
        self._timestamp_cache[domain] = tuple(timestamps)
        if not self.redis_client:
            return
        try:
            payload = json.dumps(timestamps, separators=(',', ':')).encode('utf-8')
            compressed = zlib.compress(payload, level=6)
            self.redis_client.set(
                self._timestamp_cache_key(domain),
                b64encode(compressed).decode('ascii'),
                ex=max(0, config.WAYBACK_TIMESTAMP_CACHE_TTL) or None,
            )
        except Exception:
            pass

    def _load_cached_records(self, domain):
        cache_key = self._record_cache_key(domain)
        with self._cache_lock:
            if cache_key in self._timestamp_cache:
                return [
                    {'timestamp': item[0], 'digest': item[1]}
                    for item in self._timestamp_cache[cache_key]
                ]
        if not self.redis_client:
            return None
        try:
            raw_value = self.redis_client.get(cache_key)
            if not raw_value:
                return None
            compressed = b64decode(raw_value.encode('ascii'))
            records = json.loads(zlib.decompress(compressed).decode('utf-8'))
            if isinstance(records, list):
                normalized = tuple(
                    (item.get('timestamp', ''), item.get('digest', ''))
                    for item in records if isinstance(item, dict)
                )
                with self._cache_lock:
                    self._timestamp_cache[cache_key] = normalized
                return [
                    {'timestamp': item[0], 'digest': item[1]}
                    for item in normalized if item[0]
                ]
        except Exception:
            return None
        return None

    def _save_cached_records(self, domain, records):
        cache_key = self._record_cache_key(domain)
        normalized = tuple(
            (item.get('timestamp', ''), item.get('digest', ''))
            for item in (records or [])
            if item and item.get('timestamp')
        )
        with self._cache_lock:
            self._timestamp_cache[cache_key] = normalized
        if not self.redis_client:
            return
        try:
            payload = json.dumps(
                [{'timestamp': item[0], 'digest': item[1]} for item in normalized],
                separators=(',', ':')
            ).encode('utf-8')
            compressed = zlib.compress(payload, level=6)
            self.redis_client.set(
                cache_key,
                b64encode(compressed).decode('ascii'),
                ex=max(0, config.WAYBACK_TIMESTAMP_CACHE_TTL) or None,
            )
        except Exception:
            pass

    def _fetch_cdx_records(self, domain, limit=None, fast_latest=False):
        response = None
        try:
            params = {
                'url': domain,
                'output': 'txt',
                'fl': 'timestamp,digest',
                'filter': ['statuscode:200', 'mimetype:text/html'],
            }
            if limit is not None:
                params['limit'] = str(limit)
            if fast_latest:
                params['fastLatest'] = 'true'
            response = self.session.get(
                self.cdx_api_url,
                params=params,
                timeout=config.WAYBACK_CDX_TIMEOUT,
                stream=True,
            )
            if response.status_code != 200:
                self._log_warning(f"获取快照记录失败: {response.status_code}")
                return []
            records = []
            seen = set()
            for raw_line in response.iter_lines(decode_unicode=True):
                line = (raw_line or '').strip()
                if not line:
                    continue
                parts = line.split()
                timestamp = parts[0].strip() if parts else ''
                digest = parts[1].strip() if len(parts) > 1 else ''
                if not timestamp or timestamp in seen:
                    continue
                seen.add(timestamp)
                records.append({'timestamp': timestamp, 'digest': digest})
            return records
        except Exception as e:
            self._handle_exception(e, domain)
            return []
        finally:
            try:
                if response is not None:
                    response.close()
            except Exception:
                pass

    def get_latest_snapshot_record(self, domain):
        records = self._fetch_cdx_records(domain, limit=-1, fast_latest=True)
        return records[0] if records else None

    def get_snapshot_records(self, domain):
        cached_records = self._load_cached_records(domain)
        if cached_records is not None:
            return cached_records
        records = self._fetch_cdx_records(domain)
        if records:
            self._save_cached_records(domain, records)
            self._save_cached_timestamps(domain, [item['timestamp'] for item in records])
        return records

    def check_domain(self, domain):
        """
        检测域名

        :param domain: 域名
        :return: dict - 检测结果
        """
        try:
            return self.scan_snapshots(domain)
        except Exception as e:
            return self._handle_exception(e, domain)

    def get_snapshot_timestamps(self, domain):
        cached_records = self._load_cached_records(domain)
        if cached_records is not None:
            return [item['timestamp'] for item in cached_records if item.get('timestamp')]
        cached_timestamps = self._load_cached_timestamps(domain)
        if cached_timestamps is not None:
            return cached_timestamps
        records = self.get_snapshot_records(domain)
        return [item['timestamp'] for item in records if item.get('timestamp')]

    def _fetch_snapshot_title(self, domain, timestamp):
        cached = self._load_cached_title(domain, timestamp)
        if cached is not None:
            return cached

        snapshot_url = self.snapshot_url.format(timestamp=timestamp, domain=domain)
        response = None
        try:
            response = self.session.get(
                snapshot_url,
                timeout=config.WAYBACK_SNAPSHOT_TIMEOUT,
                stream=True,
            )
            if response.status_code != 200:
                data = {'timestamp': timestamp, 'title': '', 'ok': False}
                self._save_cached_title(domain, timestamp, data)
                return data

            content_type = (response.headers.get('Content-Type') or '').lower()
            if content_type and 'text/html' not in content_type and 'application/xhtml+xml' not in content_type:
                data = {'timestamp': timestamp, 'title': '', 'ok': False}
                self._save_cached_title(domain, timestamp, data)
                return data

            chunks = []
            total_bytes = 0
            found_title = False
            for chunk in response.iter_content(chunk_size=4096, decode_unicode=True):
                if not chunk:
                    continue
                chunks.append(chunk)
                total_bytes += len(chunk.encode('utf-8', errors='ignore'))
                current_text = ''.join(chunks)
                if '</title>' in current_text.lower():
                    found_title = True
                    break
                if total_bytes >= config.WAYBACK_TITLE_MAX_BYTES:
                    break

            content = ''.join(chunks)
            title = self._extract_title(content) if found_title or content else ''
            data = {'timestamp': timestamp, 'title': title, 'ok': True}
            self._save_cached_title(domain, timestamp, data)
            return data
        except Exception:
            data = {'timestamp': timestamp, 'title': '', 'ok': False}
            self._save_cached_title(domain, timestamp, data)
            return data
        finally:
            try:
                if response is not None:
                    response.close()
            except Exception:
                pass

    def get_snapshot_years(self, domain):
        """
        获取快照年份

        :param domain: 域名
        :return: list - 快照年份列表
        """
        try:
            years = set()
            for timestamp in self.get_snapshot_timestamps(domain):
                if len(timestamp) >= 4:
                    years.add(int(timestamp[:4]))
            return sorted(years)
        except Exception as e:
            self._handle_exception(e, domain)
            return []

    def has_sensitive_content(self, domain):
        """
        检查是否包含敏感内容

        :param domain: 域名
        :return: bool - 是否包含敏感内容
        """
        try:
            result = self.scan_snapshots(domain)
            return result.get('has_sensitive_content', False)
        except Exception as e:
            self._handle_exception(e, domain)
            return False

    def scan_snapshots(self, domain, sensitive_words=None, stop_on_first_hit=True):
        sensitive_words = sensitive_words or config.load_sensitive_words()
        latest_record = self.get_latest_snapshot_record(domain)
        latest_timestamp = (latest_record or {}).get('timestamp')
        latest_digest = (latest_record or {}).get('digest', '')
        matched_word = None
        matched_timestamp = None
        matched_title = None
        fetched_snapshot_count = 0
        failed_snapshot_count = 0
        unique_title_count = 0
        duplicate_title_skipped = 0
        digest_duplicate_skipped = 0
        started_at = time.time()
        progress_interval = max(1, config.WAYBACK_PROGRESS_INTERVAL)
        title_seen = set()
        digest_seen = set()
        checked_snapshot_count = 0
        domain_concurrency = max(1, config.WAYBACK_DOMAIN_CONCURRENCY)

        if latest_timestamp:
            latest_result = self._fetch_snapshot_title(domain, latest_timestamp)
            checked_snapshot_count = 1
            if latest_result and latest_result.get('ok'):
                fetched_snapshot_count = 1
                title = latest_result.get('title', '')
                normalized_title = self._normalize_title(title)
                if latest_digest:
                    digest_seen.add(latest_digest)
                if normalized_title:
                    title_seen.add(normalized_title)
                    unique_title_count = 1
                    matched_word = self._find_sensitive_word(title, sensitive_words)
                    if matched_word and stop_on_first_hit:
                        matched_timestamp = latest_timestamp
                        matched_title = title
                        return {
                            'snapshot_years': [int(latest_timestamp[:4])] if len(latest_timestamp) >= 4 else [],
                            'has_sensitive_content': True,
                            'matched_word': matched_word,
                            'matched_timestamp': matched_timestamp,
                            'matched_title': matched_title,
                            'backlink_count': 0,
                            'backlink_count_gt_10': False,
                            'checked_snapshot_count': checked_snapshot_count,
                            'fetched_snapshot_count': fetched_snapshot_count,
                            'failed_snapshot_count': failed_snapshot_count,
                            'unique_title_count': unique_title_count,
                            'duplicate_title_skipped': duplicate_title_skipped,
                            'digest_duplicate_skipped': digest_duplicate_skipped,
                            'elapsed_seconds': round(time.time() - started_at, 2),
                        }
            else:
                failed_snapshot_count += 1

        records = sorted(self.get_snapshot_records(domain), key=lambda item: item.get('timestamp', ''), reverse=True)
        years = sorted({int(item['timestamp'][:4]) for item in records if len(item.get('timestamp', '')) >= 4})
        checked_snapshot_count = len(records)
        pending_records = []
        for item in records:
            timestamp = item.get('timestamp', '')
            digest = item.get('digest', '')
            if not timestamp:
                continue
            if latest_timestamp and timestamp == latest_timestamp:
                continue
            if digest and digest in digest_seen:
                digest_duplicate_skipped += 1
                continue
            if digest:
                digest_seen.add(digest)
            pending_records.append(item)

        with ThreadPoolExecutor(max_workers=domain_concurrency) as executor:
            pending = {}
            index = 0
            finished_count = 1 if latest_timestamp else 0
            stop_requested = False

            while (index < len(pending_records) or pending) and not stop_requested:
                while index < len(pending_records) and len(pending) < domain_concurrency and not stop_requested:
                    timestamp = pending_records[index]['timestamp']
                    future = executor.submit(self._fetch_snapshot_title, domain, timestamp)
                    pending[future] = timestamp
                    index += 1

                if not pending:
                    break

                done, _ = wait(list(pending.keys()), return_when=FIRST_COMPLETED)
                for future in done:
                    timestamp = pending.pop(future, None)
                    finished_count += 1
                    try:
                        result = future.result()
                    except Exception as e:
                        failed_snapshot_count += 1
                        self._handle_exception(e, domain)
                        continue

                    if not result or not result.get('ok'):
                        failed_snapshot_count += 1
                        continue

                    fetched_snapshot_count += 1
                    title = result.get('title', '')
                    normalized_title = self._normalize_title(title)
                    if normalized_title:
                        if normalized_title in title_seen:
                            duplicate_title_skipped += 1
                        else:
                            title_seen.add(normalized_title)
                            unique_title_count += 1
                            matched_word = self._find_sensitive_word(title, sensitive_words)
                            if matched_word:
                                matched_timestamp = timestamp
                                matched_title = title
                                if stop_on_first_hit:
                                    stop_requested = True
                    if finished_count % progress_interval == 0:
                        elapsed = round(time.time() - started_at, 2)
                        self._log_info(
                            f"{domain} 时光机进度: {finished_count}/{checked_snapshot_count}，成功 {fetched_snapshot_count}，失败 {failed_snapshot_count}，唯一标题 {unique_title_count}，标题重复跳过 {duplicate_title_skipped}，digest 重复跳过 {digest_duplicate_skipped}，耗时 {elapsed}s"
                        )
                    if config.WAYBACK_REQUEST_DELAY > 0:
                        time.sleep(config.WAYBACK_REQUEST_DELAY)

            if stop_requested:
                for future in pending:
                    future.cancel()

        return {
            'snapshot_years': years,
            'has_sensitive_content': matched_word is not None,
            'matched_word': matched_word,
            'matched_timestamp': matched_timestamp,
            'matched_title': matched_title,
            'backlink_count': 0,
            'backlink_count_gt_10': False,
            'checked_snapshot_count': checked_snapshot_count,
            'fetched_snapshot_count': fetched_snapshot_count,
            'failed_snapshot_count': failed_snapshot_count,
            'unique_title_count': unique_title_count,
            'duplicate_title_skipped': duplicate_title_skipped,
            'digest_duplicate_skipped': digest_duplicate_skipped,
            'elapsed_seconds': round(time.time() - started_at, 2),
        }

    def _check_sensitive_words(self, content):
        """
        检查敏感词

        :param content: 内容
        :return: bool - 是否包含敏感词
        """
        return self._find_sensitive_word(content, config.load_sensitive_words()) is not None

    def _find_sensitive_word(self, content, sensitive_words):
        for word in sensitive_words or []:
            if word and word in (content or ''):
                return word
        return None

    def get_backlink_count(self, domain):
        """
        当前策略仅扫描标题，不再抓取正文，友链数量默认返回 0。
        """
        return 0

    def _count_backlinks(self, content):
        return 0