#!/usr/bin/env python3
"""将探数新闻正文保存为与微信下载器相同的 Markdown 目录格式。"""

from __future__ import annotations

import base64
from dataclasses import asdict, dataclass
from pathlib import Path
import re
import shutil
import tempfile
from urllib.parse import unquote, urljoin, urlsplit

import requests
from bs4 import BeautifulSoup
from markdownify import markdownify
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry


class ArticleDownloadError(RuntimeError):
    """页面不是文章、图片下载失败等可向调用方报告的错误。"""


@dataclass
class ArticleResult:
    title: str
    source_url: str
    directory: str
    markdown: str
    images: list[str]
    warnings: list[str]

    def to_dict(self) -> dict:
        return asdict(self)


def safe_name(title: str) -> str:
    """跨平台目录/文件名；目录名与其中 Markdown 的文件名保持一致。"""
    name = re.sub(r'[<>:"/\\|?*\x00-\x1f\x7f]', '_', title)
    name = re.sub(r'\s+', ' ', name).strip(' .')
    # 同时限制 UTF-8 字节数，给 .md 和重名后缀保留空间。
    name = name.encode('utf-8')[:180].decode('utf-8', errors='ignore').rstrip(' .')
    if not name:
        name = '未命名文章'
    if re.fullmatch(r'(?i:CON|PRN|AUX|NUL|COM[1-9]|LPT[1-9])(?:\..*)?', name):
        name = '_' + name
    return name


def _session() -> requests.Session:
    session = requests.Session()
    session.headers.update({
        'User-Agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) '
                      'AppleWebKit/537.36 (KHTML, like Gecko) Chrome/130.0.0.0 Safari/537.36',
        'Accept-Language': 'zh-CN,zh;q=0.9,en;q=0.8',
    })
    retry = Retry(total=3, backoff_factor=0.6,
                  status_forcelist=[429, 500, 502, 503, 504], allowed_methods=['GET'])
    session.mount('https://', HTTPAdapter(max_retries=retry))
    session.mount('http://', HTTPAdapter(max_retries=retry))
    return session


def _get(session: requests.Session, url: str, timeout: float, limit: int,
         referer: str | None = None) -> tuple[bytes, str, str]:
    if urlsplit(url).scheme not in ('http', 'https'):
        raise ArticleDownloadError(f'不支持的下载地址：{url[:200]}')
    headers = {'Referer': referer} if referer else {}
    with session.get(url, timeout=(min(timeout, 15), timeout),
                     headers=headers, stream=True) as response:
        response.raise_for_status()
        chunks, size = [], 0
        for chunk in response.iter_content(65536):
            size += len(chunk)
            if size > limit:
                raise ArticleDownloadError(f'文件超过大小限制（{limit // 1024 // 1024} MB）：{url[:200]}')
            chunks.append(chunk)
        return b''.join(chunks), response.headers.get('Content-Type', ''), response.url


def _image_ext(data: bytes, mime: str) -> str:
    if data.startswith(b'\xff\xd8\xff'):
        return '.jpg'
    if data.startswith(b'\x89PNG\r\n\x1a\n'):
        return '.png'
    if data.startswith((b'GIF87a', b'GIF89a')):
        return '.gif'
    if data[:4] == b'RIFF' and data[8:12] == b'WEBP':
        return '.webp'
    if len(data) >= 12 and data[4:8] == b'ftyp' and data[8:12] in (b'avif', b'avis'):
        return '.avif'
    if b'<svg' in data[:1024].lower() and mime.split(';')[0] == 'image/svg+xml':
        return '.svg'
    raise ArticleDownloadError(f'返回内容不是支持的图片（Content-Type: {mime or "未知"}）')


def _image_source(img) -> str:
    for attr in ('data-src', 'data-original', 'data-actualsrc', 'data-lazy-src', 'src'):
        value = img.get(attr)
        if value and str(value).strip():
            return str(value).strip()
    srcset = img.get('data-srcset') or img.get('srcset')
    if srcset:
        return str(srcset).split(',')[-1].strip().split()[0]
    return ''


def save_article(article: dict, output_root: str | Path = '.', *, timeout: float = 30,
                 allow_missing_images: bool = False) -> ArticleResult:
    """从 API 正文保存文章；图片失败时清理半成品，不覆盖已有同名目录。"""
    url = str(article.get('weburl') or article.get('url') or '')
    title = str(article.get('title') or '').strip()
    raw_content = str(article.get('content') or '').strip()
    if not title or not raw_content:
        raise ArticleDownloadError('新闻标题或正文为空')
    if timeout <= 0:
        raise ValueError('timeout 必须大于 0')
    with _session() as session:
        final_url = url
        soup = BeautifulSoup('<article>' + raw_content + '</article>', 'html.parser')
        content = soup.article
        for tag in content.select('script, style, noscript, iframe, form, button'):
            tag.decompose()
        # 部分微信排版把插图放在行内 CSS 背景中，也转为普通图片参与下载。
        for tag in list(content.find_all(style=True)):
            for match in re.finditer(r'url\(\s*[\'\"]?(.*?)[\'\"]?\s*\)', str(tag['style'])):
                source = match.group(1).strip()
                if source and not source.startswith('#'):
                    background = soup.new_tag('img', src=source, alt='')
                    tag.insert(0, background)
        if not content.get_text(strip=True) and not content.find('img'):
            raise ArticleDownloadError('文章正文为空，可能需要浏览器执行 JavaScript')
        for link in content.find_all('a', href=True):
            href = str(link['href']).strip()
            if href and not href.startswith('#'):
                absolute = urljoin(final_url, href)
                if urlsplit(absolute).scheme in ('http', 'https', 'mailto', 'tel'):
                    link['href'] = absolute
                else:
                    del link['href']

        root = Path(output_root).expanduser().resolve()
        root.mkdir(parents=True, exist_ok=True)
        stage = Path(tempfile.mkdtemp(prefix='.article-', dir=root))
        images, warnings, downloaded = [], [], {}
        try:
            for img in list(content.find_all('img')):
                source = _image_source(img)
                if not source:
                    img.replace_with(soup.new_string('[图片缺少下载地址]'))
                    if not allow_missing_images:
                        raise ArticleDownloadError('正文中存在没有下载地址的图片')
                    warnings.append('正文中存在没有下载地址的图片')
                    continue
                image_url = urljoin(final_url, source)
                try:
                    if image_url not in downloaded:
                        if source.startswith('data:image/'):
                            header, payload = source.split(',', 1)
                            data = base64.b64decode(payload, validate=True) if ';base64' in header else unquote(payload).encode()
                            mime = header[5:].split(';')[0]
                            if len(data) > 50 * 1024 * 1024:
                                raise ArticleDownloadError('内嵌图片超过 50 MB')
                        else:
                            data, mime, _ = _get(session, image_url, timeout, 50 * 1024 * 1024, final_url)
                        ext = _image_ext(data, mime)
                        filename = f'image{len(images) + 1:03d}{ext}'
                        (stage / filename).write_bytes(data)
                        images.append(filename)
                        downloaded[image_url] = filename
                    # 清除懒加载/响应式属性，使转换器只引用本地文件。
                    alt = str(img.get('alt') or '')
                    img.attrs = {'src': downloaded[image_url], 'alt': alt}
                except (requests.RequestException, ArticleDownloadError, ValueError) as exc:
                    if not allow_missing_images:
                        raise ArticleDownloadError(f'图片下载失败：{image_url[:200]}；{exc}') from exc
                    warnings.append(f'图片下载失败：{image_url[:200]}；{exc}')
                    img.replace_with(soup.new_string(f'[图片下载失败：{img.get("alt") or "未命名图片"}]'))

            # 新闻图片后常带空的图注标题，避免生成单独的 ## 和多余空行。
            for heading in content.select('h1, h2, h3, h4, h5, h6'):
                if not heading.get_text(strip=True) and not heading.find('img'):
                    heading.decompose()
            body = markdownify(str(content), heading_style='ATX', bullets='-',
                               strip=['picture', 'source']).strip()
            # 微信富文本中各段常没有标准 p 标签；markdownify 仍会保留文字和图片顺序。
            author = str(article.get('src') or '未知作者').strip() or '未知作者'
            text = f'# {title}\n\n'
            text += body + f'\n\n来源：{author}，如有侵权，请联系删除。\n'
            basename = safe_name(title)
            # 原子认领新目录，支持多个进程并发保存同名文章。
            index = 1
            while True:
                name = basename if index == 1 else f'{basename} ({index})'
                destination = root / name
                try:
                    destination.mkdir()
                    break
                except FileExistsError:
                    index += 1
            try:
                (stage / f'{name}.md').write_text(text, encoding='utf-8')
                for item in stage.iterdir():
                    shutil.move(str(item), destination / item.name)
            except BaseException:
                shutil.rmtree(destination, ignore_errors=True)
                raise
            return ArticleResult(title, url, str(destination), str(destination / f'{name}.md'),
                                 [str(destination / x) for x in images], warnings)
        finally:
            shutil.rmtree(stage, ignore_errors=True)